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

alita_sdk/runtime/clients/client.py

@@ -68,6 +68,7 @@ class AlitaClient:
         self.bucket_url = f"{self.base_url}{self.api_path}/artifacts/buckets/{self.project_id}"
         self.configurations_url = f'{self.base_url}{self.api_path}/integrations/integrations/default/{self.project_id}?section=configurations&unsecret=true'
         self.ai_section_url = f'{self.base_url}{self.api_path}/integrations/integrations/default/{self.project_id}?section=ai'
+        self.models_url = f'{self.base_url}{self.api_path}/configurations/models/{self.project_id}?include_shared=true'
         self.image_generation_url = f"{self.base_url}{self.llm_path}/images/generations"
         self.configurations: list = configurations or []
         self.model_timeout = kwargs.get('model_timeout', 120)
@@ -175,6 +176,20 @@ class AlitaClient:
             return resp.json()
         return []
 
+    def get_available_models(self):
+        """Get list of available models from the configurations API.
+        
+        Returns:
+            List of model dictionaries with 'name' and other properties,
+            or empty list if request fails.
+        """
+        resp = requests.get(self.models_url, headers=self.headers, verify=False)
+        if resp.ok:
+            data = resp.json()
+            # API returns {"items": [...], ...}
+            return data.get('items', [])
+        return []
+
     def get_embeddings(self, embedding_model: str) -> OpenAIEmbeddings:
         """
         Get an instance of OpenAIEmbeddings configured with the project ID and auth token.
@@ -565,7 +580,7 @@ class AlitaClient:
             monitoring_meta = tasknode_task.meta.get("monitoring", {})
             return monitoring_meta["user_id"]
         except Exception as e:
-            logger.warning(f"Error: Could not determine user ID for MCP tool: {e}")
+            logger.debug(f"Error: Could not determine user ID for MCP tool: {e}")
             return None
 
     def predict_agent(self, llm: ChatOpenAI, instructions: str = "You are a helpful assistant.",

alita_sdk/runtime/langchain/constants.py

@@ -83,4 +83,5 @@ DEFAULT_MULTIMODAL_PROMPT = """
 
 ELITEA_RS = "elitea_response"
 PRINTER = "printer"
-PRINTER_NODE_RS = "printer_output"
+PRINTER_NODE_RS = "printer_output"
+PRINTER_COMPLETED_STATE = "PRINTER_COMPLETED"

alita_sdk/runtime/langchain/langraph_agent.py

@@ -19,7 +19,7 @@ from langgraph.managed.base import is_managed_value
 from langgraph.prebuilt import InjectedStore
 from langgraph.store.base import BaseStore
 
-from .constants import PRINTER_NODE_RS, PRINTER
+from .constants import PRINTER_NODE_RS, PRINTER, PRINTER_COMPLETED_STATE
 from .mixedAgentRenderes import convert_message_to_json
 from .utils import create_state, propagate_the_input_mapping, safe_format
 from ..tools.function import FunctionTool
@@ -244,11 +244,19 @@ class PrinterNode(Runnable):
         result = {}
         logger.debug(f"Initial text pattern: {self.input_mapping}")
         mapping = propagate_the_input_mapping(self.input_mapping, [], state)
+        # for printer node we expect that all the lists will be joined into strings already
+        # Join any lists that haven't been converted yet
+        for key, value in mapping.items():
+            if isinstance(value, list):
+                mapping[key] = ', '.join(str(item) for item in value)
         if mapping.get(PRINTER) is None:
             raise ToolException(f"PrinterNode requires '{PRINTER}' field in input mapping")
         formatted_output = mapping[PRINTER]
         # add info label to the printer's output
-        if formatted_output:
+        if not formatted_output == PRINTER_COMPLETED_STATE:
+            # convert formatted output to string if it's not
+            if not isinstance(formatted_output, str):
+                formatted_output = str(formatted_output)
             formatted_output += f"\n\n-----\n*How to proceed?*\n* *to resume the pipeline - type anything...*"
         logger.debug(f"Formatted output: {formatted_output}")
         result[PRINTER_NODE_RS] = formatted_output
@@ -475,10 +483,14 @@ def create_graph(
             if toolkit_name:
                 tool_name = f"{clean_string(toolkit_name)}{TOOLKIT_SPLITTER}{tool_name}"
             logger.info(f"Node: {node_id} : {node_type} - {tool_name}")
-            if node_type in ['function', 'tool', 'loop', 'loop_from_tool', 'indexer', 'subgraph', 'pipeline', 'agent']:
+            if node_type in ['function', 'toolkit', 'mcp', 'tool', 'loop', 'loop_from_tool', 'indexer', 'subgraph', 'pipeline', 'agent']:
+                if node_type == 'mcp' and tool_name not in [tool.name for tool in tools]:
+                    # MCP is not connected and node cannot be added
+                    raise ToolException(f"MCP tool '{tool_name}' not found in the provided tools. "
+                                        f"Make sure it is connected properly. Available tools: {[tool.name for tool in tools]}")
                 for tool in tools:
                     if tool.name == tool_name:
-                        if node_type == 'function':
+                        if node_type in ['function', 'toolkit', 'mcp']:
                             lg_builder.add_node(node_id, FunctionTool(
                                 tool=tool, name=node_id, return_type='dict',
                                 output_variables=node.get('output', []),
@@ -643,6 +655,7 @@ def create_graph(
                         default_output=node.get('default_output', 'END')
                     )
                 )
+                continue
             elif node_type == 'state_modifier':
                 lg_builder.add_node(node_id, StateModifierNode(
                     template=node.get('template', ''),
@@ -661,9 +674,9 @@ def create_graph(
                 # reset printer output variable to avoid carrying over
                 reset_node_id = f"{node_id}_reset"
                 lg_builder.add_node(reset_node_id, PrinterNode(
-                    input_mapping={'printer': {'type': 'fixed', 'value': ''}}
+                    input_mapping={'printer': {'type': 'fixed', 'value': PRINTER_COMPLETED_STATE}}
                 ))
-                lg_builder.add_edge(node_id, reset_node_id)
+                lg_builder.add_conditional_edges(node_id, TransitionalEdge(reset_node_id))
                 lg_builder.add_conditional_edges(reset_node_id, TransitionalEdge(clean_string(node['transition'])))
                 continue
             if node.get('transition'):
@@ -814,35 +827,63 @@ class LangGraphAgentRunnable(CompiledStateGraph):
             input['messages'] = [convert_dict_to_message(msg) for msg in chat_history]
 
         # handler for LLM node: if no input (Chat perspective), then take last human message
+        # Track if input came from messages to handle content extraction properly
+        input_from_messages = False
         if not input.get('input'):
             if input.get('messages'):
                 input['input'] = [next((msg for msg in reversed(input['messages']) if isinstance(msg, HumanMessage)),
-                                          None)]
+                                      None)]
+                if input['input'] is not None:
+                    input_from_messages = True
 
         # Append current input to existing messages instead of overwriting
         if input.get('input'):
             if isinstance(input['input'], str):
                 current_message = input['input']
             else:
+                # input can be a list of messages or a single message object
                 current_message = input.get('input')[-1]
 
             # TODO: add handler after we add 2+ inputs (filterByType, etc.)
             if isinstance(current_message, HumanMessage):
                 current_content = current_message.content
                 if isinstance(current_content, list):
-                    text_contents = [
-                        item['text'] if isinstance(item, dict) and item.get('type') == 'text'
-                        else item if isinstance(item, str)
-                        else None
-                        for item in current_content
-                    ]
-                    text_contents = [text for text in text_contents if text is not None]
-                    input['input'] = ". ".join(text_contents)
+                    # Extract text parts and keep non-text parts (images, etc.)
+                    text_contents = []
+                    non_text_parts = []
+                    
+                    for item in current_content:
+                        if isinstance(item, dict) and item.get('type') == 'text':
+                            text_contents.append(item['text'])
+                        elif isinstance(item, str):
+                            text_contents.append(item)
+                        else:
+                            # Keep image_url and other non-text content
+                            non_text_parts.append(item)
+                    
+                    # Set input to the joined text
+                    input['input'] = ". ".join(text_contents) if text_contents else ""
+                    
+                    # If this message came from input['messages'], update or remove it
+                    if input_from_messages:
+                        if non_text_parts:
+                            # Keep the message but only with non-text content (images, etc.)
+                            current_message.content = non_text_parts
+                        else:
+                            # All content was text, remove this message from the list
+                            input['messages'] = [msg for msg in input['messages'] if msg is not current_message]
+                
                 elif isinstance(current_content, str):
                     # on regenerate case
                     input['input'] = current_content
+                    # If from messages and all content is text, remove the message
+                    if input_from_messages:
+                        input['messages'] = [msg for msg in input['messages'] if msg is not current_message]
                 else:
                     input['input'] = str(current_content)
+                    # If from messages, remove since we extracted the content
+                    if input_from_messages:
+                        input['messages'] = [msg for msg in input['messages'] if msg is not current_message]
             elif isinstance(current_message, str):
                 input['input'] = current_message
             else:
@@ -852,17 +893,30 @@ class LangGraphAgentRunnable(CompiledStateGraph):
                 input['messages'] = [convert_dict_to_message(msg) for msg in input['messages']]
                 # Append to existing messages
                 # input['messages'].append(current_message)
-            else:
-                # No existing messages, create new list
-                input['messages'] = [current_message]
+            # else:
+                # NOTE: Commented out to prevent duplicates with input['input']
+                # input['messages'] = [current_message]
+        
+        # Validate that input is not empty after all processing
+        if not input.get('input'):
+            raise RuntimeError(
+                "Empty input after processing. Cannot send empty string to LLM. "
+                "This likely means the message contained only non-text content "
+                "with no accompanying text."
+            )
+        
         logging.info(f"Input: {thread_id} - {input}")
         if self.checkpointer and self.checkpointer.get_tuple(config):
             self.update_state(config, input)
-            result = super().invoke(None, config=config, *args, **kwargs)
+            if config.pop("should_continue", False):
+                invoke_input = input
+            else:
+                invoke_input = None
+            result = super().invoke(invoke_input, config=config, *args, **kwargs)
         else:
             result = super().invoke(input, config=config, *args, **kwargs)
         try:
-            if not result.get(PRINTER_NODE_RS):
+            if result.get(PRINTER_NODE_RS) == PRINTER_COMPLETED_STATE:
                 output = next((msg.content for msg in reversed(result['messages']) if not isinstance(msg, HumanMessage)),
                               result['messages'][-1].content)
             else:

alita_sdk/runtime/langchain/utils.py

@@ -2,7 +2,7 @@ import builtins
 import json
 import logging
 import re
-from pydantic import create_model, Field
+from pydantic import create_model, Field, Json
 from typing import Tuple, TypedDict, Any, Optional, Annotated
 from langchain_core.messages import AnyMessage
 from langgraph.graph import add_messages
@@ -131,7 +131,7 @@ def parse_type(type_str):
 
 
 def create_state(data: Optional[dict] = None):
-    state_dict = {'input': str, 'router_output': str,
+    state_dict = {'input': str, 'messages': 'list[str]', 'router_output': str,
                   ELITEA_RS: str, PRINTER_NODE_RS: str}  # Always include router_output
     types_dict = {}
     if not data:
@@ -208,5 +208,21 @@ def safe_format(template, mapping):
 def create_pydantic_model(model_name: str, variables: dict[str, dict]):
     fields = {}
     for var_name, var_data in variables.items():
-        fields[var_name] = (parse_type(var_data['type']), Field(description=var_data.get('description', None)))
-    return create_model(model_name, **fields)
+        fields[var_name] = (parse_pydantic_type(var_data['type']), Field(description=var_data.get('description', None)))
+    return create_model(model_name, **fields)
+
+def parse_pydantic_type(type_name: str):
+    """
+    Helper function to parse type names into Python types.
+    Extend this function to handle custom types like 'dict' -> Json[Any].
+    """
+    type_mapping = {
+        'str': str,
+        'int': int,
+        'float': float,
+        'bool': bool,
+        'dict': Json[Any],  # Map 'dict' to Pydantic's Json type
+        'list': list,
+        'any': Any
+    }
+    return type_mapping.get(type_name, Any)

alita_sdk/runtime/toolkits/artifact.py

@@ -23,11 +23,7 @@ class ArtifactToolkit(BaseToolkit):
             # client = (Any, FieldInfo(description="Client object", required=True, autopopulate=True)),
             bucket=(str, FieldInfo(
                 description="Bucket name",
-                pattern=r'^[a-z][a-z0-9-]*$',
-                json_schema_extra={
-                    'toolkit_name': True,
-                    'max_toolkit_length': ArtifactToolkit.toolkit_max_length
-                }
+                pattern=r'^[a-z][a-z0-9-]*$'
             )),
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             # indexer settings
@@ -37,7 +33,10 @@ class ArtifactToolkit(BaseToolkit):
             embedding_model=(Optional[str], Field(default=None, description="Embedding configuration.",
                                                   json_schema_extra={'configuration_model': 'embedding'})),
 
-            __config__=ConfigDict(json_schema_extra={'metadata': {"label": "Artifact", "icon_url": None}})
+            __config__=ConfigDict(json_schema_extra={'metadata': {"label": "Artifact",
+                                                                  "icon_url": None,
+                                                                  "max_length": ArtifactToolkit.toolkit_max_length
+                                                                  }})
         )
     
     @classmethod

alita_sdk/runtime/toolkits/mcp.py

@@ -498,9 +498,12 @@ class McpToolkit(BaseToolkit):
         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:
-            logger.error(f"[MCP SSE] session_id is required for SSE servers")
-            raise ValueError("session_id is required. Frontend must generate UUID.")
+            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}")
         

alita_sdk/runtime/toolkits/tools.py

@@ -110,6 +110,7 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                     toolkit_name=tool.get('toolkit_name', ''),
                     **tool['settings']).get_tools())
             elif tool['type'] == 'mcp':
+                # remote mcp tool initialization with token injection
                 settings = dict(tool['settings'])
                 url = settings.get('url')
                 headers = settings.get('headers')

alita_sdk/runtime/tools/function.py

@@ -16,6 +16,18 @@ from ..langchain.utils import propagate_the_input_mapping
 logger = logging.getLogger(__name__)
 
 
+def replace_escaped_newlines(data):
+    """
+        Replace \\n with \n in all string values recursively.
+        Required for sanitization of state variables in code node
+    """
+    if isinstance(data, dict):
+        return {key: replace_escaped_newlines(value) for key, value in data.items()}
+    elif isinstance(data, str):
+        return data.replace('\\n', '\n')
+    else:
+        return data
+
 class FunctionTool(BaseTool):
     name: str = 'FunctionalTool'
     description: str = 'This is direct call node for tools'
@@ -30,11 +42,13 @@ class FunctionTool(BaseTool):
     def _prepare_pyodide_input(self, state: Union[str, dict, ToolCall]) -> str:
         """Prepare input for PyodideSandboxTool by injecting state into the code block."""
         # add state into the code block here since it might be changed during the execution of the code
-        state_copy = deepcopy(state)
+        state_copy = replace_escaped_newlines(deepcopy(state))
 
         del state_copy['messages']  # remove messages to avoid issues with pickling without langchain-core
         # inject state into the code block as alita_state variable
-        pyodide_predata = f"#state dict\nalita_state = {state_copy}\n"
+        state_json = json.dumps(state_copy, ensure_ascii=False)
+        pyodide_predata = f'#state dict\nimport json\nalita_state = json.loads({json.dumps(state_json)})\n'
+
         return pyodide_predata
 
     def _handle_pyodide_output(self, tool_result: Any) -> dict:
@@ -94,9 +108,7 @@ class FunctionTool(BaseTool):
         # special handler for PyodideSandboxTool
         if self._is_pyodide_tool():
             code = func_args['code']
-            func_args['code'] = (f"{self._prepare_pyodide_input(state)}\n{code}"
-                                # handle new lines in the code properly
-                                 .replace('\\n','\\\\n'))
+            func_args['code'] = f"{self._prepare_pyodide_input(state)}\n{code}"
         try:
             tool_result = self.tool.invoke(func_args, config, **kwargs)
             dispatch_custom_event(
@@ -120,7 +132,8 @@ class FunctionTool(BaseTool):
                     messages_dict = {
                         "messages": [{
                             "role": "assistant",
-                            "content": dumps(tool_result) if not isinstance(tool_result, ToolException)
+                            "content": dumps(tool_result)
+                            if not isinstance(tool_result, ToolException) and not isinstance(tool_result, str)
                             else str(tool_result)
                         }]
                     }

alita_sdk/runtime/tools/llm.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from traceback import format_exc
 from typing import Any, Optional, List, Union
@@ -132,7 +133,9 @@ class LLMNode(BaseTool):
                 struct_model = create_pydantic_model(f"LLMOutput", struct_params)
                 completion = llm_client.invoke(messages, config=config)
                 if hasattr(completion, 'tool_calls') and completion.tool_calls:
-                    new_messages, _ = self.__perform_tool_calling(completion, messages, llm_client, config)
+                    new_messages, _ = self._run_async_in_sync_context(
+                        self.__perform_tool_calling(completion, messages, llm_client, config)
+                    )
                     llm = self.__get_struct_output_model(llm_client, struct_model)
                     completion = llm.invoke(new_messages, config=config)
                     result = completion.model_dump()
@@ -155,7 +158,9 @@ class LLMNode(BaseTool):
                 # Handle both tool-calling and regular responses
                 if hasattr(completion, 'tool_calls') and completion.tool_calls:
                     # Handle iterative tool-calling and execution
-                    new_messages, current_completion = self.__perform_tool_calling(completion, messages, llm_client, config)
+                    new_messages, current_completion = self._run_async_in_sync_context(
+                        self.__perform_tool_calling(completion, messages, llm_client, config)
+                    )
 
                     output_msgs = {"messages": new_messages}
                     if self.output_variables:
@@ -190,9 +195,53 @@ class LLMNode(BaseTool):
     def _run(self, *args, **kwargs):
         # Legacy support for old interface
         return self.invoke(kwargs, **kwargs)
+    
+    def _run_async_in_sync_context(self, coro):
+        """Run async coroutine from sync context.
+        
+        For MCP tools with persistent sessions, we reuse the same event loop
+        that was used to create the MCP client and sessions (set by CLI).
+        """
+        try:
+            loop = asyncio.get_running_loop()
+            # Already in async context - run in thread with new loop
+            import threading
+            
+            result_container = []
+            
+            def run_in_thread():
+                new_loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(new_loop)
+                try:
+                    result_container.append(new_loop.run_until_complete(coro))
+                finally:
+                    new_loop.close()
+            
+            thread = threading.Thread(target=run_in_thread)
+            thread.start()
+            thread.join()
+            return result_container[0] if result_container else None
+            
+        except RuntimeError:
+            # No event loop running - use/create persistent loop
+            # This loop is shared with MCP session creation for stateful tools
+            if not hasattr(self.__class__, '_persistent_loop') or \
+               self.__class__._persistent_loop is None or \
+               self.__class__._persistent_loop.is_closed():
+                self.__class__._persistent_loop = asyncio.new_event_loop()
+                logger.debug("Created persistent event loop for async tools")
+            
+            loop = self.__class__._persistent_loop
+            asyncio.set_event_loop(loop)
+            return loop.run_until_complete(coro)
+    
+    async def _arun(self, *args, **kwargs):
+        # Legacy async support
+        return self.invoke(kwargs, **kwargs)
 
-    def __perform_tool_calling(self, completion, messages, llm_client, config):
+    async def __perform_tool_calling(self, completion, messages, llm_client, config):
         # Handle iterative tool-calling and execution
+        logger.info(f"__perform_tool_calling called with {len(completion.tool_calls) if hasattr(completion, 'tool_calls') else 0} tool calls")
         new_messages = messages + [completion]
         iteration = 0
 
@@ -230,9 +279,16 @@ class LLMNode(BaseTool):
                 if tool_to_execute:
                     try:
                         logger.info(f"Executing tool '{tool_name}' with args: {tool_args}")
-                        # Pass the underlying config to the tool execution invoke method
-                        # since it may be another agent, graph, etc. to see it properly in thinking steps
-                        tool_result = tool_to_execute.invoke(tool_args, config=config)
+                        
+                        # Try async invoke first (for MCP tools), fallback to sync
+                        tool_result = None
+                        try:
+                            # Try async invocation first
+                            tool_result = await tool_to_execute.ainvoke(tool_args, config=config)
+                        except NotImplementedError:
+                            # Tool doesn't support async, use sync invoke
+                            logger.debug(f"Tool '{tool_name}' doesn't support async, using sync invoke")
+                            tool_result = tool_to_execute.invoke(tool_args, config=config)
 
                         # Create tool message with result - preserve structured content
                         from langchain_core.messages import ToolMessage
@@ -256,7 +312,9 @@ class LLMNode(BaseTool):
                         new_messages.append(tool_message)
 
                     except Exception as e:
-                        logger.error(f"Error executing tool '{tool_name}': {e}")
+                        import traceback
+                        error_details = traceback.format_exc()
+                        logger.error(f"Error executing tool '{tool_name}': {e}\n{error_details}")
                         # Create error tool message
                         from langchain_core.messages import ToolMessage
                         tool_message = ToolMessage(

alita_sdk/runtime/tools/vectorstore_base.py

@@ -1,9 +1,9 @@
 import json
-import math
 from collections import OrderedDict
 from logging import getLogger
 from typing import Any, Optional, List, Dict, Generator
 
+import math
 from langchain_core.documents import Document
 from langchain_core.messages import HumanMessage
 from langchain_core.tools import ToolException
@@ -12,7 +12,7 @@ from pydantic import BaseModel, model_validator, Field
 
 from alita_sdk.tools.elitea_base import BaseToolApiWrapper
 from alita_sdk.tools.vector_adapters.VectorStoreAdapter import VectorStoreAdapterFactory
-from ..utils.logging import dispatch_custom_event
+from ...runtime.utils.utils import IndexerKeywords
 
 logger = getLogger(__name__)
 
@@ -222,6 +222,21 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
             raise RuntimeError(f"Multiple index_meta documents found: {index_metas}")
         return index_metas[0] if index_metas else None
 
+    def get_indexed_count(self, index_name: str) -> int:
+        from sqlalchemy.orm import Session
+        from sqlalchemy import func, or_
+
+        with Session(self.vectorstore.session_maker.bind) as session:
+            return session.query(
+                self.vectorstore.EmbeddingStore.id,
+            ).filter(
+                func.jsonb_extract_path_text(self.vectorstore.EmbeddingStore.cmetadata, 'collection') == index_name,
+                or_(
+                    func.jsonb_extract_path_text(self.vectorstore.EmbeddingStore.cmetadata, 'type').is_(None),
+                    func.jsonb_extract_path_text(self.vectorstore.EmbeddingStore.cmetadata, 'type') != IndexerKeywords.INDEX_META_TYPE.value
+                )
+            ).count()
+
     def _clean_collection(self, index_name: str = ''):
         """
         Clean the vectorstore collection by deleting all indexed data.

alita_sdk/runtime/utils/mcp_sse_client.py

@@ -65,6 +65,53 @@ class McpSseClient:
             
             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}")
@@ -248,18 +295,29 @@ class McpSseClient:
         """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:
-                pass
+            except (asyncio.CancelledError, Exception) as e:
+                logger.debug(f"[MCP SSE Client] Stream task cleanup: {e}")
         
-        if self._stream_response:
-            self._stream_response.close()
+        # 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}")
         
-        if self._stream_session:
-            await self._stream_session.close()
+        # 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")
     

alita_sdk/tools/ado/repos/__init__.py

@@ -63,6 +63,7 @@ class AzureDevOpsReposToolkit(BaseToolkit):
                     "icon_url": "ado-repos-icon.svg",
                     "categories": ["code repositories"],
                     "extra_categories": ["code", "repository", "version control"],
+                    "max_length": AzureDevOpsReposToolkit.toolkit_max_length
                 }}}
         )
 

alita_sdk/tools/ado/test_plan/__init__.py

@@ -27,7 +27,6 @@ class AzureDevOpsPlansToolkit(BaseToolkit):
         AzureDevOpsPlansToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
         m = create_model(
             name_alias,
-            name=(str, Field(description="Toolkit name", json_schema_extra={'toolkit_name': True, 'max_toolkit_length': AzureDevOpsPlansToolkit.toolkit_max_length})),
             ado_configuration=(AdoConfiguration, Field(description="Ado configuration", json_schema_extra={'configuration_types': ['ado']})),
             limit=(Optional[int], Field(description="ADO plans limit used for limitation of the list with results", default=5)),
             # indexer settings
@@ -40,6 +39,7 @@ class AzureDevOpsPlansToolkit(BaseToolkit):
                 {
                     "label": "ADO plans",
                     "icon_url": "ado-plans.svg",
+                    "max_length": AzureDevOpsPlansToolkit.toolkit_max_length,
                     "categories": ["test management"],
                     "extra_categories": ["test case management", "qa"],
                     "sections": {