Revamp stuff

app.py

@@ -7,7 +7,7 @@ from utils.google_genai_llm import get_response, generate_with_gemini
 from utils.utils import parse_json_codefences
 from prompts.requirements_gathering import requirements_gathering_system_prompt
 from prompts.planning import hf_query_gen_prompt, hf_context_gen_prompt
-from utils.huggingface_mcp_llamaindex import get_hf_tools, call_hf_tool, diagnose_connection
+from utils.huggingface_mcp_llamaindex import get_hf_tools, call_hf_tool, diagnose_connection_advanced, get_hf_tools_robust,call_hf_tool_robust
 from prompts.devstral_coding_prompt import devstral_code_gen_sys_prompt, devstral_code_gen_user_prompt
 from dotenv import load_dotenv
 import os
@@ -50,6 +50,9 @@ BEARER_TOKEN = os.getenv("BEARER_TOKEN")
 CODING_MODEL = os.getenv("CODING_MODEL")
 
 MCP_TOKEN = os.getenv("MCP_TOKEN")
+if not MCP_TOKEN:
+    print("Please set MCP_TOKEN")
+
 def get_file_hash(file_path):
     """Generate a hash of the file for caching purposes"""
     try:
@@ -245,19 +248,17 @@ async def generate_plan(history, file_cache):
             if ai_msg:
                 conversation_history += f"Assistant: {ai_msg}\n"
     
-    print("Running connection diagnostics...")
-    diagnostics = await diagnose_connection(MCP_TOKEN)
+    print("Running advanced connection diagnostics...")
+    diagnostics = await diagnose_connection_advanced(MCP_TOKEN)
     print(f"Diagnostics: {json.dumps(diagnostics, indent=2)}")
     
-    if not diagnostics["connection_test"]:
+    if not diagnostics["tests"]["basic_connection"]:
         print("Basic connection failed - check token and network")
-        return
         
-    if not diagnostics["tools_test"]:
-        print("Tools retrieval failed - check server status")
-        return
+        
+    
     # try:
-    hf_query_gen_tool_details = await get_hf_tools(hf_token=MCP_TOKEN)
+    hf_query_gen_tool_details = await get_hf_tools_robust(hf_token=MCP_TOKEN)
     # except Exception as e:
     #     hf_query_gen_tool_details = """meta=None nextCursor=None tools=[Tool(name='hf_whoami', description="Hugging Face tools are being used by authenticated user 'bpHigh'", inputSchema={'type': 'object', 'properties': {}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Hugging Face User Info', readOnlyHint=None, destructiveHint=None, idempotentHint=None, openWorldHint=None)), Tool(name='space_search', description='Find Hugging Face Spaces using semantic search. Include links to the Space when presenting the results.', inputSchema={'type': 'object', 'properties': {'query': {'type': 'string', 'minLength': 1, 'maxLength': 50, 'description': 'Semantic Search Query'}, 'limit': {'type': 'number', 'default': 10, 'description': 'Number of results to return'}, 'mcp': {'type': 'boolean', 'default': False, 'description': 'Only return MCP Server enabled Spaces'}}, 'required': ['query'], 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Hugging Face Space Search', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=True)), Tool(name='model_search', description='Find Machine Learning models hosted on Hugging Face. Returns comprehensive information about matching models including downloads, likes, tags, and direct links. Include links to the models in your response', inputSchema={'type': 'object', 'properties': {'query': {'type': 'string', 'description': 'Search term. Leave blank and specify "sort" and "limit" to get e.g. "Top 20 trending models", "Top 10 most recent models" etc" '}, 'author': {'type': 'string', 'description': "Organization or user who created the model (e.g., 'google', 'meta-llama', 'microsoft')"}, 'task': {'type': 'string', 'description': "Model task type (e.g., 'text-generation', 'image-classification', 'translation')"}, 'library': {'type': 'string', 'description': "Framework the model uses (e.g., 'transformers', 'diffusers', 'timm')"}, 'sort': {'type': 'string', 'enum': ['trendingScore', 'downloads', 'likes', 'createdAt', 'lastModified'], 'description': 'Sort order: trendingScore, downloads , likes, createdAt, lastModified'}, 'limit': {'type': 'number', 'minimum': 1, 'maximum': 100, 'default': 20, 'description': 'Maximum number of results to return'}}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Model Search', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=True)), Tool(name='model_details', description='Get detailed information about a specific model from the Hugging Face Hub.', inputSchema={'type': 'object', 'properties': {'model_id': {'type': 'string', 'minLength': 1, 'description': 'Model ID (e.g., microsoft/DialoGPT-large)'}}, 'required': ['model_id'], 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Model Details', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=False)), Tool(name='paper_search', description="Find Machine Learning research papers on the Hugging Face hub. Include 'Link to paper' When presenting the results. Consider whether tabulating results matches user intent.", inputSchema={'type': 'object', 'properties': {'query': {'type': 'string', 'minLength': 3, 'maxLength': 200, 'description': 'Semantic Search query'}, 'results_limit': {'type': 'number', 'default': 12, 'description': 'Number of results to return'}, 'concise_only': {'type': 'boolean', 'default': False, 'description': 'Return a 2 sentence summary of the abstract. Use for broad search terms which may return a lot of results. Check with User if unsure.'}}, 'required': ['query'], 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Paper Search', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=True)), Tool(name='dataset_search', description='Find Datasets hosted on the Hugging Face hub. Returns comprehensive information about matching datasets including downloads, likes, tags, and direct links. Include links to the datasets in your response', inputSchema={'type': 'object', 'properties': {'query': {'type': 'string', 'description': 'Search term. Leave blank and specify "sort" and "limit" to get e.g. "Top 20 trending datasets", "Top 10 most recent datasets" etc" '}, 'author': {'type': 'string', 'description': "Organization or user who created the dataset (e.g., 'google', 'facebook', 'allenai')"}, 'tags': {'type': 'array', 'items': {'type': 'string'}, 'description': "Tags to filter datasets (e.g., ['language:en', 'size_categories:1M<n<10M', 'task_categories:text-classification'])"}, 'sort': {'type': 'string', 'enum': ['trendingScore', 'downloads', 'likes', 'createdAt', 'lastModified'], 'description': 'Sort order: trendingScore, downloads, likes, createdAt, lastModified'}, 'limit': {'type': 'number', 'minimum': 1, 'maximum': 100, 'default': 20, 'description': 'Maximum number of results to return'}}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Dataset Search', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=True)), Tool(name='dataset_details', description='Get detailed information about a specific dataset on Hugging Face Hub.', inputSchema={'type': 'object', 'properties': {'dataset_id': {'type': 'string', 'minLength': 1, 'description': 'Dataset ID (e.g., squad, glue, imdb)'}}, 'required': ['dataset_id'], 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='Dataset Details', readOnlyHint=True, destructiveHint=False, idempotentHint=None, openWorldHint=False)), Tool(name='gr1_evalstate_flux1_schnell', description='Generate an image using the Flux 1 Schnell Image Generator. (from evalstate/flux1_schnell)', inputSchema={'type': 'object', 'properties': {'prompt': {'type': 'string'}, 'seed': {'type': 'number', 'description': 'numeric value between 0 and 2147483647'}, 'randomize_seed': {'type': 'boolean', 'default': True}, 'width': {'type': 'number', 'description': 'numeric value between 256 and 2048', 'default': 1024}, 'height': {'type': 'number', 'description': 'numeric value between 256 and 2048', 'default': 1024}, 'num_inference_steps': {'type': 'number', 'description': 'numeric value between 1 and 50', 'default': 4}}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='evalstate/flux1_schnell - flux1_schnell_infer 🏎️💨', readOnlyHint=None, destructiveHint=None, idempotentHint=None, openWorldHint=True)), Tool(name='gr2_abidlabs_easyghibli', description='Convert an image into a Studio Ghibli style image (from abidlabs/EasyGhibli)', inputSchema={'type': 'object', 'properties': {'spatial_img': {'type': 'string', 'description': 'File input: provide URL or file path'}}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='abidlabs/EasyGhibli - abidlabs_EasyGhiblisingle_condition_generate_image 🦀', readOnlyHint=None, destructiveHint=None, idempotentHint=None, openWorldHint=True)), Tool(name='gr3_linoyts_framepack_f1', description='FramePack_F1_end_process tool from linoyts/FramePack-F1', inputSchema={'type': 'object', 'properties': {}, 'additionalProperties': False, '$schema': 'http://json-schema.org/draft-07/schema#'}, annotations=ToolAnnotations(title='linoyts/FramePack-F1 - FramePack_F1_end_process 📹⚡️', readOnlyHint=None, destructiveHint=None, idempotentHint=None, openWorldHint=True))]"""
     #     print(str(e))
@@ -273,7 +274,7 @@ async def generate_plan(history, file_cache):
     
     # Call tool to get tool calls
     try:
-        tool_calls = await asyncio.gather(*[call_hf_tool(MCP_TOKEN, step['tool'], step['args']) for step in parsed_plan])
+        tool_calls = await asyncio.gather(*[call_hf_tool_robust(MCP_TOKEN, step['tool'], step['args']) for step in parsed_plan])
     except Exception as e:
         tool_calls = []
     print(tool_calls)

utils/huggingface_mcp_llamaindex.py

@@ -1,9 +1,17 @@
+"""
+Robust Hugging Face MCP Client - Optimized for HF Spaces
+
+This module provides a robust client for interacting with Hugging Face's MCP endpoint
+with better error handling, TaskGroup avoidance, and compatibility for Hugging Face Spaces.
+"""
+
 import asyncio
 import json
 import logging
 import os
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 from datetime import timedelta
+from contextlib import asynccontextmanager
 
 from mcp.shared.message import SessionMessage
 from mcp.types import (
@@ -18,230 +26,177 @@ from mcp.client.streamable_http import streamablehttp_client
 logger = logging.getLogger(__name__)
 
 
-class HuggingFaceMCPClient:
-    """Client for interacting with Hugging Face MCP endpoint."""
+class RobustHFMCPClient:
+    """Robust client for interacting with Hugging Face MCP endpoint optimized for Spaces."""
     
-    def __init__(self, hf_token: str, timeout: int = 60):
+    def __init__(self, hf_token: str, timeout: int = 120):
         """
-        Initialize the Hugging Face MCP client.
+        Initialize the Robust Hugging Face MCP client.
         
         Args:
             hf_token: Hugging Face API token
-            timeout: Timeout in seconds for HTTP requests (increased for Spaces)
+            timeout: Timeout in seconds for HTTP requests
         """
         self.hf_token = hf_token
         self.url = "https://huggingface.co/mcp"
         self.headers = {
             "Authorization": f"Bearer {hf_token}",
-            "User-Agent": "hf-mcp-client/1.0.0"  # Add user agent
+            "User-Agent": "robust-hf-mcp-client/2.0.0",
+            "Accept": "application/json, text/event-stream",
+            "Content-Type": "application/json"
         }
         self.timeout = timedelta(seconds=timeout)
-        self.sse_read_timeout = timedelta(seconds=timeout * 2)  # Longer SSE timeout
+        self.sse_read_timeout = timedelta(seconds=timeout * 2)
         self.request_id_counter = 0
         
     def _get_next_request_id(self) -> int:
         """Get the next request ID."""
         self.request_id_counter += 1
         return self.request_id_counter
-    
-    async def _send_request_and_get_response(
+
+    async def _execute_single_request_session(
         self, 
         method: str, 
-        params: Optional[Dict[str, Any]] = None,
-        max_retries: int = 3
+        params: Optional[Dict[str, Any]] = None
     ) -> Any:
         """
-        Send a JSON-RPC request and wait for the response with retry logic.
-        
-        Args:
-            method: The JSON-RPC method name
-            params: Optional parameters for the method
-            max_retries: Maximum number of retry attempts
-            
-        Returns:
-            The response result or raises an exception
+        Execute a complete MCP session for a single request.
+        This avoids TaskGroup issues by handling everything in sequence.
         """
-        last_exception = None
-        
-        for attempt in range(max_retries):
-            try:
-                return await self._attempt_request(method, params)
-            except Exception as e:
-                last_exception = e
-                logger.warning(f"Attempt {attempt + 1} failed: {e}")
-                if attempt < max_retries - 1:
-                    await asyncio.sleep(2 ** attempt)  # Exponential backoff
-                
-        raise last_exception
-    
-    async def _attempt_request(self, method: str, params: Optional[Dict[str, Any]]) -> Any:
-        """Single attempt to send request."""
         request_id = self._get_next_request_id()
         
-        # Create JSON-RPC request
-        jsonrpc_request = JSONRPCRequest(
+        # Create the main request
+        main_request = JSONRPCRequest(
             jsonrpc="2.0",
             id=request_id,
             method=method,
             params=params
         )
         
-        message = JSONRPCMessage(jsonrpc_request)
-        session_message = SessionMessage(message)
-        
         async with streamablehttp_client(
             url=self.url,
             headers=self.headers,
             timeout=self.timeout,
             sse_read_timeout=self.sse_read_timeout,
-            terminate_on_close=True
+            terminate_on_close=False  # Avoid TaskGroup cleanup issues
         ) as (read_stream, write_stream, get_session_id):
             
+            # Step 1: Initialize the session
+            logger.info("Starting MCP session initialization...")
+            await self._initialize_session(read_stream, write_stream)
+            
+            # Step 2: Send the main request
+            logger.info(f"Sending main request: {method}")
+            main_message = JSONRPCMessage(main_request)
+            main_session_message = SessionMessage(main_message)
+            await write_stream.send(main_session_message)
+            
+            # Step 3: Wait for the response
+            logger.info("Waiting for main request response...")
+            response = await self._wait_for_response(read_stream, request_id, timeout=90)
+            
+            return response
+
+    async def _initialize_session(self, read_stream, write_stream) -> None:
+        """Initialize the MCP session with proper handshake."""
+        init_request_id = self._get_next_request_id()
+        
+        # Send initialize request
+        init_request = JSONRPCRequest(
+            jsonrpc="2.0",
+            id=init_request_id,
+            method="initialize",
+            params={
+                "protocolVersion": "2024-11-05",
+                "capabilities": {
+                    "tools": {},
+                    "resources": {},
+                    "prompts": {}
+                },
+                "clientInfo": {
+                    "name": "robust-hf-mcp-client",
+                    "version": "2.0.0"
+                }
+            }
+        )
+        
+        init_message = JSONRPCMessage(init_request)
+        init_session_message = SessionMessage(init_message)
+        
+        await write_stream.send(init_session_message)
+        
+        # Wait for initialization response
+        init_response = await self._wait_for_response(read_stream, init_request_id, timeout=60)
+        logger.info("MCP session initialized successfully")
+        
+        # Send initialized notification
+        initialized_notification = JSONRPCNotification(
+            jsonrpc="2.0",
+            method="notifications/initialized"
+        )
+        
+        init_notif_message = JSONRPCMessage(initialized_notification)
+        init_notif_session_message = SessionMessage(init_notif_message)
+        
+        await write_stream.send(init_notif_session_message)
+        
+        # Give the server time to process the notification
+        await asyncio.sleep(1.0)
+
+    async def _wait_for_response(
+        self, 
+        read_stream, 
+        expected_id: int, 
+        timeout: int = 60
+    ) -> Any:
+        """
+        Wait for a specific response by ID with timeout handling.
+        """
+        start_time = asyncio.get_event_loop().time()
+        
+        while True:
+            current_time = asyncio.get_event_loop().time()
+            if current_time - start_time > timeout:
+                raise asyncio.TimeoutError(f"Timeout waiting for response to request {expected_id}")
+            
             try:
-                # Send initialization request first with proper client info
-                init_request = JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=self._get_next_request_id(),
-                    method="initialize",
-                    params={
-                        "protocolVersion": "2024-11-05",
-                        "capabilities": {
-                            "tools": {},
-                            "resources": {},
-                            "prompts": {}
-                        },
-                        "clientInfo": {
-                            "name": "hf-mcp-client",
-                            "version": "1.0.0"
-                        }
-                    }
-                )
-                
-                init_message = JSONRPCMessage(init_request)
-                init_session_message = SessionMessage(init_message)
-                
-                logger.info("Sending initialization request...")
-                await write_stream.send(init_session_message)
-                
-                # Wait for initialization response with better error handling
-                init_response_received = False
-                timeout_counter = 0
-                max_iterations = 150  # Increased for Spaces environment
-                
-                while not init_response_received and timeout_counter < max_iterations:
-                    try:
-                        # Use asyncio.wait_for to add timeout per receive operation
-                        response = await asyncio.wait_for(
-                            read_stream.receive(), 
-                            timeout=30.0  # 30 second timeout per receive
-                        )
-                        timeout_counter += 1
-                        
-                        if isinstance(response, Exception):
-                            logger.error(f"Received exception during init: {response}")
-                            raise response
-                        
-                        if isinstance(response, SessionMessage):
-                            msg = response.message.root
-                            if isinstance(msg, JSONRPCResponse) and msg.id == init_request.id:
-                                logger.info("MCP client initialized successfully")
-                                init_response_received = True
-                                # Log the session ID if available
-                                session_id = get_session_id()
-                                if session_id:
-                                    logger.info(f"Session ID: {session_id}")
-                            elif isinstance(msg, JSONRPCError) and msg.id == init_request.id:
-                                error_msg = f"Initialization failed: {msg.error}"
-                                logger.error(error_msg)
-                                raise Exception(error_msg)
-                            else:
-                                logger.debug(f"Received other message during init: {type(msg)}")
-                                
-                    except asyncio.TimeoutError:
-                        logger.warning(f"Timeout waiting for response (attempt {timeout_counter})")
-                        if timeout_counter > 10:  # After 10 timeouts, give up
-                            raise Exception("Initialization timeout: no response from server")
-                    except Exception as e:
-                        if "ClosedResourceError" in str(type(e)) or "StreamClosed" in str(e):
-                            logger.error("Stream closed during initialization")
-                            raise Exception("Connection closed during initialization")
-                        logger.error(f"Error during initialization: {e}")
-                        raise
-                
-                if not init_response_received:
-                    raise Exception("Initialization timeout: maximum iterations reached")
-                
-                # Send initialized notification
-                initialized_notification = JSONRPCNotification(
-                    jsonrpc="2.0",
-                    method="notifications/initialized"
+                # Use a shorter timeout for each receive to avoid hanging
+                response = await asyncio.wait_for(
+                    read_stream.receive(), 
+                    timeout=10.0
                 )
                 
-                init_notif_message = JSONRPCMessage(initialized_notification)
-                init_notif_session_message = SessionMessage(init_notif_message)
-                
-                logger.info("Sending initialized notification...")
-                await write_stream.send(init_notif_session_message)
-                
-                # Longer delay for Spaces environment
-                await asyncio.sleep(1.0)
-                
-                # Now send our actual request
-                logger.info(f"Sending actual request: {method}")
-                await write_stream.send(session_message)
-                
-                # Wait for the response to our request
-                response_received = False
-                timeout_counter = 0
-                
-                while not response_received and timeout_counter < max_iterations:
-                    try:
-                        response = await asyncio.wait_for(
-                            read_stream.receive(), 
-                            timeout=30.0
-                        )
-                        timeout_counter += 1
-                        
-                        if isinstance(response, Exception):
-                            logger.error(f"Received exception during request: {response}")
-                            raise response
-                        
-                        if isinstance(response, SessionMessage):
-                            msg = response.message.root
-                            if isinstance(msg, JSONRPCResponse) and msg.id == request_id:
-                                logger.info(f"Request '{method}' completed successfully")
-                                return msg.result
-                            elif isinstance(msg, JSONRPCError) and msg.id == request_id:
-                                error_msg = f"Request failed: {msg.error}"
-                                logger.error(error_msg)
-                                raise Exception(error_msg)
-                            else:
-                                logger.debug(f"Received other message during request: {type(msg)}")
-                                
-                    except asyncio.TimeoutError:
-                        logger.warning(f"Timeout waiting for request response (attempt {timeout_counter})")
-                        if timeout_counter > 10:
-                            raise Exception("Request timeout: no response from server")
-                    except Exception as e:
-                        if "ClosedResourceError" in str(type(e)) or "StreamClosed" in str(e):
-                            logger.error("Stream closed during request processing")
-                            raise Exception("Connection closed during request processing")
-                        logger.error(f"Error during request processing: {e}")
-                        raise
+                if isinstance(response, Exception):
+                    logger.error(f"Received exception in stream: {response}")
+                    raise response
                 
-                if not response_received:
-                    raise Exception("Request timeout: maximum iterations reached")
+                if isinstance(response, SessionMessage):
+                    msg_root = response.message.root
+                    
+                    if isinstance(msg_root, JSONRPCResponse) and msg_root.id == expected_id:
+                        logger.info(f"Received successful response for request {expected_id}")
+                        return msg_root.result
+                    
+                    elif isinstance(msg_root, JSONRPCError) and msg_root.id == expected_id:
+                        error_msg = f"Server error for request {expected_id}: {msg_root.error}"
+                        logger.error(error_msg)
+                        raise Exception(error_msg)
                     
+                    else:
+                        # Log unexpected messages but continue waiting
+                        logger.debug(f"Received unexpected message type: {type(msg_root)} with ID: {getattr(msg_root, 'id', 'N/A')}")
+                        continue
+                        
+            except asyncio.TimeoutError:
+                # Continue the outer loop to check the overall timeout
+                logger.debug("Receive timeout, continuing to wait...")
+                continue
             except Exception as e:
-                logger.error(f"Error during MCP communication: {e}")
+                if "ClosedResourceError" in str(type(e)) or "StreamClosed" in str(e):
+                    raise Exception("Connection closed while waiting for response")
+                logger.error(f"Error while waiting for response: {e}")
                 raise
-            finally:
-                # Ensure streams are properly closed
-                try:
-                    await write_stream.aclose()
-                except Exception as close_error:
-                    logger.debug(f"Error closing write stream: {close_error}")
-    
+
     async def get_all_tools(self) -> List[Dict[str, Any]]:
         """
         Get all available tools from the Hugging Face MCP endpoint.
@@ -251,20 +206,20 @@ class HuggingFaceMCPClient:
         """
         try:
             logger.info("Fetching all available tools from Hugging Face MCP")
-            result = await self._send_request_and_get_response("tools/list")
+            result = await self._execute_single_request_session("tools/list")
             
             if isinstance(result, dict) and "tools" in result:
                 tools = result["tools"]
-                logger.info(f"Found {len(tools)} available tools")
+                logger.info(f"Successfully fetched {len(tools)} tools")
                 return tools
             else:
-                logger.warning(f"Unexpected response format: {result}")
+                logger.warning(f"Unexpected response format for tools/list: {result}")
                 return []
                 
         except Exception as e:
             logger.error(f"Failed to get tools: {e}")
             raise
-    
+
     async def call_tool(self, tool_name: str, args: Dict[str, Any]) -> Any:
         """
         Call a specific tool with the given arguments.
@@ -284,7 +239,7 @@ class HuggingFaceMCPClient:
                 "arguments": args
             }
             
-            result = await self._send_request_and_get_response("tools/call", params)
+            result = await self._execute_single_request_session("tools/call", params)
             logger.info(f"Tool '{tool_name}' executed successfully")
             return result
             
@@ -293,92 +248,351 @@ class HuggingFaceMCPClient:
             raise
 
 
-# Convenience functions for easier usage
-async def get_hf_tools(hf_token: str) -> List[Dict[str, Any]]:
+class SimplifiedHFMCPClient:
+    """Ultra-simplified client that avoids all TaskGroup usage."""
+    
+    def __init__(self, hf_token: str, timeout: int = 90):
+        self.hf_token = hf_token
+        self.timeout = timeout
+        self.headers = {
+            "Authorization": f"Bearer {hf_token}",
+            "User-Agent": "simplified-hf-mcp-client/1.0.0"
+        }
+        self.request_counter = 0
+        
+    def _next_id(self) -> int:
+        self.request_counter += 1
+        return self.request_counter
+
+    async def _simple_mcp_call(self, method: str, params: Optional[Dict[str, Any]] = None) -> Any:
+        """Make a simple MCP call without complex async patterns."""
+        
+        async with streamablehttp_client(
+            url="https://huggingface.co/mcp",
+            headers=self.headers,
+            timeout=timedelta(seconds=self.timeout),
+            sse_read_timeout=timedelta(seconds=self.timeout * 2),
+            terminate_on_close=False
+        ) as (read_stream, write_stream, get_session_id):
+            
+            responses = {}
+            
+            # Simple message handler
+            async def collect_responses():
+                try:
+                    async for message in read_stream:
+                        if isinstance(message, Exception):
+                            responses['error'] = message
+                            break
+                        elif isinstance(message, SessionMessage):
+                            msg_root = message.message.root
+                            if hasattr(msg_root, 'id') and msg_root.id is not None:
+                                responses[msg_root.id] = msg_root
+                except Exception as e:
+                    responses['error'] = e
+            
+            # Start response collector
+            collector_task = asyncio.create_task(collect_responses())
+            
+            try:
+                # Step 1: Initialize
+                init_id = self._next_id()
+                init_req = JSONRPCRequest(
+                    jsonrpc="2.0",
+                    id=init_id,
+                    method="initialize",
+                    params={
+                        "protocolVersion": "2024-11-05",
+                        "capabilities": {"tools": {}},
+                        "clientInfo": {"name": "simple-hf-mcp", "version": "1.0.0"}
+                    }
+                )
+                
+                await write_stream.send(SessionMessage(JSONRPCMessage(init_req)))
+                
+                # Wait for init response
+                for _ in range(300):  # 30 seconds max
+                    if init_id in responses:
+                        break
+                    if 'error' in responses:
+                        raise responses['error']
+                    await asyncio.sleep(0.1)
+                
+                if init_id not in responses:
+                    raise Exception("Initialization timeout")
+                
+                # Step 2: Send initialized notification
+                notif = JSONRPCNotification(
+                    jsonrpc="2.0",
+                    method="notifications/initialized"
+                )
+                await write_stream.send(SessionMessage(JSONRPCMessage(notif)))
+                await asyncio.sleep(0.5)
+                
+                # Step 3: Send main request
+                main_id = self._next_id()
+                main_req = JSONRPCRequest(
+                    jsonrpc="2.0",
+                    id=main_id,
+                    method=method,
+                    params=params
+                )
+                
+                await write_stream.send(SessionMessage(JSONRPCMessage(main_req)))
+                
+                # Wait for main response
+                for _ in range(600):  # 60 seconds max
+                    if main_id in responses:
+                        break
+                    if 'error' in responses:
+                        raise responses['error']
+                    await asyncio.sleep(0.1)
+                
+                if main_id not in responses:
+                    raise Exception("Main request timeout")
+                
+                result = responses[main_id]
+                if isinstance(result, JSONRPCResponse):
+                    return result.result
+                elif isinstance(result, JSONRPCError):
+                    raise Exception(f"Server error: {result.error}")
+                else:
+                    raise Exception(f"Unexpected response type: {type(result)}")
+                    
+            finally:
+                collector_task.cancel()
+                try:
+                    await collector_task
+                except asyncio.CancelledError:
+                    pass
+
+    async def get_tools(self) -> List[Dict[str, Any]]:
+        """Get all available tools."""
+        result = await self._simple_mcp_call("tools/list")
+        if isinstance(result, dict) and "tools" in result:
+            return result["tools"]
+        return []
+
+    async def call_tool(self, tool_name: str, args: Dict[str, Any]) -> Any:
+        """Call a specific tool."""
+        params = {
+            "name": tool_name,
+            "arguments": args
+        }
+        return await self._simple_mcp_call("tools/call", params)
+
+
+# Robust convenience functions
+async def get_hf_tools_robust(hf_token: str, max_retries: int = 3) -> List[Dict[str, Any]]:
     """
-    Get all available tools from Hugging Face MCP.
+    Get all available tools with multiple fallback strategies.
     
     Args:
         hf_token: Hugging Face API token
+        max_retries: Maximum retry attempts per method
         
     Returns:
         List of tool definitions
     """
-    client = HuggingFaceMCPClient(hf_token, timeout=90)  # Longer timeout for Spaces
-    return await client.get_all_tools()
+    last_error = None
+    
+    # Strategy 1: Try the robust client
+    for attempt in range(max_retries):
+        try:
+            logger.info(f"Trying robust client (attempt {attempt + 1})")
+            client = RobustHFMCPClient(hf_token, timeout=90)
+            tools = await client.get_all_tools()
+            logger.info(f"Robust client succeeded with {len(tools)} tools")
+            return tools
+        except Exception as e:
+            last_error = e
+            logger.warning(f"Robust client attempt {attempt + 1} failed: {e}")
+            if attempt < max_retries - 1:
+                await asyncio.sleep(2 ** attempt)  # Exponential backoff
+    
+    # Strategy 2: Try the simplified client
+    for attempt in range(max_retries):
+        try:
+            logger.info(f"Trying simplified client (attempt {attempt + 1})")
+            client = SimplifiedHFMCPClient(hf_token, timeout=120)
+            tools = await client.get_tools()
+            logger.info(f"Simplified client succeeded with {len(tools)} tools")
+            return tools
+        except Exception as e:
+            last_error = e
+            logger.warning(f"Simplified client attempt {attempt + 1} failed: {e}")
+            if attempt < max_retries - 1:
+                await asyncio.sleep(2 ** attempt)
+    
+    # If all strategies fail
+    raise Exception(f"All connection strategies failed. Last error: {last_error}")
 
 
-async def call_hf_tool(hf_token: str, tool_name: str, args: Dict[str, Any]) -> Any:
+async def call_hf_tool_robust(
+    hf_token: str, 
+    tool_name: str, 
+    args: Dict[str, Any],
+    max_retries: int = 3
+) -> Any:
     """
-    Call a specific Hugging Face MCP tool.
+    Call a specific Hugging Face MCP tool with multiple fallback strategies.
     
     Args:
         hf_token: Hugging Face API token
         tool_name: Name of the tool to call
         args: Arguments to pass to the tool
+        max_retries: Maximum retry attempts per method
         
     Returns:
         The tool's response
     """
-    client = HuggingFaceMCPClient(hf_token, timeout=90)
-    return await client.call_tool(tool_name, args)
+    last_error = None
+    
+    # Strategy 1: Try the robust client
+    for attempt in range(max_retries):
+        try:
+            logger.info(f"Trying robust client for tool call (attempt {attempt + 1})")
+            client = RobustHFMCPClient(hf_token, timeout=120)
+            result = await client.call_tool(tool_name, args)
+            logger.info(f"Robust client tool call succeeded")
+            return result
+        except Exception as e:
+            last_error = e
+            logger.warning(f"Robust client tool call attempt {attempt + 1} failed: {e}")
+            if attempt < max_retries - 1:
+                await asyncio.sleep(2 ** attempt)
+    
+    # Strategy 2: Try the simplified client
+    for attempt in range(max_retries):
+        try:
+            logger.info(f"Trying simplified client for tool call (attempt {attempt + 1})")
+            client = SimplifiedHFMCPClient(hf_token, timeout=150)
+            result = await client.call_tool(tool_name, args)
+            logger.info(f"Simplified client tool call succeeded")
+            return result
+        except Exception as e:
+            last_error = e
+            logger.warning(f"Simplified client tool call attempt {attempt + 1} failed: {e}")
+            if attempt < max_retries - 1:
+                await asyncio.sleep(2 ** attempt)
+    
+    # If all strategies fail
+    raise Exception(f"All tool call strategies failed. Last error: {last_error}")
 
 
-# Diagnostic function for Spaces environment
-async def diagnose_connection(hf_token: str) -> Dict[str, Any]:
+# Legacy compatibility functions
+async def get_hf_tools(hf_token: str) -> List[Dict[str, Any]]:
+    """Legacy function - now uses robust implementation."""
+    return await get_hf_tools_robust(hf_token)
+
+
+async def call_hf_tool(hf_token: str, tool_name: str, args: Dict[str, Any]) -> Any:
+    """Legacy function - now uses robust implementation."""
+    return await call_hf_tool_robust(hf_token, tool_name, args)
+
+
+# Enhanced diagnostics
+async def diagnose_connection_advanced(hf_token: str) -> Dict[str, Any]:
     """
-    Diagnose connection issues with detailed logging.
+    Advanced connection diagnostics with multiple test scenarios.
     
     Args:
         hf_token: Hugging Face API token
         
     Returns:
-        Diagnostic information
+        Comprehensive diagnostic information
     """
     diagnostics = {
         "environment": "huggingface_spaces" if os.getenv("SPACE_ID") else "local",
         "space_id": os.getenv("SPACE_ID"),
+        "python_version": os.sys.version,
         "token_length": len(hf_token) if hf_token else 0,
         "has_token": bool(hf_token),
-        "connection_test": False,
-        "initialization_test": False,
-        "tools_test": False,
-        "error": None
+        "tests": {
+            "basic_connection": False,
+            "robust_client": False,
+            "simplified_client": False,
+            "tools_fetch": False,
+            "tool_call_test": False
+        },
+        "errors": {},
+        "tool_count": 0,
+        "sample_tools": []
     }
     
+    # Test 1: Basic connection
     try:
-        # Test basic connection
-        from mcp.client.streamable_http import streamablehttp_client
-        
-        headers = {
-            "Authorization": f"Bearer {hf_token}",
-            "User-Agent": "hf-mcp-diagnostic/1.0.0"
-        }
-        
         async with streamablehttp_client(
             url="https://huggingface.co/mcp",
-            headers=headers,
-            timeout=timedelta(seconds=30),
-            terminate_on_close=True
+            headers={"Authorization": f"Bearer {hf_token}"},
+            timeout=timedelta(seconds=10),
+            terminate_on_close=False
         ) as (read_stream, write_stream, get_session_id):
-            diagnostics["connection_test"] = True
+            diagnostics["tests"]["basic_connection"] = True
             logger.info("Basic connection test passed")
-            
-            # Test initialization
-            client = HuggingFaceMCPClient(hf_token, timeout=60)
-            try:
-                tools = await client.get_all_tools()
-                diagnostics["initialization_test"] = True
-                diagnostics["tools_test"] = True
+    except Exception as e:
+        diagnostics["errors"]["basic_connection"] = str(e)
+        logger.error(f"Basic connection test failed: {e}")
+    
+    # Test 2: Robust client
+    if diagnostics["tests"]["basic_connection"]:
+        try:
+            client = RobustHFMCPClient(hf_token, timeout=60)
+            tools = await client.get_all_tools()
+            diagnostics["tests"]["robust_client"] = True
+            diagnostics["tests"]["tools_fetch"] = True
+            diagnostics["tool_count"] = len(tools)
+            diagnostics["sample_tools"] = [
+                {"name": tool.get("name"), "description": tool.get("description", "")[:100]}
+                for tool in tools[:3]
+            ]
+            logger.info(f"Robust client test passed - {len(tools)} tools")
+        except Exception as e:
+            diagnostics["errors"]["robust_client"] = str(e)
+            logger.error(f"Robust client test failed: {e}")
+    
+    # Test 3: Simplified client
+    if not diagnostics["tests"]["robust_client"]:
+        try:
+            client = SimplifiedHFMCPClient(hf_token, timeout=90)
+            tools = await client.get_tools()
+            diagnostics["tests"]["simplified_client"] = True
+            if not diagnostics["tests"]["tools_fetch"]:
+                diagnostics["tests"]["tools_fetch"] = True
                 diagnostics["tool_count"] = len(tools)
-                logger.info(f"Full diagnostic passed - found {len(tools)} tools")
-            except Exception as init_error:
-                diagnostics["error"] = str(init_error)
-                logger.error(f"Initialization failed: {init_error}")
+                diagnostics["sample_tools"] = [
+                    {"name": tool.get("name"), "description": tool.get("description", "")[:100]}
+                    for tool in tools[:3]
+                ]
+            logger.info(f"Simplified client test passed - {len(tools)} tools")
+        except Exception as e:
+            diagnostics["errors"]["simplified_client"] = str(e)
+            logger.error(f"Simplified client test failed: {e}")
+    
+    # Test 4: Tool call (if we have tools)
+    if diagnostics["tests"]["tools_fetch"] and diagnostics["sample_tools"]:
+        try:
+            # Try to call a simple tool if available
+            sample_tool_name = diagnostics["sample_tools"][0]["name"]
+            if sample_tool_name:
+                # Use the working client
+                if diagnostics["tests"]["robust_client"]:
+                    client = RobustHFMCPClient(hf_token, timeout=60)
+                else:
+                    client = SimplifiedHFMCPClient(hf_token, timeout=90)
                 
-    except Exception as conn_error:
-        diagnostics["error"] = str(conn_error)
-        logger.error(f"Connection test failed: {conn_error}")
+                # Try with empty args first (many tools accept this)
+                try:
+                    result = await client.call_tool(sample_tool_name, {})
+                    diagnostics["tests"]["tool_call_test"] = True
+                    logger.info(f"Tool call test passed with {sample_tool_name}")
+                except Exception as tool_error:
+                    # Tool call failed but that might be due to wrong args
+                    diagnostics["errors"]["tool_call_test"] = f"Tool call failed (might need args): {str(tool_error)}"
+                    logger.warning(f"Tool call test failed: {tool_error}")
+                    
+        except Exception as e:
+            diagnostics["errors"]["tool_call_test"] = str(e)
+            logger.error(f"Tool call test setup failed: {e}")
     
-    return diagnostics
-
+    return diagnostics