alita-sdk 0.3.448__py3-none-any.whl → 0.3.450__py3-none-any.whl

alita_sdk/runtime/toolkits/mcp.py

@@ -7,6 +7,7 @@ Following MCP specification: https://modelcontextprotocol.io/specification/2025-
 import logging
 import re
 import requests
+import asyncio
 from typing import List, Optional, Any, Dict, Literal, ClassVar, Union
 
 from langchain_core.tools import BaseToolkit, BaseTool
@@ -17,6 +18,7 @@ from ..tools.mcp_remote_tool import McpRemoteTool
 from ..tools.mcp_inspect_tool import McpInspectTool
 from ...tools.utils import TOOLKIT_SPLITTER, clean_string
 from ..models.mcp_models import McpConnectionConfig
+from ..utils.mcp_sse_client import McpSseClient
 from ..utils.mcp_oauth import (
     McpAuthorizationRequired,
     canonical_resource,
@@ -459,42 +461,79 @@ class McpToolkit(BaseToolkit):
         timeout: int
     ) -> List[Dict[str, Any]]:
         """
-        Synchronously discover tools and prompts from MCP server using HTTP requests.
+        Discover tools and prompts from MCP server using SSE client.
         Returns list of tool/prompt dictionaries with name, description, and inputSchema.
         Prompts are converted to tools that can be invoked.
         """
-        all_tools = []
-        
-        # Use session_id from UI (passed via connection_config)
-        # For SSE-based MCP servers, frontend generates a UUID and sends it with the OAuth token
         session_id = connection_config.session_id
         
         if not session_id:
             logger.warning(f"[MCP Session] No session_id provided for '{toolkit_name}' - server may require it")
             logger.warning(f"[MCP Session] Frontend should generate a UUID and include it with mcp_tokens")
         
-        # Discover regular tools
-        tools_data = cls._discover_mcp_endpoint(
-            endpoint="tools/list",
-            toolkit_name=toolkit_name,
-            connection_config=connection_config,
-            timeout=timeout,
-            session_id=session_id
+        # Run async discovery in sync context
+        try:
+            all_tools = asyncio.run(
+                cls._discover_tools_async(
+                    toolkit_name=toolkit_name,
+                    connection_config=connection_config,
+                    timeout=timeout
+                )
+            )
+            return all_tools, session_id
+        except Exception as e:
+            logger.error(f"[MCP SSE] Discovery failed for '{toolkit_name}': {e}")
+            raise
+    
+    @classmethod
+    async def _discover_tools_async(
+        cls,
+        toolkit_name: str,
+        connection_config: McpConnectionConfig,
+        timeout: int
+    ) -> List[Dict[str, Any]]:
+        """
+        Async implementation of tool discovery using SSE client.
+        """
+        all_tools = []
+        session_id = connection_config.session_id
+        
+        # Generate temporary session_id if not provided (for OAuth flow)
+        # The real session_id should come from frontend after OAuth completes
+        if not session_id:
+            import uuid
+            session_id = str(uuid.uuid4())
+            logger.info(f"[MCP SSE] Generated temporary session_id for OAuth: {session_id}")
+        
+        logger.info(f"[MCP SSE] Discovering from {connection_config.url} with session {session_id}")
+        
+        # Prepare headers
+        headers = {}
+        if connection_config.headers:
+            headers.update(connection_config.headers)
+        
+        # Create SSE client
+        client = McpSseClient(
+            url=connection_config.url,
+            session_id=session_id,
+            headers=headers,
+            timeout=timeout
         )
-        all_tools.extend(tools_data)
-        logger.info(f"Discovered {len(tools_data)} tools from MCP toolkit '{toolkit_name}'")
         
-        # Discover prompts and convert them to tools
+        # Initialize MCP session
+        await client.initialize()
+        logger.info(f"[MCP SSE] Session initialized for '{toolkit_name}'")
+        
+        # Discover tools
+        tools = await client.list_tools()
+        all_tools.extend(tools)
+        logger.info(f"[MCP SSE] Discovered {len(tools)} tools from '{toolkit_name}'")
+        
+        # Discover prompts
         try:
-            prompts_data = cls._discover_mcp_endpoint(
-                endpoint="prompts/list",
-                toolkit_name=toolkit_name,
-                connection_config=connection_config,
-                timeout=timeout,
-                session_id=session_id
-            )
+            prompts = await client.list_prompts()
             # Convert prompts to tool format
-            for prompt in prompts_data:
+            for prompt in prompts:
                 prompt_tool = {
                     "name": f"prompt_{prompt.get('name', 'unnamed')}",
                     "description": prompt.get('description', f"Execute prompt: {prompt.get('name')}"),
@@ -519,152 +558,12 @@ class McpToolkit(BaseToolkit):
                     "_mcp_prompt_name": prompt.get('name')
                 }
                 all_tools.append(prompt_tool)
-            logger.info(f"Discovered {len(prompts_data)} prompts from MCP toolkit '{toolkit_name}'")
-        except Exception as e:
-            logger.warning(f"Failed to discover prompts from MCP toolkit '{toolkit_name}': {e}")
-        
-        logger.info(f"Total discovered {len(all_tools)} tools+prompts from MCP toolkit '{toolkit_name}'")
-        return all_tools, session_id
-
-    @classmethod
-    def _discover_mcp_endpoint(
-        cls,
-        endpoint: str,
-        toolkit_name: str,
-        connection_config: McpConnectionConfig,
-        timeout: int,
-        session_id: Optional[str] = None
-    ) -> List[Dict[str, Any]]:
-        """
-        Discover items from a specific MCP endpoint (tools/list or prompts/list).
-        Returns list of dictionaries.
-        """
-        import time
-
-        # MCP protocol request
-        mcp_request = {
-            "jsonrpc": "2.0",
-            "id": f"discover_{endpoint.replace('/', '_')}_{int(time.time())}",
-            "method": endpoint,
-            "params": {}
-        }
-
-        headers = {
-            "Content-Type": "application/json",
-            "Accept": "application/json, text/event-stream"
-        }
-        if connection_config.headers:
-            headers.update(connection_config.headers)
-
-        # Add sessionId to URL if provided (for stateful SSE servers)
-        url = connection_config.url
-        if session_id:
-            separator = '&' if '?' in url else '?'
-            url = f"{url}{separator}sessionId={session_id}"
-            logger.info(f"[MCP Session] Using session {session_id} for {endpoint} request")
-        else:
-            logger.debug(f"[MCP Session] No session ID available for {endpoint} request - server may not require sessions")
-
-        try:
-            logger.debug(f"Sending MCP {endpoint} request to {url}")
-            response = requests.post(
-                url,
-                json=mcp_request,
-                headers=headers,
-                timeout=timeout
-            )
-
-            auth_header = response.headers.get('WWW-Authenticate', '')
-            if response.status_code == 401:
-                resource_metadata_url = extract_resource_metadata_url(auth_header, connection_config.url)
-                metadata = fetch_resource_metadata(resource_metadata_url, timeout=timeout) if resource_metadata_url else None
-                
-                # If we couldn't get metadata from the resource_metadata endpoint,
-                # infer authorization servers from the WWW-Authenticate header and server URL
-                if not metadata or not metadata.get('authorization_servers'):
-                    inferred_servers = infer_authorization_servers_from_realm(auth_header, connection_config.url)
-                    if inferred_servers:
-                        if not metadata:
-                            metadata = {}
-                        metadata['authorization_servers'] = inferred_servers
-                        logger.info(f"Inferred authorization servers for {connection_config.url}: {inferred_servers}")
-                        
-                        # Fetch OAuth authorization server metadata from the inferred server
-                        # This avoids CORS issues in the frontend
-                        from alita_sdk.runtime.utils.mcp_oauth import fetch_oauth_authorization_server_metadata
-                        auth_server_metadata = fetch_oauth_authorization_server_metadata(inferred_servers[0], timeout=timeout)
-                        if auth_server_metadata:
-                            metadata['oauth_authorization_server'] = auth_server_metadata
-                            logger.info(f"Fetched OAuth metadata for {inferred_servers[0]}")
-                
-                raise McpAuthorizationRequired(
-                    message=f"MCP server {connection_config.url} requires OAuth authorization",
-                    server_url=canonical_resource(connection_config.url),
-                    resource_metadata_url=resource_metadata_url,
-                    www_authenticate=auth_header,
-                    resource_metadata=metadata,
-                    status=response.status_code,
-                    tool_name=toolkit_name,
-                )
-
-            if response.status_code != 200:
-                logger.error(f"MCP server returned non-200 status: {response.status_code}")
-                raise Exception(f"HTTP {response.status_code}: {response.text}")
-
-            # Check content type and parse accordingly
-            content_type = response.headers.get('Content-Type', '')
-            
-            if 'text/event-stream' in content_type:
-                # Parse SSE (Server-Sent Events) format
-                data = cls._parse_sse_response(response.text)
-            elif 'application/json' in content_type:
-                # Parse regular JSON
-                try:
-                    data = response.json()
-                except ValueError as json_err:
-                    raise Exception(f"Invalid JSON response: {json_err}. Response text: {response.text[:200]}")
-            else:
-                raise Exception(f"Unexpected Content-Type: {content_type}. Response: {response.text[:200]}")
-
-            if "error" in data:
-                raise Exception(f"MCP Error: {data['error']}")
-
-            # Parse MCP response - different endpoints return different keys
-            result = data.get("result", {})
-            if endpoint == "tools/list":
-                return result.get("tools", [])
-            elif endpoint == "prompts/list":
-                return result.get("prompts", [])
-            else:
-                return result.get("items", [])
-
+            logger.info(f"[MCP SSE] Discovered {len(prompts)} prompts from '{toolkit_name}'")
         except Exception as e:
-            logger.error(f"Failed to discover from {endpoint} on MCP toolkit '{toolkit_name}': {e}")
-            raise
-
-    @staticmethod
-    def _parse_sse_response(sse_text: str) -> Dict[str, Any]:
-        """
-        Parse Server-Sent Events (SSE) format response.
-        SSE format: event: message\ndata: {json}\n\n
-        """
-        import json
+            logger.warning(f"[MCP SSE] Failed to discover prompts: {e}")
         
-        lines = sse_text.strip().split('\n')
-        data_line = None
-        
-        for line in lines:
-            if line.startswith('data:'):
-                data_line = line[5:].strip()  # Remove 'data:' prefix
-                break
-        
-        if not data_line:
-            raise Exception(f"No data found in SSE response: {sse_text[:200]}")
-        
-        try:
-            return json.loads(data_line)
-        except json.JSONDecodeError as e:
-            raise Exception(f"Failed to parse SSE data as JSON: {e}. Data: {data_line[:200]}")
+        logger.info(f"[MCP SSE] Total discovered {len(all_tools)} items from '{toolkit_name}'")
+        return all_tools
 
     @classmethod
     def _create_tool_from_dict(

alita_sdk/runtime/tools/mcp_remote_tool.py

@@ -20,6 +20,7 @@ from ..utils.mcp_oauth import (
     fetch_resource_metadata_async,
     infer_authorization_servers_from_realm,
 )
+from ..utils.mcp_sse_client import McpSseClient
 
 logger = logging.getLogger(__name__)
 
@@ -83,172 +84,68 @@ class McpRemoteTool(McpServerTool):
         return asyncio.run(self._execute_remote_tool(kwargs))
 
     async def _execute_remote_tool(self, kwargs: Dict[str, Any]) -> str:
-        """Execute the actual remote MCP tool call."""
-        import aiohttp
+        """Execute the actual remote MCP tool call using SSE client."""
         from ...tools.utils import TOOLKIT_SPLITTER
         
+        # Check for session_id requirement
+        if not self.session_id:
+            logger.error(f"[MCP Session] Missing session_id for tool '{self.name}'")
+            raise Exception("sessionId required. Frontend must generate UUID and send with mcp_tokens.")
+        
         # Use the original tool name from discovery for MCP server invocation
-        # The MCP server doesn't know about our optimized names
         tool_name_for_server = self.original_tool_name
-        
-        # Fallback: extract from optimized name if original not stored (backwards compatibility)
         if not tool_name_for_server:
             tool_name_for_server = self.name.rsplit(TOOLKIT_SPLITTER, 1)[-1] if TOOLKIT_SPLITTER in self.name else self.name
-            logger.warning(f"original_tool_name not set for '{self.name}', using extracted name: {tool_name_for_server}")
+            logger.warning(f"original_tool_name not set for '{self.name}', using extracted: {tool_name_for_server}")
         
-        # Build the MCP request based on whether this is a prompt or tool
-        if self.is_prompt:
-            # For prompts, use prompts/get endpoint
-            mcp_request = {
-                "jsonrpc": "2.0",
-                "id": f"prompt_get_{int(time.time())}_{uuid.uuid4().hex[:8]}",
-                "method": "prompts/get",
-                "params": {
-                    "name": self.prompt_name or tool_name_for_server.replace("prompt_", ""),
-                    "arguments": kwargs.get("arguments", kwargs)
-                }
-            }
-        else:
-            # For regular tools, use tools/call endpoint
-            mcp_request = {
-                "jsonrpc": "2.0",
-                "id": f"tool_call_{int(time.time())}_{uuid.uuid4().hex[:8]}",
-                "method": "tools/call",
-                "params": {
-                    "name": tool_name_for_server,
-                    "arguments": kwargs
-                }
-            }
-
-        # Set up headers
-        headers = {
-            "Content-Type": "application/json",
-            "Accept": "application/json, text/event-stream"
-        }
-        if self.server_headers:
-            headers.update(self.server_headers)
-
-        # Add sessionId to URL if this is a stateful SSE server
-        url = self.server_url
-        if self.session_id:
-            separator = '&' if '?' in url else '?'
-            url = f"{url}{separator}sessionId={self.session_id}"
-            logger.debug(f"Using session URL: {url}")
-
-        # Execute the HTTP request
-        timeout = aiohttp.ClientTimeout(total=self.tool_timeout_sec)
-        async with aiohttp.ClientSession(timeout=timeout) as session:
-            try:
-                logger.debug(f"Calling remote MCP tool '{tool_name_for_server}' (optimized name: '{self.name}') at {url}")
-                logger.debug(f"Request: {json.dumps(mcp_request, indent=2)}")
+        logger.info(f"[MCP SSE] Executing tool '{tool_name_for_server}' with session {self.session_id}")
+        
+        try:
+            # Prepare headers
+            headers = {}
+            if self.server_headers:
+                headers.update(self.server_headers)
+            
+            # Create SSE client
+            client = McpSseClient(
+                url=self.server_url,
+                session_id=self.session_id,
+                headers=headers,
+                timeout=self.tool_timeout_sec
+            )
+            
+            # Execute tool call via SSE
+            result = await client.call_tool(tool_name_for_server, kwargs)
+            
+            # Format the result
+            if isinstance(result, dict):
+                # Check for content array (common in MCP responses)
+                if "content" in result:
+                    content_items = result["content"]
+                    if isinstance(content_items, list):
+                        # Extract text from content items
+                        text_parts = []
+                        for item in content_items:
+                            if isinstance(item, dict):
+                                if item.get("type") == "text" and "text" in item:
+                                    text_parts.append(item["text"])
+                                elif "text" in item:
+                                    text_parts.append(item["text"])
+                                else:
+                                    text_parts.append(json.dumps(item))
+                            else:
+                                text_parts.append(str(item))
+                        return "\n".join(text_parts)
                 
-                async with session.post(url, json=mcp_request, headers=headers) as response:
-                    auth_header = response.headers.get('WWW-Authenticate') or response.headers.get('Www-Authenticate')
-                    if response.status == 401:
-                        resource_metadata_url = extract_resource_metadata_url(auth_header, self.server_url)
-                        metadata = None
-                        if resource_metadata_url:
-                            metadata = await fetch_resource_metadata_async(
-                                resource_metadata_url,
-                                session=session,
-                                timeout=self.tool_timeout_sec,
-                            )
-                        
-                        # If we couldn't get metadata from the resource_metadata endpoint,
-                        # infer authorization servers from the WWW-Authenticate header and server URL
-                        if not metadata or not metadata.get('authorization_servers'):
-                            inferred_servers = infer_authorization_servers_from_realm(auth_header, self.server_url)
-                            if inferred_servers:
-                                if not metadata:
-                                    metadata = {}
-                                metadata['authorization_servers'] = inferred_servers
-                                logger.info(f"Inferred authorization servers for {self.server_url}: {inferred_servers}")
-                                
-                                # Fetch OAuth authorization server metadata from the inferred server
-                                # This avoids CORS issues in the frontend
-                                from alita_sdk.runtime.utils.mcp_oauth import fetch_oauth_authorization_server_metadata
-                                auth_server_metadata = fetch_oauth_authorization_server_metadata(inferred_servers[0], timeout=self.tool_timeout_sec)
-                                if auth_server_metadata:
-                                    metadata['oauth_authorization_server'] = auth_server_metadata
-                                    logger.info(f"Fetched OAuth metadata for {inferred_servers[0]}")
-                        
-                        raise McpAuthorizationRequired(
-                            message=f"MCP server {self.server_url} requires OAuth authorization",
-                            server_url=canonical_resource(self.server_url),
-                            resource_metadata_url=resource_metadata_url,
-                            www_authenticate=auth_header,
-                            resource_metadata=metadata,
-                            status=response.status,
-                            tool_name=self.name,
-                        )
-
-                    # Check for errors first
-                    data = None
-                    if response.status != 200:
-                        error_text = await response.text()
-                        
-                        # Check if this is a "Missing sessionId" error
-                        if response.status == 404 and "sessionId" in error_text:
-                            logger.error(f"[MCP Session] Server requires session but none provided")
-                            logger.error(f"[MCP Session] Frontend must generate a UUID and send it with mcp_tokens")
-                            logger.error(f"[MCP Session] Example: mcp_tokens = {{'server_url': {{'access_token': '...', 'session_id': crypto.randomUUID()}}}}")
-                            raise Exception(f"HTTP {response.status}: {error_text} - sessionId required but not provided by frontend")
-                        else:
-                            raise Exception(f"HTTP {response.status}: {error_text}")
-                    
-                    # Parse response if not already done in retry logic
-                    if data is None:
-                        content_type = response.headers.get('Content-Type', '')
-                        if 'text/event-stream' in content_type:
-                            # Parse SSE format
-                            text = await response.text()
-                            data = self._parse_sse(text)
-                        else:
-                            # Parse regular JSON
-                            data = await response.json()
-
-                    logger.debug(f"Response: {json.dumps(data, indent=2)}")
-
-                    # Check for MCP error
-                    if "error" in data:
-                        error = data["error"]
-                        error_msg = error.get("message", str(error))
-                        raise Exception(f"MCP Error: {error_msg}")
-
-                    # Extract result
-                    result = data.get("result", {})
-                    
-                    # Format the result based on content type
-                    if isinstance(result, dict):
-                        # Check for content array (common in MCP responses)
-                        if "content" in result:
-                            content_items = result["content"]
-                            if isinstance(content_items, list):
-                                # Extract text from content items
-                                text_parts = []
-                                for item in content_items:
-                                    if isinstance(item, dict):
-                                        if item.get("type") == "text" and "text" in item:
-                                            text_parts.append(item["text"])
-                                        elif "text" in item:
-                                            text_parts.append(item["text"])
-                                        else:
-                                            text_parts.append(json.dumps(item))
-                                    else:
-                                        text_parts.append(str(item))
-                                return "\n".join(text_parts)
-                        
-                        # Return formatted JSON if no content field
-                        return json.dumps(result, indent=2)
-                    
-                    # Return as string for other types
-                    return str(result)
-
-            except asyncio.TimeoutError:
-                raise Exception(f"Tool execution timed out after {self.tool_timeout_sec}s")
-            except Exception as e:
-                logger.error(f"Error calling remote MCP tool '{tool_name_for_server}': {e}", exc_info=True)
-                raise
+                # Return formatted JSON if no content field
+                return json.dumps(result, indent=2)
+            
+            # Return as string for other types
+            return str(result)
+            
+        except Exception as e:
+            logger.error(f"[MCP SSE] Tool execution failed: {e}", exc_info=True)
+            raise
 
     def _parse_sse(self, text: str) -> Dict[str, Any]:
         """Parse Server-Sent Events (SSE) format response."""

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

{alita_sdk-0.3.448.dist-info → alita_sdk-0.3.450.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.448
+Version: 0.3.450
 Summary: SDK for building langchain agents using resources from Alita
 Author-email: Artem Rozumenko <artyom.rozumenko@gmail.com>, Mikalai Biazruchka <mikalai_biazruchka@epam.com>, Roman Mitusov <roman_mitusov@epam.com>, Ivan Krakhmaliuk <lifedj27@gmail.com>, Artem Dubrovskiy <ad13box@gmail.com>
 License-Expression: Apache-2.0
@@ -24,6 +24,7 @@ Requires-Dist: fastapi==0.115.9
 Requires-Dist: httpcore==1.0.7
 Requires-Dist: urllib3>=2
 Requires-Dist: certifi==2024.8.30
+Requires-Dist: aiohttp>=3.9.0
 Provides-Extra: runtime
 Requires-Dist: streamlit>=1.28.0; extra == "runtime"
 Requires-Dist: langchain_core<0.4.0,>=0.3.76; extra == "runtime"

{alita_sdk-0.3.448.dist-info → alita_sdk-0.3.450.dist-info}/RECORD

@@ -103,7 +103,7 @@ alita_sdk/runtime/toolkits/application.py,sha256=HHAKgwKOckxc7EQG-AV7rz4POOzQJKF
 alita_sdk/runtime/toolkits/artifact.py,sha256=YChNCX4QhVpaQG7Jk4TS-Wl0Aruc4slQ2K21zh9nNO0,3176
 alita_sdk/runtime/toolkits/configurations.py,sha256=kIDAlnryPQfbZyFxV-9SzN2-Vefzx06TX1BBdIIpN90,141
 alita_sdk/runtime/toolkits/datasource.py,sha256=qk78OdPoReYPCWwahfkKLbKc4pfsu-061oXRryFLP6I,2498
-alita_sdk/runtime/toolkits/mcp.py,sha256=2RpT00q6qKkXHPMkFN6wnqUHYGjYpDAmZ5dWokEXA_s,42172
+alita_sdk/runtime/toolkits/mcp.py,sha256=wa-47obeWm8WrIiaUw5BR40dsR5sLdQKZHzG7MCRnRM,37105
 alita_sdk/runtime/toolkits/prompt.py,sha256=WIpTkkVYWqIqOWR_LlSWz3ug8uO9tm5jJ7aZYdiGRn0,1192
 alita_sdk/runtime/toolkits/subgraph.py,sha256=wwUK8JjPXkGzyVZ3tAukmvST6eGbqx_U11rpnmbrvtg,2105
 alita_sdk/runtime/toolkits/tools.py,sha256=74R0SadMiUJFUe3U5_PJcxzw5m1KMMBiVfVeiowZK6I,13434
@@ -122,7 +122,7 @@ alita_sdk/runtime/tools/llm.py,sha256=iRG_wU4T01LRsjEMPZe5Uah7LiMqDc-vspwkMuQtlt
 alita_sdk/runtime/tools/loop.py,sha256=uds0WhZvwMxDVFI6MZHrcmMle637cQfBNg682iLxoJA,8335
 alita_sdk/runtime/tools/loop_output.py,sha256=U4hO9PCQgWlXwOq6jdmCGbegtAxGAPXObSxZQ3z38uk,8069
 alita_sdk/runtime/tools/mcp_inspect_tool.py,sha256=38X8euaxDbEGjcfp6ElvExZalpZun6QEr6ZEW4nU5pQ,11496
-alita_sdk/runtime/tools/mcp_remote_tool.py,sha256=0-wqqq2ZY9cgGeTd6ao17e393YgD99n3X-k2pXlaR2g,13352
+alita_sdk/runtime/tools/mcp_remote_tool.py,sha256=Bpgu5r97QzDvNert8Vv424DlolMR0rXsac3Qrg2iAig,7059
 alita_sdk/runtime/tools/mcp_server_tool.py,sha256=-xO3H6BM63KaIV1CdcQKPVE0WPigiqOgFZDX7m2_yGs,4419
 alita_sdk/runtime/tools/pgvector_search.py,sha256=NN2BGAnq4SsDHIhUcFZ8d_dbEOM8QwB0UwpsWCYruXU,11692
 alita_sdk/runtime/tools/prompt.py,sha256=nJafb_e5aOM1Rr3qGFCR-SKziU9uCsiP2okIMs9PppM,741
@@ -137,6 +137,7 @@ alita_sdk/runtime/utils/constants.py,sha256=Xntx1b_uxUzT4clwqHA_U6K8y5bBqf_4lSQw
 alita_sdk/runtime/utils/evaluate.py,sha256=iM1P8gzBLHTuSCe85_Ng_h30m52hFuGuhNXJ7kB1tgI,1872
 alita_sdk/runtime/utils/logging.py,sha256=svPyiW8ztDfhqHFITv5FBCj8UhLxz6hWcqGIY6wpJJE,3331
 alita_sdk/runtime/utils/mcp_oauth.py,sha256=Ynoa_C_G5WXL_tlcdol2wBLgQauyvIPX0isCJKsvWMs,6151
+alita_sdk/runtime/utils/mcp_sse_client.py,sha256=cSOyfnOoxVorfIePQ4k-BmOutBOeL3YDhTNxWtFREA0,16251
 alita_sdk/runtime/utils/save_dataframe.py,sha256=i-E1wp-t4wb17Zq3nA3xYwgSILjoXNizaQAA9opWvxY,1576
 alita_sdk/runtime/utils/streamlit.py,sha256=yIb4YIGH8HRAXZtZlywjxI07Xdcb5eUt7rMA-elFTdc,107261
 alita_sdk/runtime/utils/toolkit_runtime.py,sha256=MU63Fpxj0b5_r1IUUc0Q3-PN9VwL7rUxp2MRR4tmYR8,5136
@@ -360,8 +361,8 @@ alita_sdk/tools/zephyr_scale/api_wrapper.py,sha256=kT0TbmMvuKhDUZc0i7KO18O38JM9S
 alita_sdk/tools/zephyr_squad/__init__.py,sha256=0ne8XLJEQSLOWfzd2HdnqOYmQlUliKHbBED5kW_Vias,2895
 alita_sdk/tools/zephyr_squad/api_wrapper.py,sha256=kmw_xol8YIYFplBLWTqP_VKPRhL_1ItDD0_vXTe_UuI,14906
 alita_sdk/tools/zephyr_squad/zephyr_squad_cloud_client.py,sha256=R371waHsms4sllHCbijKYs90C-9Yu0sSR3N4SUfQOgU,5066
-alita_sdk-0.3.448.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-alita_sdk-0.3.448.dist-info/METADATA,sha256=FsoPHn8uhMeNXFvo4DkRBzyweY-GNMGH_-xJtnAVX2c,19071
-alita_sdk-0.3.448.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-alita_sdk-0.3.448.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
-alita_sdk-0.3.448.dist-info/RECORD,,
+alita_sdk-0.3.450.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+alita_sdk-0.3.450.dist-info/METADATA,sha256=ef9IMngx_rnicjlwr5pbK6rWNCbb8oGxTS4NbiAatuw,19101
+alita_sdk-0.3.450.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+alita_sdk-0.3.450.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
+alita_sdk-0.3.450.dist-info/RECORD,,