agent.py

import asyncio
import os
import sys
import threading
import time
import re
import atexit
from contextlib import asynccontextmanager
from typing import Any, Optional, List, Dict, Tuple, Callable  
from smolagents import CodeAgent, MCPClient
from smolagents.models import Model
from inference import initialize, generate_content
from workflow_vizualizer import track_workflow_step, track_communication, complete_workflow_step

# Global session management
_session_initialized = False
_session_lock = threading.Lock()
_session_start_time = None

# Enhanced global caching for Phase 2 optimizations with async support
_global_tools_cache = {}
_global_tools_timestamp = None
_global_model_instance = None
_global_model_lock = threading.Lock()
_global_connection_pool = {}
_global_connection_lock = threading.Lock()

# Managed event loop system
_managed_event_loop = None
_event_loop_lock = threading.Lock()
_event_loop_manager = None  # Global event loop manager instance


@asynccontextmanager
async def managed_event_loop():
    """Proper async context manager for event loop lifecycle."""
    global _managed_event_loop
    
    try:
        # Create new event loop if needed
        if _managed_event_loop is None or _managed_event_loop.is_closed():
            _managed_event_loop = asyncio.new_event_loop()
            asyncio.set_event_loop(_managed_event_loop)
        
        print("✅ Event loop initialized and set as current")
        yield _managed_event_loop
        
    except Exception as e:
        print(f"❌ Event loop error: {e}")
        raise
    finally:
        # Don't close the loop here - let it be managed at a higher level
        pass


async def safe_async_call(coroutine, timeout=30):
    """Safely execute async calls with proper error handling."""
    try:
        return await asyncio.wait_for(coroutine, timeout=timeout)
    except asyncio.TimeoutError:
        print(f"⏱️ Async call timed out after {timeout}s")
        raise
    except RuntimeError as e:
        if "Event loop is closed" in str(e):
            print("🔄 Event loop closed - attempting to create new one")
            # Create new event loop and retry
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            return await asyncio.wait_for(coroutine, timeout=timeout)
        raise


class AsyncEventLoopManager:
    def __init__(self):
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._thread: Optional[threading.Thread] = None
        self._loop = asyncio.new_event_loop()
        self._thread = threading.Thread(target=self._run_loop, daemon=True)
        self._thread.start()
        print("AsyncEventLoopManager: Initialized and thread started.")
    
    def _run_loop(self):
        if self._loop is None:
            print("AsyncEventLoopManager: _run_loop called but loop is None.")
            return
        asyncio.set_event_loop(self._loop)
        try:
            print("AsyncEventLoopManager: Event loop running.")
            self._loop.run_forever()
        except Exception as e:
            print(f"AsyncEventLoopManager: Exception in event loop: {e}")
        finally:
            # Ensure the loop is stopped if it was running.
            # The actual closing is handled by the shutdown() method.
            if self._loop and self._loop.is_running():
                self._loop.stop()
            print("AsyncEventLoopManager: Event loop stopped in _run_loop finally.")

    def run_async(self, coro):
        """Run a coroutine in the event loop from another thread."""
        coro_name = getattr(coro, '__name__', str(coro))
        if self._loop is None:
            print(f"AsyncEventLoopManager: Loop object is None. Cannot run coroutine {coro_name}.")
            raise RuntimeError("Event loop manager is not properly initialized (loop missing).")

        if self._loop.is_closed():
            print(f"AsyncEventLoopManager: Loop is CLOSED. Cannot schedule coroutine {coro_name}.")
            raise RuntimeError(f"Event loop is closed. Cannot run {coro_name}.")

        if self._thread is None or not self._thread.is_alive():
            print(f"AsyncEventLoopManager: Event loop thread is not alive or None. Cannot run coroutine {coro_name}.")
            raise RuntimeError("Event loop thread is not alive or None.")

        try:
            future = asyncio.run_coroutine_threadsafe(coro, self._loop)
            return future.result(timeout=30)  # Assuming a 30s timeout
        except RuntimeError as e:
            print(f"AsyncEventLoopManager: RuntimeError during run_coroutine_threadsafe for {coro_name}: {e}")
            raise
        except asyncio.TimeoutError:
            print(f"AsyncEventLoopManager: Timeout waiting for coroutine {coro_name} result.")
            raise
        except Exception as e:
            print(f"AsyncEventLoopManager: Error submitting coroutine {coro_name}: {e}")
            raise

    def shutdown(self):
        """Stop and close the event loop."""
        print("AsyncEventLoopManager: Shutdown initiated.")
        if self._loop and not self._loop.is_closed():
            if self._loop.is_running():
                self._loop.call_soon_threadsafe(self._loop.stop)
                print("AsyncEventLoopManager: Stop signal sent to running event loop.")
            else:
                print("AsyncEventLoopManager: Event loop was not running, but attempting to stop.")
                # If not running, stop might not be necessary or might error, 
                # but call_soon_threadsafe should be safe.
                try:
                    self._loop.call_soon_threadsafe(self._loop.stop)
                except RuntimeError as e:
                    print(f"AsyncEventLoopManager: Info - could not send stop to non-running loop: {e}")

        if self._thread and self._thread.is_alive():
            self._thread.join(timeout=10)
            if self._thread.is_alive():
                print("AsyncEventLoopManager: Thread did not join in time during shutdown.")
            else:
                print("AsyncEventLoopManager: Thread joined.")
        else:
            print("AsyncEventLoopManager: Thread already stopped, not initialized, or None at shutdown.")

        # Explicitly close the loop here after the thread has finished.
        if self._loop and not self._loop.is_closed():
            try:
                # Ensure all tasks are cancelled before closing
                # Gather all tasks:
                if sys.version_info >= (3, 7): # gather works on all tasks in 3.7+
                    tasks = asyncio.all_tasks(self._loop)
                    for task in tasks:
                        task.cancel()
                    # Wait for tasks to cancel - this should be done within the loop's thread ideally
                    # but since we are shutting down from outside, this is a best effort.
                    # self._loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
                self._loop.close()
                print("AsyncEventLoopManager: Event loop closed in shutdown.")
            except Exception as e:
                print(f"AsyncEventLoopManager: Exception while closing loop: {e}")
        elif self._loop and self._loop.is_closed():
            print("AsyncEventLoopManager: Event loop was already closed.")
        else:
            print("AsyncEventLoopManager: No loop to close or loop was None.")
    
        self._loop = None
        self._thread = None
        print("AsyncEventLoopManager: Shutdown process complete.")

def get_event_loop_manager():
    """Get or create the global event loop manager."""
    global _event_loop_manager

    with _event_loop_lock:
        # Check if manager exists and its loop and thread are valid
        manager_valid = False
        if _event_loop_manager is not None:
            # More robust check: loop exists, is not closed, thread exists and is alive
            if _event_loop_manager._loop is not None and \
               not _event_loop_manager._loop.is_closed() and \
               _event_loop_manager._thread is not None and \
               _event_loop_manager._thread.is_alive():
                manager_valid = True
            else:
                print("get_event_loop_manager: Existing manager found but its loop or thread is invalid. Recreating.")
                try:
                    _event_loop_manager.shutdown() # Attempt to clean up the old one
                except Exception as e:
                    print(f"get_event_loop_manager: Error shutting down invalid manager: {e}")
                _event_loop_manager = None # Ensure it's None so a new one is created

        if _event_loop_manager is None: # Covers both initial creation and recreation
            print("get_event_loop_manager: Creating new AsyncEventLoopManager instance.")
            _event_loop_manager = AsyncEventLoopManager()
        else:
            print("get_event_loop_manager: Reusing existing valid AsyncEventLoopManager instance.")
    return _event_loop_manager

def shutdown_event_loop_manager():
    """Shutdown the global event loop manager."""
    global _event_loop_manager
    with _event_loop_lock:
        if _event_loop_manager is not None:
            print("shutdown_event_loop_manager: Shutting down global event loop manager.")
            try:
                _event_loop_manager.shutdown()
            except Exception as e:
                print(f"shutdown_event_loop_manager: Error during shutdown: {e}")
            finally:
                _event_loop_manager = None
        else:
            print("shutdown_event_loop_manager: No active event loop manager to shut down.")

class AsyncMCPClientWrapper:
    """Wrapper for async MCP client operations."""

    def __init__(self, url: str):
        self.url = url
        self._mcp_client = None
        self._tools = None
        self._tools_cache_time = None
        self._cache_ttl = 300  # 5 minutes cache
        self._connected = False

    async def ensure_connected(self):
        """Ensure async connection is established."""
        if not self._connected or self._mcp_client is None:
            try:
                # Create MCP client with SSE transport for Gradio
                self._mcp_client = MCPClient({"url": self.url, "transport": "sse"})
                # Attempt a lightweight operation to confirm connectivity, e.g., get_tools or a custom ping
                # For now, we assume MCPClient constructor success implies basic connectivity.
                # If get_tools is lightweight enough, it can be called here.
                # await self._mcp_client.get_tools() # Example, if get_tools is async and suitable
                self._connected = True
                print(f"✅ Connected to MCP server: {self.url}")
            except Exception as e:
                self._connected = False
                print(f"❌ Failed to connect to {self.url}: {e}")
                raise

    async def get_tools(self):
        """Get tools asynchronously."""
        current_time = time.time()

        # Check instance cache
        if (self._tools is not None and
            self._tools_cache_time is not None and
            current_time - self._tools_cache_time < self._cache_ttl):
            return self._tools

        # Fetch fresh tools
        await self.ensure_connected() # Ensures client is connected

        if self._mcp_client is None: # Should be caught by ensure_connected, but as a safeguard
            raise RuntimeError("MCP client not connected")

        try:
            # Assuming MCPClient.get_tools() is a synchronous method based on original structure
            # If it were async, it would be `await self._mcp_client.get_tools()`
            self._tools = self._mcp_client.get_tools()
            self._tools_cache_time = current_time
            tool_names = [tool.name for tool in self._tools] if self._tools else []
            print(f"🔧 Fetched {len(tool_names)} tools from {self.url}: {tool_names}")

            return self._tools
        except Exception as e:
            print(f"❌ Error fetching tools from {self.url}: {e}")
            # Potentially reset connection status if error indicates a connection problem
            # self._connected = False 
            raise

    async def disconnect(self):
        """Gracefully disconnect."""
        if self._mcp_client and self._connected:
            try:
                # Assuming MCPClient.disconnect() is synchronous
                # If it were async, it would be `await self._mcp_client.disconnect()`
                self._mcp_client.disconnect()
            except Exception as e:
                print(f"Error during MCPClient disconnect for {self.url}: {e}")
                # Log error but continue to mark as disconnected
                pass # Fall through to set _connected = False
        self._connected = False
        self._mcp_client = None
        print(f"🔌 Disconnected from MCP server: {self.url}")

class AsyncPersistentMCPClient:
    """Async-aware persistent MCP client that survives multiple requests."""

    def __init__(self, url: str):
        self.url = url
        self._wrapper = AsyncMCPClientWrapper(url)
        self._loop_manager = None

    def ensure_connected(self):
        """Sync wrapper for async connection."""
        if self._loop_manager is None:
            self._loop_manager = get_event_loop_manager()

        conn_step = track_communication("agent", "mcp_client", "connection_ensure", f"Ensuring connection to {self.url}")
        try:
            # Ensure we have a valid loop manager
            if self._loop_manager is None:
                self._loop_manager = get_event_loop_manager()

            # Additional safety check
            if self._loop_manager is None:
                raise RuntimeError("Failed to create event loop manager")

              # Pass the coroutine object itself, not its result
            self._loop_manager.run_async(self._wrapper.ensure_connected())
            complete_workflow_step(conn_step, "completed", details={"url": self.url})
        except Exception as e:
            complete_workflow_step(conn_step, "error", details={"error": str(e)})
            raise

    def get_client(self):
        """Get the underlying MCP client."""
        self.ensure_connected()
        return self._wrapper._mcp_client

    def get_tools(self):
        """Get tools with enhanced caching and async support."""
        global _global_tools_cache, _global_tools_timestamp
        current_time = time.time()

        if self._loop_manager is None:
            self._loop_manager = get_event_loop_manager()

        # Phase 2 Optimization: Check server-specific global cache first
        with _global_connection_lock:
            server_cache_key = self.url
            server_cache = _global_tools_cache.get(server_cache_key, {})

            if (server_cache and _global_tools_timestamp and
                current_time - _global_tools_timestamp < 300):
                # Track global cache hit
                cache_step = track_communication("mcp_client", "mcp_server", "cache_hit_global", f"Using global cached tools for {self.url}")
                complete_workflow_step(cache_step, "completed", details={
                    "tools": list(server_cache.keys()),
                    "cache_type": "global_server_specific",
                    "server_url": self.url,
                    "cache_age": current_time - _global_tools_timestamp
                })
                return list(server_cache.values())

        # Fetch fresh tools using async
        tools_step = track_communication("mcp_client", "mcp_server", "get_tools", f"Fetching tools from {self.url} (cache refresh)")
        try:
            # Ensure we have a valid loop manager
            if self._loop_manager is None:
                self._loop_manager = get_event_loop_manager()

            # Additional safety check
            if self._loop_manager is None:
                raise RuntimeError("Failed to create event loop manager")

            # Pass the coroutine object itself
            tools = self._loop_manager.run_async(self._wrapper.get_tools())

            # Update global cache
            with _global_connection_lock:
                if tools:
                    if server_cache_key not in _global_tools_cache:
                        _global_tools_cache[server_cache_key] = {}

                    _global_tools_cache[server_cache_key] = {tool.name: tool for tool in tools}
                    _global_tools_timestamp = current_time

                    total_tools = sum(len(server_tools) for server_tools in _global_tools_cache.values())
                    print(f"🔄 Global tools cache updated for {self.url}: {len(tools)} tools")
                    print(f"   Total cached tools across all servers: {total_tools}")

            tool_names = [tool.name for tool in tools] if tools else []
            complete_workflow_step(tools_step, "completed", details={
                "tools": tool_names,
                "count": len(tool_names),
                "server_url": self.url,
                "cache_status": "refreshed_server_specific",
                "global_cache_servers": len(_global_tools_cache)
            })
            return tools

        except Exception as e:
            complete_workflow_step(tools_step, "error", details={"error": str(e), "server_url": self.url})
            raise

    def disconnect(self):
        """Gracefully disconnect."""
        if self._loop_manager and self._wrapper:
            try:
                # Ensure we have a valid loop manager
                if self._loop_manager is None:
                    self._loop_manager = get_event_loop_manager()

                # Additional safety check
                if self._loop_manager is None:
                    raise RuntimeError("Failed to create event loop manager")

                # Pass the coroutine object itself
                self._loop_manager.run_async(self._wrapper.disconnect())
            except RuntimeError as e:
                # Handle cases where the loop might already be closed or unable to run tasks
                print(f"AsyncPersistentMCPClient: Error running disconnect for {self.url} in async loop: {e}")
            except Exception as e:
                print(f"AsyncPersistentMCPClient: General error during disconnect for {self.url}: {e}")

def get_mcp_client(url: str = "https://NLarchive-Agent-client-multi-mcp-SKT.hf.space/gradio_api/mcp/sse") -> AsyncPersistentMCPClient:
    """Get or create an MCP client with enhanced global connection pooling."""
    # Phase 2 Optimization: Use global connection pool
    with _global_connection_lock:
        if url not in _global_connection_pool:
            conn_step = track_communication("agent", "mcp_client", "connection_create", f"Creating new global connection to {url}")
            _global_connection_pool[url] = AsyncPersistentMCPClient(url)
            complete_workflow_step(conn_step, "completed", details={"url": url, "pool_size": len(_global_connection_pool)})
        else:
            # Track connection reuse
            reuse_step = track_communication("agent", "mcp_client", "connection_reuse", f"Reusing global connection to {url}")
            complete_workflow_step(reuse_step, "completed", details={"url": url, "pool_size": len(_global_connection_pool)})

        return _global_connection_pool[url]

def get_global_model() -> 'CachedLocalInferenceModel':
    """Get or create global model instance for Phase 2 optimization."""
    global _global_model_instance

    with _global_model_lock:
        if _global_model_instance is None:
            model_step = track_workflow_step("model_init_global", "Initializing global model instance")

            # CRITICAL FIX: Create and assign BEFORE initialization
            _global_model_instance = CachedLocalInferenceModel()

            # Now initialize the model
            try:
                _global_model_instance.ensure_initialized()
                complete_workflow_step(model_step, "completed", details={"model_type": "global_cached"})
                print(f"🤖 Global model instance created and initialized")
            except Exception as e:
                # If initialization fails, reset global instance
                _global_model_instance = None
                complete_workflow_step(model_step, "error", details={"error": str(e)})
                raise
        else:
            # Track model reuse
            reuse_step = track_workflow_step("model_reuse", "Reusing global model instance")
            complete_workflow_step(reuse_step, "completed", details={"model_type": "global_cached"})

    return _global_model_instance

def reset_global_state():
    """Reset global state for testing purposes with server-specific cache awareness."""
    global _global_tools_cache, _global_tools_timestamp, _global_model_instance, _global_connection_pool, _event_loop_manager

    with _global_connection_lock:
        # FIXED: Clear server-specific cache structure (don't rebind!)
        _global_tools_cache.clear()  # Now clears {url: {tool_name: tool}} structure
        _global_tools_timestamp = None

        # Disconnect all connections but keep pool structure
        for client in _global_connection_pool.values():
            try:
                client.disconnect()
            except:
                pass

    with _global_model_lock:
        # Don't reset model instance - it should persist
        pass

    print("🔄 Global state reset for testing (server-specific cache cleared)")

# Enhanced LocalInferenceModel with workflow tracking
class CachedLocalInferenceModel(Model):
    """Model with enhanced caching and session persistence."""

    def __init__(self):
        super().__init__()
        self._response_cache = {}
        self._cache_hits = 0
        self._cache_misses = 0
        self._model_ready = False

    def ensure_initialized(self):
        """Lazy initialization of the model."""
        if not self._model_ready:
            init_step = track_workflow_step("model_init", "Initializing inference model (lazy)")
            try:
                initialize()
                self._model_ready = True
                complete_workflow_step(init_step, "completed")
            except Exception as e:
                complete_workflow_step(init_step, "error", details={"error": str(e)})
                raise

    def generate(self, messages: Any, **kwargs: Any) -> Any:
        self.ensure_initialized()

        prompt = self._format_messages(messages)

        # Enhanced cache with hash-based lookup
        cache_key = hash(prompt)
        if cache_key in self._response_cache:
            self._cache_hits += 1
            cached_response = self._response_cache[cache_key]

            # Track cache hit
            cache_step = track_communication("agent", "llm_service", "cache_hit", "Using cached response")
            complete_workflow_step(cache_step, "completed", details={
                "cache_hits": self._cache_hits,
                "cache_misses": self._cache_misses,
                "cache_ratio": self._cache_hits / (self._cache_hits + self._cache_misses)
            })

            return ModelResponse(cached_response.content, prompt)

        self._cache_misses += 1

        # Track LLM call
        llm_step = track_communication("agent", "llm_service", "generate_request", "Generating new response")

        try:
            enhanced_prompt = self._enhance_prompt_for_tools(prompt)

            response_text = generate_content(
                prompt=enhanced_prompt,
                model_name=kwargs.get('model_name'),
                allow_fallbacks=True,
                generation_config={
                    'temperature': kwargs.get('temperature', 0.3),
                    'max_output_tokens': kwargs.get('max_tokens', 512)
                }
            )

            # Validate and fix response format
            if not self._is_valid_code_response(response_text):
                response_text = self._fix_response_format(response_text, prompt)

            response = ModelResponse(str(response_text), prompt)

            # Smart cache management (keep most recent 10 responses)
            if len(self._response_cache) >= 10:
                # Remove oldest entry (simple FIFO)
                oldest_key = next(iter(self._response_cache))
                del self._response_cache[oldest_key]

            self._response_cache[cache_key] = response

            complete_workflow_step(llm_step, "completed", details={
                "cache_status": "new",
                "input_tokens": response.token_usage.input_tokens,
                "output_tokens": response.token_usage.output_tokens,
                "model": response.model
            })

            return response

        except Exception as e:
            fallback_response = self._create_fallback_response(prompt, str(e))
            complete_workflow_step(llm_step, "error", details={"error": str(e)})
            return ModelResponse(fallback_response, prompt)

    def _enhance_prompt_for_tools(self, prompt: str) -> str:
        """Enhance the prompt with better tool usage examples."""
        if "sentiment" in prompt.lower():
            tool_example = """
IMPORTANT: When calling sentiment_analysis, use keyword arguments only:
Correct: sentiment_analysis(text="your text here")  
Wrong: sentiment_analysis("your text here")

Example:
```py
text = "this is horrible"
result = sentiment_analysis(text=text)
final_answer(result)
```"""
            return prompt + "\n" + tool_example
        return prompt

    def _format_messages(self, messages: Any) -> str:
        """Convert messages to a single prompt string."""
        if isinstance(messages, str):
            return messages
        elif isinstance(messages, list):
            prompt_parts = []
            for msg in messages:
                if isinstance(msg, dict):
                    if 'content' in msg:
                        content = msg['content']
                        role = msg.get('role', 'user')
                        if isinstance(content, list):
                            text_parts = [part.get('text', '') for part in content if part.get('type') == 'text']
                            content = ' '.join(text_parts)
                        prompt_parts.append(f"{role}: {content}")
                    elif 'text' in msg:
                        prompt_parts.append(msg['text'])
                elif hasattr(msg, 'content'):
                    prompt_parts.append(str(msg.content))
                else:
                    prompt_parts.append(str(msg))
            return '\n'.join(prompt_parts)
        else:
            return str(messages)

    def _is_valid_code_response(self, response: str) -> bool:
        """Check if response contains valid code block format."""
        code_pattern = r'```(?:py|python)?\s*\n(.*?)\n```'
        return bool(re.search(code_pattern, response, re.DOTALL))

    def _fix_response_format(self, response: str, original_prompt: str) -> str:
        """Try to fix response format to match expected pattern."""
        # Attempt to remove or comment out "Thoughts:" if not in a code block already
        # This is a common source of SyntaxError if the LLM includes it directly
        if "Thoughts:" in response and not "```" in response.split("Thoughts:")[0]:
            # If "Thoughts:" appears before any code block, comment it out
            response = response.replace("Thoughts:", "# Thoughts:", 1)

        if "sentiment" in original_prompt.lower():
            text_to_analyze = "neutral text"
            if "this is horrible" in original_prompt:
                text_to_analyze = "this is horrible"
            elif "awful" in original_prompt:
                text_to_analyze = "awful"

            return f"""Thoughts: I need to analyze the sentiment of the given text using the sentiment_analysis tool.
Code:
```py
text = "{text_to_analyze}"
result = sentiment_analysis(text=text)
final_answer(result)
```<end_code>"""

        if "```" in response and ("Thoughts:" in response or "Code:" in response):
            return response
        
        clean_response = response.replace('"', '\\"').replace('\n', '\\n')
        return f"""Thoughts: Processing the user's request.
Code:
```py
result = "{clean_response}"
final_answer(result)
```<end_code>"""

    def _create_fallback_response(self, prompt: str, error_msg: str) -> str:
        """Create a valid fallback response when the model fails."""
        return f"""Thoughts: The AI service is experiencing issues, providing a fallback response.
Code:
```py
error_message = "I apologize, but the AI service is temporarily experiencing high load. Please try again in a moment."
final_answer(error_message)
```<end_code>"""

class TokenUsage:
    def __init__(self, input_tokens: int = 0, output_tokens: int = 0):
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens
        self.total_tokens = input_tokens + output_tokens
        self.prompt_tokens = input_tokens
        self.completion_tokens = output_tokens

class ModelResponse:
    def __init__(self, content: str, prompt: str = ""):
        self.content = content
        self.text = content
        estimated_input_tokens = len(prompt.split()) if prompt else 0
        estimated_output_tokens = len(content.split()) if content else 0
        self.token_usage = TokenUsage(estimated_input_tokens, estimated_output_tokens)
        self.finish_reason = 'stop'
        self.model = 'local-inference'
    
    def __str__(self):
        return self.content

# Global variables
_mcp_client = None
_tools = None
_model = None
_agent = None
_initialized = False
_initialization_lock = threading.Lock()

def initialize_agent():
    """Initialize the agent components with Hugging Face Spaces MCP servers."""
    global _mcp_client, _tools, _model, _agent, _initialized
    
    with _initialization_lock:
        if _initialized:
            skip_step = track_workflow_step("agent_init_skip", "Agent already initialized - using cached instance")
            complete_workflow_step(skip_step, "completed", details={"optimization": "session_persistence"})
            return
        
        try:
            print("Initializing MCP agent...")
            
            agent_init_step = track_workflow_step("agent_init", "Initializing MCP agent components")
            
            # Get clients for Hugging Face Spaces servers
            all_tools = []
            tool_names = set()
            
            # Semantic Search & Keywords server
            try:
                semantic_client = get_mcp_client("https://nlarchive-mcp-semantic-keywords.hf.space/gradio_api/mcp/sse")
                semantic_tools = semantic_client.get_tools()
                for tool in semantic_tools:
                    if tool.name not in tool_names:
                        all_tools.append(tool)
                        tool_names.add(tool.name)
                print(f"Connected to semantic server: {len(semantic_tools)} tools - {[t.name for t in semantic_tools]}")
            except Exception as e:
                print(f"WARNING: Semantic server unavailable: {e}")
            
            # Token Counter server
            try:
                token_client = get_mcp_client("https://nlarchive-mcp-gr-token-counter.hf.space/gradio_api/mcp/sse")
                token_tools = token_client.get_tools()
                for tool in token_tools:
                    if tool.name not in tool_names:
                        all_tools.append(tool)
                        tool_names.add(tool.name)
                print(f"Connected to token counter server: {len(token_tools)} tools - {[t.name for t in token_tools]}")
            except Exception as e:
                print(f"WARNING: Token counter server unavailable: {e}")
            
            # Sentiment Analysis server
            try:
                sentiment_client = get_mcp_client("https://nlarchive-mcp-sentiment.hf.space/gradio_api/mcp/sse")
                sentiment_tools = sentiment_client.get_tools()
                for tool in sentiment_tools:
                    if tool.name not in tool_names:
                        all_tools.append(tool)
                        tool_names.add(tool.name)
                print(f"Connected to sentiment analysis server: {len(sentiment_tools)} tools - {[t.name for t in sentiment_tools]}")
            except Exception as e:
                print(f"WARNING: Sentiment analysis server unavailable: {e}")
                        
            _tools = all_tools
            _model = get_global_model()
            
            # Create agent with unique tools only
            _agent = CodeAgent(tools=_tools, model=_model)
            
            complete_workflow_step(agent_init_step, "completed", details={
                "tools_count": len(_tools),
                "unique_tool_names": list(tool_names),
                "servers_connected": 3
            })
            
            _initialized = True
            print(f"Agent initialized with {len(_tools)} unique tools: {list(tool_names)}")
            
        except Exception as e:
            print(f"Agent initialization failed: {e}")
            _model = get_global_model()
            _agent = CodeAgent(tools=[], model=_model)
            _initialized = True
            print("Agent initialized in fallback mode")

def is_agent_initialized() -> bool:
    """Check if the agent is initialized."""
    return _initialized

def run_agent(message: str) -> str:
    """Send message through the agent with comprehensive tracking."""
    if not _initialized:
        initialize_agent()
    if _agent is None:
        raise RuntimeError("Agent not properly initialized")
    
    # Track agent processing
    process_step = track_workflow_step("agent_process", f"Processing: {message}")
    
    try:
        # Enhanced tool tracking
        tool_step: Optional[str] = None
        detected_tools = []
        
        # Detect potential tool usage
        if any(keyword in message.lower() for keyword in ['sentiment', 'analyze', 'feeling']):
            detected_tools.append('sentiment_analysis')
        if any(keyword in message.lower() for keyword in ['token', 'count']):
            detected_tools.extend(['count_tokens_openai_gpt4', 'count_tokens_bert_family'])
        if any(keyword in message.lower() for keyword in ['semantic', 'similar', 'keyword']):
            detected_tools.extend(['semantic_similarity', 'extract_semantic_keywords'])
        
        if detected_tools:
            tool_step = track_communication("agent", "mcp_server", "tool_call", 
                                          f"Executing tools {detected_tools} for: {message[:50]}...")
        
        result = _agent.run(message)
        
        # Complete tool step if it was tracked
        if tool_step is not None:
            complete_workflow_step(tool_step, "completed", details={
                "result": str(result)[:100],
                "detected_tools": detected_tools
            })
        
        complete_workflow_step(process_step, "completed", details={
            "result_length": len(str(result)),
            "detected_tools": detected_tools
        })
        
        return str(result)
        
    except Exception as e:
        error_msg = str(e)
        print(f"Agent execution error: {error_msg}")
        
        complete_workflow_step(process_step, "error", details={"error": error_msg})
        
        # Enhanced error responses
        if "503" in error_msg or "overloaded" in error_msg.lower():
            return "I apologize, but the AI service is currently experiencing high demand. Please try again in a few moments."
        elif "rate limit" in error_msg.lower():
            return "The service is currently rate-limited. Please wait a moment before trying again."
        elif "event loop" in error_msg.lower():
            return "There was an async processing issue. The system is recovering. Please try again."
        else:
            return "I encountered an error while processing your request. Please try rephrasing your question or try again later."

def disconnect():
    """Cleanly disconnect connections with global pool management."""
    global _mcp_client, _initialized
    disconnect_step = track_workflow_step("agent_disconnect", "Disconnecting MCP client")
    
    try:
        # Phase 2 Optimization: Preserve global connections for reuse
        with _global_connection_lock:
            preserved_connections = 0
            for url, client in _global_connection_pool.items():
                try:
                    # Keep connections alive but mark as idle
                    if hasattr(client, '_last_used'):
                        client._last_used = time.time()
                    preserved_connections += 1
                except:
                    pass
        
        complete_workflow_step(disconnect_step, "completed", details={
            "preserved_connections": preserved_connections,
            "optimization": "connection_persistence"
        })
    except Exception as e:
        complete_workflow_step(disconnect_step, "error", details={"error": str(e)})
    finally:
        # Don't reset global state - preserve for next session
        _initialized = False

def initialize_session():
    """Initialize the persistent session - alias for initialize_agent."""
    initialize_agent()

def is_session_initialized() -> bool:
    """Check if the persistent session is initialized - alias for is_agent_initialized."""
    return is_agent_initialized()

# Make sure these are exported for imports
__all__ = [
    'run_agent', 'initialize_agent', 'is_agent_initialized', 'disconnect',
    'initialize_session', 'is_session_initialized',
    'get_mcp_client', 'get_global_model', 'reset_global_state',
    '_global_tools_cache', '_global_connection_pool', '_global_model_instance',
    '_global_connection_lock', '_global_model_lock'
]

# Register cleanup function
def cleanup_global_resources():
    """Cleanup function for graceful shutdown."""
    global _global_connection_pool, _event_loop_manager, _global_connection_lock, _event_loop_lock
    
    print("Cleaning up global resources...")
    
    with _global_connection_lock:
        for client in _global_connection_pool.values():
            try:
                client.disconnect()
            except:
                pass
        _global_connection_pool.clear()
    
    # Shutdown event loop manager
    with _event_loop_lock:
        if _event_loop_manager:
            try:
                _event_loop_manager.shutdown()
            except:
                pass
            _event_loop_manager = None

# Register cleanup on exit
atexit.register(cleanup_global_resources)