alita-sdk 0.3.462__py3-none-any.whl → 0.3.627__py3-none-any.whl

alita_sdk/runtime/utils/mcp_oauth.py

@@ -43,6 +43,23 @@ class McpAuthorizationRequired(ToolException):
         }
 
 
+def extract_authorization_uri(www_authenticate: Optional[str]) -> Optional[str]:
+    """
+    Extract authorization_uri from WWW-Authenticate header.
+    This points directly to the OAuth authorization server metadata URL.
+    Should be used before falling back to resource_metadata.
+    """
+    if not www_authenticate:
+        return None
+    
+    # Look for authorization_uri="<url>" in the header
+    match = re.search(r'authorization_uri\s*=\s*\"?([^\", ]+)\"?', www_authenticate)
+    if match:
+        return match.group(1)
+    
+    return None
+
+
 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.
@@ -62,15 +79,33 @@ def extract_resource_metadata_url(www_authenticate: Optional[str], server_url: O
     # 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]]:
+def fetch_oauth_authorization_server_metadata(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.
+    
+    Args:
+        url: Either a full well-known URL (e.g., https://api.figma.com/.well-known/oauth-authorization-server)
+             or a base URL (e.g., https://api.figma.com) where we'll try discovery endpoints.
+        timeout: Request timeout in seconds.
+    
+    Returns:
+        OAuth authorization server metadata dict, or None if not found.
     """
+    # If the URL is already a .well-known endpoint, try it directly first
+    if '/.well-known/' in url:
+        try:
+            resp = requests.get(url, timeout=timeout)
+            if resp.status_code == 200:
+                return resp.json()
+        except Exception as exc:
+            logger.debug(f"Failed to fetch OAuth metadata from {url}: {exc}")
+        # If direct fetch failed, don't try other endpoints
+        return None
+    
+    # Otherwise, try standard discovery endpoints
     discovery_endpoints = [
-        f"{base_url}/.well-known/oauth-authorization-server",
-        f"{base_url}/.well-known/openid-configuration",
+        f"{url}/.well-known/oauth-authorization-server",
+        f"{url}/.well-known/openid-configuration",
     ]
     
     for endpoint in discovery_endpoints:
@@ -162,3 +197,165 @@ def canonical_resource(server_url: str) -> str:
     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: Optional[str] = None,
+    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 (optional for DCR/public clients)
+        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
+    
+    Note:
+        client_id may be optional for:
+        - Dynamic Client Registration (DCR): client_id may be in the code
+        - OIDC public clients: some providers don't require it
+        - Some MCP servers handle auth differently
+    """
+    # Build the token request body
+    token_body = {
+        "grant_type": "authorization_code",
+        "code": code,
+        "redirect_uri": redirect_uri,
+    }
+    
+    if client_id:
+        token_body["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}")
+
+
+def refresh_oauth_token(
+    token_endpoint: str,
+    refresh_token: str,
+    client_id: Optional[str] = None,
+    client_secret: Optional[str] = None,
+    scope: Optional[str] = None,
+    timeout: int = 30,
+) -> Dict[str, Any]:
+    """
+    Refresh an OAuth access token using a refresh token.
+    
+    Args:
+        token_endpoint: OAuth token endpoint URL
+        refresh_token: Refresh token from previous authorization
+        client_id: OAuth client ID (optional for DCR/public clients)
+        client_secret: OAuth client secret (optional for public clients)
+        scope: OAuth scope (optional)
+        timeout: Request timeout in seconds
+        
+    Returns:
+        Token response from OAuth provider containing access_token, etc.
+        May also include a new refresh_token depending on the provider.
+        
+    Raises:
+        requests.RequestException: If the HTTP request fails
+        ValueError: If the token refresh fails
+    
+    Note:
+        client_id may be optional for:
+        - Dynamic Client Registration (DCR): client_id embedded in refresh_token
+        - OIDC public clients: some providers don't require it
+        - Some MCP servers handle auth differently
+    """
+    token_body = {
+        "grant_type": "refresh_token",
+        "refresh_token": refresh_token,
+    }
+    
+    if client_id:
+        token_body["client_id"] = client_id
+    if client_secret:
+        token_body["client_secret"] = client_secret
+    if scope:
+        token_body["scope"] = scope
+
+    logger.info(f"MCP OAuth: refreshing token at {token_endpoint}")
+    
+    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 refresh 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 refresh failed - {response.status_code}: {error_msg}")
+        raise ValueError(f"Token refresh failed: {error_msg}")
+

alita_sdk/runtime/utils/mcp_sse_client.py

@@ -71,6 +71,7 @@ class McpSseClient:
                     McpAuthorizationRequired,
                     canonical_resource,
                     extract_resource_metadata_url,
+                    extract_authorization_uri,
                     fetch_resource_metadata_async,
                     infer_authorization_servers_from_realm,
                     fetch_oauth_authorization_server_metadata
@@ -79,13 +80,41 @@ class McpSseClient:
                 auth_header = self._stream_response.headers.get('WWW-Authenticate', '')
                 resource_metadata_url = extract_resource_metadata_url(auth_header, self.url)
                 
+                # First, try authorization_uri from WWW-Authenticate header (preferred)
+                authorization_uri = extract_authorization_uri(auth_header)
+                
                 metadata = None
-                if resource_metadata_url:
-                    metadata = await fetch_resource_metadata_async(
-                        resource_metadata_url,
-                        session=self._stream_session,
-                        timeout=30
-                    )
+                if authorization_uri:
+                    # Fetch OAuth metadata directly from authorization_uri
+                    auth_server_metadata = fetch_oauth_authorization_server_metadata(authorization_uri, timeout=30)
+                    if auth_server_metadata:
+                        # Extract base authorization server URL from the issuer or the well-known URL
+                        base_auth_server = auth_server_metadata.get('issuer')
+                        if not base_auth_server and '/.well-known/' in authorization_uri:
+                            base_auth_server = authorization_uri.split('/.well-known/')[0]
+                        
+                        metadata = {
+                            'authorization_servers': [base_auth_server] if base_auth_server else [authorization_uri],
+                            'oauth_authorization_server': auth_server_metadata
+                        }
+                        logger.info(f"[MCP SSE Client] Using authorization_uri: {authorization_uri}, base: {base_auth_server}")
+                
+                # Fall back to resource_metadata if authorization_uri didn't work
+                if not metadata:
+                    if resource_metadata_url:
+                        metadata = await fetch_resource_metadata_async(
+                            resource_metadata_url,
+                            session=self._stream_session,
+                            timeout=30
+                        )
+                        # If we got resource_metadata, also fetch oauth_authorization_server
+                        if metadata and metadata.get('authorization_servers'):
+                            auth_server_metadata = fetch_oauth_authorization_server_metadata(
+                                metadata['authorization_servers'][0], timeout=30
+                            )
+                            if auth_server_metadata:
+                                metadata['oauth_authorization_server'] = auth_server_metadata
+                                logger.info(f"[MCP SSE Client] Fetched OAuth metadata from resource_metadata")
                 
                 # Infer authorization servers if not in metadata
                 if not metadata or not metadata.get('authorization_servers'):
@@ -345,7 +374,7 @@ class McpSseClient:
                     "sampling": {}
                 },
                 "clientInfo": {
-                    "name": "Alita MCP Client",
+                    "name": "ELITEA MCP Client",
                     "version": "1.0.0"
                 }
             }

alita_sdk/runtime/utils/mcp_tools_discovery.py

@@ -0,0 +1,124 @@
+"""
+MCP Tools Discovery Utility.
+Provides a standalone function to discover tools from remote MCP servers.
+Supports both SSE (Server-Sent Events) and Streamable HTTP transports with auto-detection.
+"""
+
+import asyncio
+import logging
+from typing import Any, Dict, List, Optional
+
+from .mcp_oauth import McpAuthorizationRequired
+from .mcp_client import McpClient
+
+logger = logging.getLogger(__name__)
+
+
+def discover_mcp_tools(
+    url: str,
+    headers: Optional[Dict[str, str]] = None,
+    timeout: int = 60,
+    session_id: Optional[str] = None,
+) -> List[Dict[str, Any]]:
+    """
+    Discover available tools from a remote MCP server.
+    
+    This function connects to a remote MCP server and retrieves the list of 
+    available tools using the MCP protocol. Automatically detects and uses
+    the appropriate transport (SSE or Streamable HTTP).
+    
+    Args:
+        url: MCP server HTTP URL (http:// or https://)
+        headers: Optional HTTP headers for authentication
+        timeout: Request timeout in seconds (default: 60)
+        session_id: Optional session ID for stateful connections
+        
+    Returns:
+        List of tool definitions, each containing:
+        - name: Tool name
+        - description: Tool description
+        - inputSchema: JSON schema for tool input parameters
+        
+    Raises:
+        McpAuthorizationRequired: If the server requires OAuth authorization (401)
+        Exception: For other connection or protocol errors
+        
+    Example:
+        >>> tools = discover_mcp_tools(
+        ...     url="https://mcp.example.com/sse",
+        ...     headers={"Authorization": "Bearer token123"}
+        ... )
+        >>> print(f"Found {len(tools)} tools")
+    """
+    logger.info(f"[MCP Discovery] Starting tool discovery from {url}")
+    
+    try:
+        # Run the async discovery in a new event loop
+        tools_list = asyncio.run(
+            _discover_tools_async(url, headers, timeout, session_id)
+        )
+        logger.info(f"[MCP Discovery] Successfully discovered {len(tools_list)} tools from {url}")
+        return tools_list
+        
+    except McpAuthorizationRequired:
+        # Re-raise auth exceptions directly
+        logger.info(f"[MCP Discovery] Authorization required for {url}")
+        raise
+        
+    except Exception as e:
+        logger.error(f"[MCP Discovery] Failed to discover tools from {url}: {e}")
+        raise
+
+
+async def _discover_tools_async(
+    url: str,
+    headers: Optional[Dict[str, str]],
+    timeout: int,
+    session_id: Optional[str],
+) -> List[Dict[str, Any]]:
+    """
+    Async implementation of tool discovery using unified MCP client.
+    """
+    all_tools = []
+    
+    # Create unified MCP client (auto-detects transport)
+    client = McpClient(
+        url=url,
+        session_id=session_id,
+        headers=headers,
+        timeout=timeout
+    )
+    
+    async with client:
+        # Initialize MCP session
+        await client.initialize()
+        logger.debug(f"[MCP Discovery] Session initialized (transport={client.detected_transport})")
+        
+        # Get tools list
+        tools = await client.list_tools()
+        logger.debug(f"[MCP Discovery] Received {len(tools)} tools")
+        
+        # Convert tools to standard format
+        for tool in tools:
+            tool_def = {
+                'name': tool.get('name'),
+                'description': tool.get('description', ''),
+                'inputSchema': tool.get('inputSchema', {}),
+            }
+            all_tools.append(tool_def)
+    
+    return all_tools
+
+
+async def discover_mcp_tools_async(
+    url: str,
+    headers: Optional[Dict[str, str]] = None,
+    timeout: int = 60,
+    session_id: Optional[str] = None,
+) -> List[Dict[str, Any]]:
+    """
+    Async version of discover_mcp_tools.
+    
+    See discover_mcp_tools for full documentation.
+    """
+    return await _discover_tools_async(url, headers, timeout, session_id)

alita_sdk/runtime/utils/serialization.py

@@ -0,0 +1,155 @@
+"""
+Serialization utilities for safe JSON encoding of complex objects.
+
+Handles Pydantic models, LangChain messages, datetime objects, and other
+non-standard types that may appear in state variables.
+"""
+import json
+import logging
+from datetime import datetime, date
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def _convert_to_serializable(obj: Any, _seen: set = None) -> Any:
+    """
+    Recursively convert an object to JSON-serializable primitives.
+
+    Handles nested dicts and lists that may contain non-serializable objects.
+    Uses a seen set to prevent infinite recursion with circular references.
+
+    Args:
+        obj: Any object to convert
+        _seen: Internal set to track seen object ids (for circular reference detection)
+
+    Returns:
+        JSON-serializable representation of the object
+    """
+    # Initialize seen set for circular reference detection
+    if _seen is None:
+        _seen = set()
+
+    # Check for circular references (only for mutable objects)
+    obj_id = id(obj)
+    if isinstance(obj, (dict, list, set)) and obj_id in _seen:
+        return f"<circular reference: {type(obj).__name__}>"
+
+    # Primitives - return as-is
+    if obj is None or isinstance(obj, (str, int, float, bool)):
+        return obj
+
+    # Add to seen set for mutable containers
+    if isinstance(obj, (dict, list, set)):
+        _seen = _seen | {obj_id}  # Create new set to avoid mutation issues
+
+    # Dict - recursively process all values
+    if isinstance(obj, dict):
+        return {
+            _convert_to_serializable(k, _seen): _convert_to_serializable(v, _seen)
+            for k, v in obj.items()
+        }
+
+    # List/tuple - recursively process all items
+    if isinstance(obj, (list, tuple)):
+        return [_convert_to_serializable(item, _seen) for item in obj]
+
+    # Set - convert to list and process
+    if isinstance(obj, set):
+        return [_convert_to_serializable(item, _seen) for item in obj]
+
+    # Bytes - decode to string
+    if isinstance(obj, bytes):
+        try:
+            return obj.decode('utf-8')
+        except UnicodeDecodeError:
+            return obj.decode('utf-8', errors='replace')
+
+    # Datetime objects
+    if isinstance(obj, datetime):
+        return obj.isoformat()
+    if isinstance(obj, date):
+        return obj.isoformat()
+
+    # Pydantic BaseModel (v2) - check for model_dump method
+    if hasattr(obj, 'model_dump') and callable(getattr(obj, 'model_dump')):
+        try:
+            return _convert_to_serializable(obj.model_dump(), _seen)
+        except Exception as e:
+            logger.debug(f"Failed to call model_dump on {type(obj).__name__}: {e}")
+
+    # Pydantic BaseModel (v1) - check for dict method
+    if hasattr(obj, 'dict') and callable(getattr(obj, 'dict')) and hasattr(obj, '__fields__'):
+        try:
+            return _convert_to_serializable(obj.dict(), _seen)
+        except Exception as e:
+            logger.debug(f"Failed to call dict on {type(obj).__name__}: {e}")
+
+    # LangChain BaseMessage - extract key fields
+    if hasattr(obj, 'type') and hasattr(obj, 'content'):
+        try:
+            result = {
+                "type": obj.type,
+                "content": _convert_to_serializable(obj.content, _seen),
+            }
+            if hasattr(obj, 'additional_kwargs') and obj.additional_kwargs:
+                result["additional_kwargs"] = _convert_to_serializable(obj.additional_kwargs, _seen)
+            if hasattr(obj, 'name') and obj.name:
+                result["name"] = obj.name
+            return result
+        except Exception as e:
+            logger.debug(f"Failed to extract message fields from {type(obj).__name__}: {e}")
+
+    # Objects with __dict__ attribute (custom classes)
+    if hasattr(obj, '__dict__'):
+        try:
+            return _convert_to_serializable(obj.__dict__, _seen)
+        except Exception as e:
+            logger.debug(f"Failed to serialize __dict__ of {type(obj).__name__}: {e}")
+
+    # UUID objects
+    if hasattr(obj, 'hex') and hasattr(obj, 'int'):
+        return str(obj)
+
+    # Enum objects
+    if hasattr(obj, 'value') and hasattr(obj, 'name') and hasattr(obj.__class__, '__members__'):
+        return obj.value
+
+    # Last resort - convert to string
+    try:
+        return str(obj)
+    except Exception:
+        return f"<non-serializable: {type(obj).__name__}>"
+
+
+def safe_serialize(obj: Any, **kwargs) -> str:
+    """
+    Safely serialize any object to a JSON string.
+
+    Pre-processes the entire object tree to convert non-serializable
+    objects before passing to json.dumps. This ensures nested dicts
+    and lists with non-standard objects are handled correctly.
+
+    Args:
+        obj: Any object to serialize
+        **kwargs: Additional arguments passed to json.dumps
+            (e.g., indent, sort_keys)
+
+    Returns:
+        JSON string representation of the object
+
+    Example:
+        >>> from pydantic import BaseModel
+        >>> class User(BaseModel):
+        ...     name: str
+        >>> state = {"user": User(name="Alice"), "count": 5}
+        >>> safe_serialize(state)
+        '{"user": {"name": "Alice"}, "count": 5}'
+    """
+    # Pre-process the entire object tree
+    serializable = _convert_to_serializable(obj)
+
+    # Set defaults
+    kwargs.setdefault('ensure_ascii', False)
+
+    return json.dumps(serializable, **kwargs)

alita_sdk/runtime/utils/streamlit.py

@@ -287,7 +287,6 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                     model_config={
                         "temperature": 0.1,
                         "max_tokens": 1000,
-                        "top_p": 1.0
                     }
                 )
             except Exception as e:
@@ -1256,7 +1255,6 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                     model_config={
                         "temperature": 0.1,
                         "max_tokens": 1000,
-                        "top_p": 1.0
                     }
                 )
             except Exception as e:
@@ -1387,20 +1385,18 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                 help="Maximum number of tokens in the AI response"
                             )
                             
-                            top_p = st.slider(
-                                "Top-p:",
-                                min_value=0.1,
-                                max_value=1.0,
-                                value=1.0,
-                                step=0.1,
-                                help="Controls diversity via nucleus sampling"
+                            reasoning_effort = st.selectbox(
+                                "Reasoning effort:",
+                                options=['null', 'low', 'medium', 'high'],
+                                index=0,
+                                help="Higher effort better reasoning, slower response"
                             )
                         
                         # Create LLM config
                         llm_config = {
                             'max_tokens': max_tokens,
                             'temperature': temperature,
-                            'top_p': top_p
+                            'reasoning_effort': reasoning_effort
                         }
                         
                         col1, col2 = st.columns([3, 1])

alita_sdk/runtime/utils/toolkit_utils.py

@@ -12,7 +12,9 @@ logger = logging.getLogger(__name__)
 
 def instantiate_toolkit_with_client(toolkit_config: Dict[str, Any], 
                                    llm_client: Any, 
-                                   alita_client: Optional[Any] = None) -> List[Any]:
+                                   alita_client: Optional[Any] = None,
+                                   mcp_tokens: Optional[Dict[str, Any]] = None,
+                                   use_prefix: bool = False) -> List[Any]:
     """
     Instantiate a toolkit with LLM client support.
     
@@ -22,7 +24,11 @@ def instantiate_toolkit_with_client(toolkit_config: Dict[str, Any],
     Args:
         toolkit_config: Configuration dictionary for the toolkit
         llm_client: LLM client instance for tools that need LLM capabilities
-        client: Optional additional client instance
+        alita_client: Optional additional client instance
+        mcp_tokens: Optional dictionary of MCP OAuth tokens by server URL
+        use_prefix: If True, tools get prefixed with toolkit_name to prevent collisions
+                   (for agent use). If False, tools use base names only (for testing interface).
+                   Default False for backward compatibility with testing.
     
     Returns:
         List of instantiated tools from the toolkit
@@ -52,16 +58,19 @@ def instantiate_toolkit_with_client(toolkit_config: Dict[str, Any],
         toolkit_type = toolkit_config.get('type', toolkit_name.lower())
 
         # Create a tool configuration dict with required fields
+        # Note: MCP toolkit always requires toolkit_name, other toolkits respect use_prefix flag
+        # Note: 'name' is always set for provider-based toolkits (used by provider_worker.utils.tools)
         tool_config = {
             'id': toolkit_config.get('id', random.randint(1, 1000000)),
             'type': toolkit_config.get('type', toolkit_type),
             'settings': settings,
-            'toolkit_name': toolkit_name
+            'name': toolkit_name,  # Always pass name for provider toolkits
+            'toolkit_name': toolkit_name if (use_prefix or toolkit_type == 'mcp') else None
         }
         
         # Get tools using the toolkit configuration with clients
-        # Parameter order: get_tools(tools_list, alita_client, llm, memory_store)
-        tools = get_tools([tool_config], alita_client, llm_client)
+        # Parameter order: get_tools(tools_list, alita_client, llm, memory_store, debug_mode, mcp_tokens)
+        tools = get_tools([tool_config], alita_client, llm_client, mcp_tokens=mcp_tokens)
         
         if not tools:
             logger.warning(f"No tools returned for toolkit {toolkit_name}")
@@ -73,9 +82,11 @@ def instantiate_toolkit_with_client(toolkit_config: Dict[str, Any],
     except Exception as e:
         # Re-raise McpAuthorizationRequired without logging as error
         from ..utils.mcp_oauth import McpAuthorizationRequired
+        
         if isinstance(e, McpAuthorizationRequired):
             logger.info(f"Toolkit {toolkit_name} requires MCP OAuth authorization")
             raise
+        
         # Log and re-raise other errors
         logger.error(f"Error instantiating toolkit {toolkit_name} with client: {str(e)}")
         raise