alita-sdk 0.3.423__py3-none-any.whl → 0.3.449__py3-none-any.whl

alita_sdk/runtime/utils/mcp_oauth.py

@@ -0,0 +1,164 @@
+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

alita_sdk/runtime/utils/mcp_sse_client.py

@@ -0,0 +1,347 @@
+"""
+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}")
+            
+            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...")
+        
+        if self._stream_task and not self._stream_task.done():
+            self._stream_task.cancel()
+            try:
+                await self._stream_task
+            except asyncio.CancelledError:
+                pass
+        
+        if self._stream_response:
+            self._stream_response.close()
+        
+        if self._stream_session:
+            await self._stream_session.close()
+        
+        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

alita_sdk/runtime/utils/streamlit.py

@@ -868,10 +868,24 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                 label = f"{'🔒 ' if is_secret else ''}{'*' if is_required else ''}{field_name.replace('_', ' ').title()}"
                                 
                                 if field_type == 'string':
-                                    if is_secret:
+                                    # Check if this is an enum field
+                                    if field_schema.get('enum'):
+                                        # Dropdown for enum values
+                                        options = field_schema['enum']
+                                        default_index = 0
+                                        if default_value and str(default_value) in options:
+                                            default_index = options.index(str(default_value))
+                                        toolkit_config_values[field_name] = st.selectbox(
+                                            label,
+                                            options=options,
+                                            index=default_index,
+                                            help=field_description,
+                                            key=f"config_{field_name}_{selected_toolkit_idx}"
+                                        )
+                                    elif is_secret:
                                         toolkit_config_values[field_name] = st.text_input(
                                             label,
-                                            value=str(default_value) if default_value else '', 
+                                            value=str(default_value) if default_value else '',
                                             help=field_description,
                                             type="password",
                                             key=f"config_{field_name}_{selected_toolkit_idx}"
@@ -879,7 +893,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                     else:
                                         toolkit_config_values[field_name] = st.text_input(
                                             label,
-                                            value=str(default_value) if default_value else '', 
+                                            value=str(default_value) if default_value else '',
                                             help=field_description,
                                             key=f"config_{field_name}_{selected_toolkit_idx}"
                                         )
@@ -971,6 +985,23 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                                 key=f"config_{field_name}_{selected_toolkit_idx}"
                                             )
                                         toolkit_config_values[field_name] = [line.strip() for line in array_input.split('\n') if line.strip()]
+                                elif field_type == 'object':
+                                    # Handle object/dict types (like headers)
+                                    obj_input = st.text_area(
+                                        f"{label} (JSON object)",
+                                        value=json.dumps(default_value) if isinstance(default_value, dict) else str(default_value) if default_value else '',
+                                        help=f"{field_description} - Enter as JSON object, e.g. {{\"Authorization\": \"Bearer token\"}}",
+                                        placeholder='{"key": "value"}',
+                                        key=f"config_{field_name}_{selected_toolkit_idx}"
+                                    )
+                                    try:
+                                        if obj_input.strip():
+                                            toolkit_config_values[field_name] = json.loads(obj_input)
+                                        else:
+                                            toolkit_config_values[field_name] = None
+                                    except json.JSONDecodeError as e:
+                                        st.error(f"Invalid JSON format for {field_name}: {e}")
+                                        toolkit_config_values[field_name] = None
                         else:
                             st.info("This toolkit doesn't require additional configuration.")
                         

alita_sdk/runtime/utils/toolkit_utils.py

@@ -29,13 +29,14 @@ def instantiate_toolkit_with_client(toolkit_config: Dict[str, Any],
         
     Raises:
         ValueError: If required configuration or client is missing
+        McpAuthorizationRequired: If MCP server requires OAuth authorization
         Exception: If toolkit instantiation fails
     """
+    toolkit_name = toolkit_config.get('toolkit_name', 'unknown')
     try:
         from ..toolkits.tools import get_tools
         
-        toolkit_name = toolkit_config.get('toolkit_name')
-        if not toolkit_name:
+        if not toolkit_name or toolkit_name == 'unknown':
             raise ValueError("toolkit_name is required in configuration")
         
         if not llm_client:
@@ -46,11 +47,14 @@ def instantiate_toolkit_with_client(toolkit_config: Dict[str, Any],
         # Log the configuration being used
         logger.info(f"Instantiating toolkit {toolkit_name} with LLM client")
         logger.debug(f"Toolkit {toolkit_name} configuration: {toolkit_config}")
-        
+
+        # Use toolkit type from config, or fall back to lowercase toolkit name
+        toolkit_type = toolkit_config.get('type', toolkit_name.lower())
+
         # Create a tool configuration dict with required fields
         tool_config = {
             'id': toolkit_config.get('id', random.randint(1, 1000000)),
-            'type': toolkit_config.get('type', toolkit_name.lower()),
+            'type': toolkit_config.get('type', toolkit_type),
             'settings': settings,
             'toolkit_name': toolkit_name
         }
@@ -67,6 +71,12 @@ def instantiate_toolkit_with_client(toolkit_config: Dict[str, Any],
         return tools
             
     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
 

alita_sdk/tools/__init__.py

@@ -131,6 +131,11 @@ def get_tools(tools_list, alita, llm, store: Optional[BaseStore] = None, *args,
                 logger.error(f"Error getting ADO repos tools: {e}")
             continue
 
+        # Skip MCP toolkit - it's handled by runtime/toolkits/tools.py to avoid duplicate loading
+        if tool_type == 'mcp':
+            logger.debug(f"Skipping MCP toolkit '{tool.get('toolkit_name')}' - handled by runtime toolkit system")
+            continue
+
         # Handle standard tools
         if tool_type in AVAILABLE_TOOLS and 'get_tools' in AVAILABLE_TOOLS[tool_type]:
             try:

alita_sdk/tools/chunkers/sematic/proposal_chunker.py

@@ -6,7 +6,7 @@ from langchain_core.prompts import ChatPromptTemplate
 from langchain.text_splitter import TokenTextSplitter
 
 from typing import Optional, List
-from langchain_core.pydantic_v1 import BaseModel
+from pydantic import BaseModel
 from ..utils import tiktoken_length
 
 logger = getLogger(__name__)

alita_sdk/tools/gitlab/api_wrapper.py

@@ -115,6 +115,11 @@ class GitLabAPIWrapper(CodeIndexerToolkit):
         """Remove trailing slash from URL if present."""
         return url.rstrip('/') if url else url
 
+    @model_validator(mode='before')
+    @classmethod
+    def validate_toolkit_before(cls, values: Dict) -> Dict:
+        return super().validate_toolkit(values)
+
     @model_validator(mode='after')
     def validate_toolkit(self):
         try: