alita-sdk 0.3.444__tar.gz → 0.3.446__tar.gz

{alita_sdk-0.3.444 → alita_sdk-0.3.446}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.444
+Version: 0.3.446
 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

{alita_sdk-0.3.444 → alita_sdk-0.3.446}/alita_sdk/runtime/models/mcp_models.py

@@ -19,6 +19,10 @@ class McpConnectionConfig(BaseModel):
         default=None,
         description="HTTP headers for the connection (JSON object)"
     )
+    session_id: Optional[str] = Field(
+        default=None,
+        description="MCP session ID for stateful SSE servers (managed by client)"
+    )
 
     @validator('url')
     def validate_url(cls, v):

{alita_sdk-0.3.444 → alita_sdk-0.3.446}/alita_sdk/runtime/toolkits/mcp.py

@@ -340,9 +340,14 @@ class McpToolkit(BaseToolkit):
             logger.error(f"Headers must be a dictionary or JSON string, got: {type(headers)}")
             raise ValueError(f"Headers must be a dictionary or JSON string, got: {type(headers)}")
 
+        # Extract session_id from kwargs if provided
+        session_id = kwargs.get('session_id')
+        if session_id:
+            logger.info(f"[MCP Session] Using provided session ID for toolkit '{toolkit_name}': {session_id}")
+        
         # Create MCP connection configuration
         try:
-            connection_config = McpConnectionConfig(url=url, headers=parsed_headers)
+            connection_config = McpConnectionConfig(url=url, headers=parsed_headers, session_id=session_id)
         except Exception as e:
             logger.error(f"Invalid MCP connection configuration: {e}")
             raise ValueError(f"Invalid MCP connection configuration: {e}")
@@ -382,7 +387,7 @@ class McpToolkit(BaseToolkit):
             logger.info(f"Discovering tools from MCP toolkit '{toolkit_name}' at {connection_config.url}")
 
             # Use synchronous HTTP discovery for toolkit initialization
-            tool_metadata_list = cls._discover_tools_sync(
+            tool_metadata_list, session_id = cls._discover_tools_sync(
                 toolkit_name=toolkit_name,
                 connection_config=connection_config,
                 timeout=timeout
@@ -397,13 +402,15 @@ class McpToolkit(BaseToolkit):
                 ]
 
             # Create BaseTool instances from discovered metadata
+            # Use session_id from connection_config (passed from UI) instead of discovery
             for tool_metadata in tool_metadata_list:
                 server_tool = cls._create_tool_from_dict(
                     tool_dict=tool_metadata,
                     toolkit_name=toolkit_name,
                     connection_config=connection_config,
                     timeout=timeout,
-                    client=client
+                    client=client,
+                    session_id=connection_config.session_id  # Use session from connection config
                 )
 
                 if server_tool:
@@ -455,12 +462,19 @@ class McpToolkit(BaseToolkit):
         """
         all_tools = []
         
+        # Note: We don't initialize MCP session here because:
+        # 1. Discovery happens before we have OAuth tokens
+        # 2. Sessions are initialized on-demand during first tool call (when we have auth)
+        # 3. Sessions are stored in UI sessionStorage and passed back via mcp_tokens
+        session_id = None
+        
         # Discover regular tools
         tools_data = cls._discover_mcp_endpoint(
             endpoint="tools/list",
             toolkit_name=toolkit_name,
             connection_config=connection_config,
-            timeout=timeout
+            timeout=timeout,
+            session_id=session_id
         )
         all_tools.extend(tools_data)
         logger.info(f"Discovered {len(tools_data)} tools from MCP toolkit '{toolkit_name}'")
@@ -471,7 +485,8 @@ class McpToolkit(BaseToolkit):
                 endpoint="prompts/list",
                 toolkit_name=toolkit_name,
                 connection_config=connection_config,
-                timeout=timeout
+                timeout=timeout,
+                session_id=session_id
             )
             # Convert prompts to tool format
             for prompt in prompts_data:
@@ -504,7 +519,107 @@ class McpToolkit(BaseToolkit):
             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
+        return all_tools, session_id
+
+    @classmethod
+    def _initialize_mcp_session(
+        cls,
+        toolkit_name: str,
+        connection_config: McpConnectionConfig,
+        timeout: int,
+        extra_headers: Optional[Dict[str, str]] = None
+    ) -> Optional[str]:
+        """
+        Initialize an MCP session for stateful SSE servers.
+        Returns sessionId if successful, None if server doesn't require sessions.
+        
+        Note: This session should be stored by the caller (e.g., UI sessionStorage)
+        and passed back for subsequent requests. Sessions should NOT be cached
+        on the backend for security reasons (multi-tenant environment).
+        
+        Args:
+            toolkit_name: Name of the toolkit
+            connection_config: MCP connection configuration
+            timeout: Request timeout in seconds
+            extra_headers: Additional headers (e.g., Authorization) to include
+        """
+        import time
+
+        mcp_request = {
+            "jsonrpc": "2.0",
+            "id": f"initialize_{int(time.time())}",
+            "method": "initialize",
+            "params": {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {
+                    "roots": {"listChanged": True},
+                    "sampling": {}
+                },
+                "clientInfo": {
+                    "name": "Alita MCP Client",
+                    "version": "1.0.0"
+                }
+            }
+        }
+
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json, text/event-stream"
+        }
+        if connection_config.headers:
+            headers.update(connection_config.headers)
+        if extra_headers:
+            headers.update(extra_headers)
+            logger.debug(f"[MCP Session] Added extra headers: {list(extra_headers.keys())}")
+
+        try:
+            logger.info(f"[MCP Session] Attempting to initialize session for {connection_config.url}")
+            logger.debug(f"[MCP Session] Initialize request: {mcp_request}")
+            response = requests.post(
+                connection_config.url,
+                json=mcp_request,
+                headers=headers,
+                timeout=timeout
+            )
+
+            logger.info(f"[MCP Session] Initialize response status: {response.status_code}")
+            if response.status_code == 200:
+                # Parse the response to extract sessionId
+                content_type = response.headers.get('Content-Type', '')
+                logger.debug(f"[MCP Session] Response Content-Type: {content_type}")
+                
+                if 'text/event-stream' in content_type:
+                    data = cls._parse_sse_response(response.text)
+                elif 'application/json' in content_type:
+                    data = response.json()
+                else:
+                    logger.warning(f"[MCP Session] Unexpected Content-Type during initialize: {content_type}")
+                    logger.debug(f"[MCP Session] Response text: {response.text[:500]}")
+                    return None
+
+                logger.debug(f"[MCP Session] Parsed response: {data}")
+                
+                # Extract sessionId from response
+                result = data.get("result", {})
+                session_id = result.get("sessionId")
+                if session_id:
+                    logger.info(f"[MCP Session] ✓ Session initialized for '{toolkit_name}': {session_id}")
+                    logger.info(f"[MCP Session] This sessionId should be stored by the caller and sent back for subsequent requests")
+                    return session_id
+                else:
+                    logger.info(f"[MCP Session] No sessionId in initialize response for '{toolkit_name}' - server may not require sessions")
+                    logger.debug(f"[MCP Session] Full result: {result}")
+                    return None
+            else:
+                logger.warning(f"[MCP Session] Initialize returned {response.status_code} for '{toolkit_name}' - server may not support sessions")
+                logger.debug(f"[MCP Session] Response: {response.text[:500]}")
+                return None
+
+        except Exception as e:
+            logger.warning(f"[MCP Session] Failed to initialize MCP session for '{toolkit_name}': {e} - proceeding without session")
+            import traceback
+            logger.debug(f"[MCP Session] Traceback: {traceback.format_exc()}")
+            return None
 
     @classmethod
     def _discover_mcp_endpoint(
@@ -512,7 +627,8 @@ class McpToolkit(BaseToolkit):
         endpoint: str,
         toolkit_name: str,
         connection_config: McpConnectionConfig,
-        timeout: int
+        timeout: int,
+        session_id: Optional[str] = None
     ) -> List[Dict[str, Any]]:
         """
         Discover items from a specific MCP endpoint (tools/list or prompts/list).
@@ -535,10 +651,19 @@ class McpToolkit(BaseToolkit):
         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 {connection_config.url}")
+            logger.debug(f"Sending MCP {endpoint} request to {url}")
             response = requests.post(
-                connection_config.url,
+                url,
                 json=mcp_request,
                 headers=headers,
                 timeout=timeout
@@ -643,7 +768,8 @@ class McpToolkit(BaseToolkit):
         toolkit_name: str,
         connection_config: McpConnectionConfig,
         timeout: int,
-        client
+        client,
+        session_id: Optional[str] = None
     ) -> Optional[BaseTool]:
         """Create a BaseTool from a tool/prompt dictionary (from direct HTTP discovery)."""
         try:
@@ -680,7 +806,8 @@ class McpToolkit(BaseToolkit):
                 tool_timeout_sec=timeout,
                 is_prompt=is_prompt,
                 prompt_name=tool_dict.get("_mcp_prompt_name") if is_prompt else None,
-                original_tool_name=tool_dict.get('name')  # Store original name for MCP server invocation
+                original_tool_name=tool_dict.get('name'),  # Store original name for MCP server invocation
+                session_id=session_id  # Pass session ID for stateful SSE servers
             )
         except Exception as e:
             logger.error(f"Failed to create MCP tool '{tool_dict.get('name')}' from toolkit '{toolkit_name}': {e}")

{alita_sdk-0.3.444 → alita_sdk-0.3.446}/alita_sdk/runtime/toolkits/tools.py

@@ -113,13 +113,41 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                 settings = dict(tool['settings'])
                 url = settings.get('url')
                 headers = settings.get('headers')
-                token = None
+                token_data = None
+                session_id = None
                 if mcp_tokens and url:
-                    token = mcp_tokens.get(canonical_resource(url))
-                if token:
+                    canonical_url = canonical_resource(url)
+                    logger.info(f"[MCP Auth] Looking for token for URL: {url}")
+                    logger.info(f"[MCP Auth] Canonical URL: {canonical_url}")
+                    logger.info(f"[MCP Auth] Available tokens: {list(mcp_tokens.keys())}")
+                    token_data = mcp_tokens.get(canonical_url)
+                    if token_data:
+                        logger.info(f"[MCP Auth] Found token data for {canonical_url}")
+                        # Handle both old format (string) and new format (dict with access_token and session_id)
+                        if isinstance(token_data, dict):
+                            access_token = token_data.get('access_token')
+                            session_id = token_data.get('session_id')
+                            logger.info(f"[MCP Auth] Token data: access_token={'present' if access_token else 'missing'}, session_id={session_id or 'none'}")
+                        else:
+                            # Backward compatibility: treat as plain token string
+                            access_token = token_data
+                            logger.info(f"[MCP Auth] Using legacy token format (string)")
+                    else:
+                        access_token = None
+                        logger.warning(f"[MCP Auth] No token found for {canonical_url}")
+                else:
+                    access_token = None
+                    
+                if access_token:
                     merged_headers = dict(headers) if headers else {}
-                    merged_headers.setdefault('Authorization', f'Bearer {token}')
+                    merged_headers.setdefault('Authorization', f'Bearer {access_token}')
                     settings['headers'] = merged_headers
+                    logger.info(f"[MCP Auth] Added Authorization header for {url}")
+                    
+                # Pass session_id to MCP toolkit if available
+                if session_id:
+                    settings['session_id'] = session_id
+                    logger.info(f"[MCP Auth] Passing session_id to toolkit: {session_id}")
                 tools.extend(McpToolkit.get_toolkit(
                     toolkit_name=tool.get('toolkit_name', ''),
                     client=alita_client,

{alita_sdk-0.3.444 → alita_sdk-0.3.446}/alita_sdk/runtime/tools/mcp_remote_tool.py

@@ -36,6 +36,22 @@ class McpRemoteTool(McpServerTool):
     original_tool_name: Optional[str] = Field(default=None, description="Original tool name from MCP server (before optimization)")
     is_prompt: bool = False  # Flag to indicate if this is a prompt tool
     prompt_name: Optional[str] = None  # Original prompt name if this is a prompt
+    session_id: Optional[str] = Field(default=None, description="MCP session ID for stateful SSE servers")
+    
+    def model_post_init(self, __context: Any) -> None:
+        """Update metadata with session info after model initialization."""
+        super().model_post_init(__context)
+        self._update_metadata_with_session()
+    
+    def _update_metadata_with_session(self):
+        """Update the metadata dict with current session information."""
+        if self.session_id:
+            if self.metadata is None:
+                self.metadata = {}
+            self.metadata.update({
+                'mcp_session_id': self.session_id,
+                'mcp_server_url': canonical_resource(self.server_url)
+            })
     
     def __getstate__(self):
         """Custom serialization for pickle compatibility."""
@@ -112,14 +128,21 @@ class McpRemoteTool(McpServerTool):
         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 {self.server_url}")
+                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)}")
                 
-                async with session.post(self.server_url, json=mcp_request, headers=headers) as response:
+                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)
@@ -159,19 +182,68 @@ class McpRemoteTool(McpServerTool):
                             tool_name=self.name,
                         )
 
+                    # Check for errors first
+                    data = None
                     if response.status != 200:
                         error_text = await response.text()
-                        raise Exception(f"HTTP {response.status}: {error_text}")
-
-                    # Handle both JSON and SSE responses
-                    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()
+                        
+                        # Check if this is a "Missing sessionId" error
+                        if response.status == 404 and "sessionId" in error_text and not self.session_id:
+                            logger.info(f"[MCP Session] Server requires session but none provided - attempting to initialize")
+                            # Try to initialize session with current auth headers
+                            from ..toolkits.mcp import McpToolkit
+                            from ..models.mcp_models import McpConnectionConfig
+                            
+                            connection_config = McpConnectionConfig(
+                                url=self.server_url,
+                                headers=self.server_headers
+                            )
+                            
+                            # Initialize session with auth headers
+                            session_id = McpToolkit._initialize_mcp_session(
+                                toolkit_name=self.server,
+                                connection_config=connection_config,
+                                timeout=self.tool_timeout_sec,
+                                extra_headers=headers  # Include Authorization header
+                            )
+                            
+                            if session_id:
+                                # Retry the request with session
+                                self.session_id = session_id
+                                # Update metadata so callbacks can access session info
+                                self._update_metadata_with_session()
+                                separator = '&' if '?' in self.server_url else '?'
+                                retry_url = f"{self.server_url}{separator}sessionId={session_id}"
+                                logger.info(f"[MCP Session] Retrying tool call with new session: {session_id}")
+                                logger.info(f"[MCP Session] Created session - server: {canonical_resource(self.server_url)}, session_id: {session_id}")
+                                
+                                async with session.post(retry_url, json=mcp_request, headers=headers) as retry_response:
+                                    if retry_response.status != 200:
+                                        retry_error = await retry_response.text()
+                                        raise Exception(f"HTTP {retry_response.status} (after session init): {retry_error}")
+                                    
+                                    # Process successful retry response
+                                    content_type = retry_response.headers.get('Content-Type', '')
+                                    if 'text/event-stream' in content_type:
+                                        text = await retry_response.text()
+                                        data = self._parse_sse(text)
+                                    else:
+                                        data = await retry_response.json()
+                            else:
+                                raise Exception(f"HTTP {response.status}: {error_text} (session initialization failed)")
+                        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)}")
 
@@ -224,3 +296,12 @@ class McpRemoteTool(McpServerTool):
                 json_str = line[5:].strip()
                 return json.loads(json_str)
         raise ValueError("No data found in SSE response")
+    
+    def get_session_metadata(self) -> dict:
+        """Return session metadata to be included in tool responses."""
+        if self.session_id:
+            return {
+                'mcp_session_id': self.session_id,
+                'mcp_server_url': canonical_resource(self.server_url)
+            }
+        return {}

{alita_sdk-0.3.444 → alita_sdk-0.3.446}/alita_sdk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.444
+Version: 0.3.446
 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

{alita_sdk-0.3.444 → alita_sdk-0.3.446}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "alita_sdk"
-version = "0.3.444"
+version = "0.3.446"
 description = "SDK for building langchain agents using resources from Alita"
 readme = "README.md"
 requires-python = ">=3.10"