massgen 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

massgen/backend/gemini_mcp_manager.py

@@ -1,545 +0,0 @@
-# -*- coding: utf-8 -*-
-"""
-Lightweight MCP manager for Gemini that wraps backend_utils with Gemini-specific streaming status, session conversion, and planning-mode checks.
-"""
-
-import asyncio
-from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple
-
-from ..logger_config import log_backend_activity, logger
-from ..mcp_tools import (
-    MCPCircuitBreakerManager,
-    MCPClient,
-    MCPConfigurationError,
-    MCPConfigValidator,
-    MCPConnectionError,
-    MCPError,
-    MCPErrorHandler,
-    MCPExecutionManager,
-    MCPResourceManager,
-    MCPRetryHandler,
-    MCPServerError,
-    MCPSetupManager,
-    MCPTimeoutError,
-    MCPValidationError,
-)
-from ..mcp_tools.hooks import convert_sessions_to_permission_sessions
-from .base import StreamChunk
-
-
-class GeminiMCPManager:
-    def __init__(self, backend_instance) -> None:
-        self.backend = backend_instance
-        # References to backend state/config
-        self.mcp_servers = getattr(self.backend, "mcp_servers", [])
-        self.allowed_tools = getattr(self.backend, "allowed_tools", None)
-        self.exclude_tools = getattr(self.backend, "exclude_tools", None)
-        self.agent_id = getattr(self.backend, "agent_id", None)
-
-        # Counters (will be updated on backend instance)
-        self._mcp_tool_calls_count = getattr(self.backend, "_mcp_tool_calls_count", 0)
-        self._mcp_tool_failures = getattr(self.backend, "_mcp_tool_failures", 0)
-        self._mcp_tool_successes = getattr(self.backend, "_mcp_tool_successes", 0)
-
-        # MCP client and init state mirror backend
-        self._mcp_client: Optional[MCPClient] = getattr(self.backend, "_mcp_client", None)
-        self._mcp_initialized: bool = getattr(self.backend, "_mcp_initialized", False)
-
-        # Circuit breaker and filesystem manager references
-        self._mcp_tools_circuit_breaker = getattr(self.backend, "_mcp_tools_circuit_breaker", None)
-        self.filesystem_manager = getattr(self.backend, "filesystem_manager", None)
-
-    async def setup_mcp_with_status_stream(self, agent_id: Optional[str] = None) -> AsyncGenerator[StreamChunk, None]:
-        """Initialize MCP client with status streaming."""
-        if not self.mcp_servers or self._mcp_initialized:
-            if False:
-                yield  # make this an async generator
-            return
-
-        status_queue: asyncio.Queue[StreamChunk] = asyncio.Queue()
-        _agent_id = agent_id or self.agent_id
-
-        async def status_callback(status: str, details: Dict[str, Any]) -> None:
-            """Callback to queue status updates as StreamChunks."""
-            chunk = StreamChunk(
-                type="mcp_status",
-                status=status,
-                content=details.get("message", ""),
-                source="mcp_tools",
-            )
-            await status_queue.put(chunk)
-
-        setup_task = asyncio.create_task(self.setup_mcp_internal(_agent_id, status_callback))
-
-        while not setup_task.done():
-            try:
-                chunk = await asyncio.wait_for(status_queue.get(), timeout=0.1)
-                yield chunk
-            except asyncio.TimeoutError:
-                continue
-
-        try:
-            await setup_task
-        except Exception as e:
-            yield StreamChunk(
-                type="mcp_status",
-                status="error",
-                content=f"MCP setup failed: {e}",
-                source="mcp_tools",
-            )
-
-    async def setup_mcp_tools(self, agent_id: Optional[str] = None) -> None:
-        """Initialize MCP client (sessions only) - wrapper that consumes status stream."""
-        async for _ in self.setup_mcp_with_status_stream(agent_id):
-            pass
-
-    async def setup_mcp_internal(
-        self,
-        agent_id: Optional[str] = None,
-        status_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None,
-    ) -> None:
-        """Internal MCP setup logic leveraging backend_utils managers."""
-        if not self.mcp_servers or self._mcp_initialized:
-            return
-
-        _agent_id = agent_id or self.agent_id
-
-        try:
-            # Validate MCP configuration
-            backend_config = {
-                "mcp_servers": self.mcp_servers,
-                "allowed_tools": self.allowed_tools,
-                "exclude_tools": self.exclude_tools,
-            }
-
-            try:
-                validated_config = MCPConfigValidator.validate_backend_mcp_config(backend_config)
-                self.mcp_servers = validated_config.get("mcp_servers", self.mcp_servers)
-                if status_callback:
-                    await status_callback("info", {"message": f"MCP configuration validated: {len(self.mcp_servers)} servers"})
-                log_backend_activity("gemini", "MCP configuration validated", {"server_count": len(self.mcp_servers)}, agent_id=_agent_id)
-            except MCPConfigurationError as e:
-                if status_callback:
-                    await status_callback("error", {"message": f"Invalid MCP configuration: {e}"})
-                raise RuntimeError(f"Invalid MCP configuration: {e}") from e
-            except MCPValidationError as e:
-                if status_callback:
-                    await status_callback("error", {"message": f"MCP validation error: {e}"})
-                raise RuntimeError(f"MCP validation error: {e}") from e
-            except Exception as e:
-                # Validation unavailable or unexpected error; continue with normalization path
-                log_backend_activity("gemini", "MCP validation unavailable or error", {"error": str(e)}, agent_id=_agent_id)
-
-            # Normalize servers
-            normalized_servers = MCPSetupManager.normalize_mcp_servers(self.mcp_servers, backend_name="gemini", agent_id=_agent_id)
-            if status_callback:
-                await status_callback("info", {"message": f"Setting up MCP sessions for {len(normalized_servers)} servers"})
-
-            # Apply circuit breaker filtering
-            if self._mcp_tools_circuit_breaker:
-                filtered_servers = MCPCircuitBreakerManager.apply_circuit_breaker_filtering(
-                    normalized_servers,
-                    self._mcp_tools_circuit_breaker,
-                    backend_name="gemini",
-                    agent_id=_agent_id,
-                )
-            else:
-                filtered_servers = normalized_servers
-
-            if not filtered_servers:
-                log_backend_activity("gemini", "All MCP servers blocked by circuit breaker", {}, agent_id=_agent_id)
-                if status_callback:
-                    await status_callback("warning", {"message": "All MCP servers blocked by circuit breaker"})
-                return
-
-            # Extract tool filtering parameters
-            allowed_tools = backend_config.get("allowed_tools")
-            exclude_tools = backend_config.get("exclude_tools")
-
-            # Setup MCP client via resource manager (handles retries and filtering)
-            client = await MCPResourceManager.setup_mcp_client(
-                servers=filtered_servers,
-                allowed_tools=allowed_tools,
-                exclude_tools=exclude_tools,
-                circuit_breaker=self._mcp_tools_circuit_breaker,
-                timeout_seconds=30,
-                backend_name="gemini",
-                agent_id=_agent_id,
-            )
-
-            if not client:
-                # Treat as connection failure
-                self._mcp_client = None
-                self.backend._mcp_client = None
-                if status_callback:
-                    await status_callback("error", {"message": "MCP connection failed: no servers connected"})
-                log_backend_activity("gemini", "MCP connection failed: no servers connected", {}, agent_id=_agent_id)
-                return
-
-            # Assign on success
-            self._mcp_client = client
-            self.backend._mcp_client = client
-            self._mcp_initialized = True
-            self.backend._mcp_initialized = True
-
-            log_backend_activity("gemini", "MCP sessions initialized successfully", {}, agent_id=_agent_id)
-            if status_callback:
-                # Attempt to list connected servers
-                try:
-                    names = client.get_server_names()
-                except Exception:
-                    names = []
-                await status_callback(
-                    "success",
-                    {"message": f"MCP sessions initialized successfully with {len(names)} servers"},
-                )
-
-        except Exception as e:
-            # Enhanced error mapping using backend_utils
-            log_type, user_message, _ = MCPErrorHandler.get_error_details(e)
-            log_backend_activity("gemini", f"MCP {log_type} during setup", {"error": str(e)}, agent_id=_agent_id)
-            self._mcp_client = None
-            self.backend._mcp_client = None
-            self._mcp_initialized = False
-            self.backend._mcp_initialized = False
-            if status_callback:
-                await status_callback("error", {"message": f"MCP session setup failed: {e}"})
-
-    async def handle_mcp_retry_error(self, error: Exception, retry_count: int, max_retries: int) -> Tuple[bool, AsyncGenerator[StreamChunk, None]]:
-        """Delegate retry error handling to backend_utils with StreamChunk emission."""
-        return await MCPRetryHandler.handle_retry_error(
-            error=error,
-            retry_count=retry_count,
-            max_retries=max_retries,
-            stream_chunk_class=StreamChunk,
-            backend_name="gemini",
-            agent_id=self.agent_id,
-        )
-
-    async def handle_mcp_error_and_fallback(self, error: Exception) -> AsyncGenerator[StreamChunk, None]:
-        """Handle MCP errors with specific messaging and fallback behavior."""
-        # increment backend failure counter
-        try:
-            self.backend._mcp_tool_failures += 1
-        except Exception:
-            pass
-
-        async for chunk in MCPRetryHandler.handle_error_and_fallback(
-            error=error,
-            tool_call_count=getattr(self.backend, "_mcp_tool_calls_count", 0),
-            stream_chunk_class=StreamChunk,
-            backend_name="gemini",
-            agent_id=self.agent_id,
-        ):
-            yield chunk
-
-    async def execute_mcp_function_with_retry(self, function_name: str, args: Dict[str, Any], functions: Dict, agent_id: Optional[str] = None) -> Any:
-        """Execute MCP function with retry and circuit breaker recording."""
-        _agent_id = agent_id or self.agent_id
-
-        async def stats_callback(action: str) -> int:
-            if action == "increment_calls":
-                self.backend._mcp_tool_calls_count += 1
-                return self.backend._mcp_tool_calls_count
-            elif action == "increment_failures":
-                self.backend._mcp_tool_failures += 1
-                return self.backend._mcp_tool_failures
-            return 0
-
-        async def circuit_breaker_callback(event: str, error_msg: str) -> None:
-            try:
-                if event == "failure":
-                    # Record failure for all configured servers
-                    servers = MCPSetupManager.normalize_mcp_servers(self.mcp_servers, backend_name="gemini", agent_id=_agent_id)
-                    if self._mcp_tools_circuit_breaker:
-                        await MCPCircuitBreakerManager.record_event(
-                            servers,
-                            self._mcp_tools_circuit_breaker,
-                            "failure",
-                            error_message=error_msg,
-                            backend_name="gemini",
-                            agent_id=_agent_id,
-                        )
-                else:
-                    # Record success only for connected servers
-                    connected_names: List[str] = []
-                    try:
-                        if self._mcp_client:
-                            connected_names = self._mcp_client.get_server_names()
-                    except Exception:
-                        connected_names = []
-
-                    if connected_names and self._mcp_tools_circuit_breaker:
-                        servers_to_record = [{"name": name} for name in connected_names]
-                        await MCPCircuitBreakerManager.record_event(
-                            servers_to_record,
-                            self._mcp_tools_circuit_breaker,
-                            "success",
-                            backend_name="gemini",
-                            agent_id=_agent_id,
-                        )
-            except Exception:
-                # Never fail the call due to circuit breaker recording issues
-                pass
-
-        return await MCPExecutionManager.execute_function_with_retry(
-            function_name=function_name,
-            args=args,
-            functions=functions,
-            max_retries=3,
-            stats_callback=stats_callback,
-            circuit_breaker_callback=circuit_breaker_callback,
-            logger_instance=logger,
-        )
-
-    async def setup_mcp_sessions_with_retry(self, agent_id: Optional[str] = None, max_retries: int = 5) -> Tuple[bool, AsyncGenerator[StreamChunk, None]]:
-        """Attempt to setup sessions with retries; returns (connected, status_chunks)."""
-        _agent_id = agent_id or self.agent_id
-        chunks: List[StreamChunk] = []
-        success: bool = False
-
-        for retry_count in range(1, max_retries + 1):
-            # Emit retry status
-            chunks.append(
-                StreamChunk(
-                    type="mcp_status",
-                    status="mcp_retry",
-                    content=f"Retrying MCP connection (attempt {retry_count}/{max_retries})",
-                    source="mcp_tools",
-                ),
-            )
-
-            try:
-                # Apply circuit breaker filtering before attempts
-                normalized = MCPSetupManager.normalize_mcp_servers(self.mcp_servers, backend_name="gemini", agent_id=_agent_id)
-                if self._mcp_tools_circuit_breaker:
-                    filtered = MCPCircuitBreakerManager.apply_circuit_breaker_filtering(
-                        normalized,
-                        self._mcp_tools_circuit_breaker,
-                        backend_name="gemini",
-                        agent_id=_agent_id,
-                    )
-                else:
-                    filtered = normalized
-
-                if not filtered:
-                    chunks.append(
-                        StreamChunk(
-                            type="mcp_status",
-                            status="mcp_blocked",
-                            content="All MCP servers blocked by circuit breaker",
-                            source="mcp_tools",
-                        ),
-                    )
-                    success = False
-                    break
-
-                # Validate config for allowed/excluded tools
-                try:
-                    MCPConfigValidator.validate_backend_mcp_config({"mcp_servers": filtered})
-                except Exception:
-                    pass
-
-                allowed_items = self.allowed_tools
-                exclude_items = self.exclude_tools
-                allowed_count = len(allowed_items) if isinstance(allowed_items, (list, tuple, set)) else (1 if allowed_items else 0)
-                exclude_count = len(exclude_items) if isinstance(exclude_items, (list, tuple, set)) else (1 if exclude_items else 0)
-                logger.debug(
-                    "[GeminiMCPManager] Retry filter settings | allowed=%d excluded=%d",
-                    allowed_count,
-                    exclude_count,
-                )
-
-                client = await MCPResourceManager.setup_mcp_client(
-                    servers=filtered,
-                    allowed_tools=self.allowed_tools,
-                    exclude_tools=self.exclude_tools,
-                    circuit_breaker=self._mcp_tools_circuit_breaker,
-                    timeout_seconds=30,
-                    backend_name="gemini",
-                    agent_id=_agent_id,
-                )
-
-                if client:
-                    # Assign to backend/manager and mark success
-                    self._mcp_client = client
-                    self.backend._mcp_client = client
-                    self._mcp_initialized = True
-                    self.backend._mcp_initialized = True
-
-                    # Record success event
-                    if self._mcp_tools_circuit_breaker:
-                        await MCPCircuitBreakerManager.record_event(
-                            filtered,
-                            self._mcp_tools_circuit_breaker,
-                            "success",
-                            backend_name="gemini",
-                            agent_id=_agent_id,
-                        )
-
-                    chunks.append(
-                        StreamChunk(
-                            type="mcp_status",
-                            status="mcp_connected",
-                            content=f"MCP connection successful on attempt {retry_count}",
-                            source="mcp_tools",
-                        ),
-                    )
-                    success = True
-                    break
-
-                # Client not connected; handle as transient error with retry logic
-                should_continue, error_chunks = await self.handle_mcp_retry_error(RuntimeError("No servers connected"), retry_count, max_retries)
-                # Drain any error chunks into list
-                async for ch in error_chunks:
-                    chunks.append(ch)
-                if not should_continue:
-                    success = False
-                    break
-
-            except (MCPConnectionError, MCPTimeoutError, MCPServerError, MCPError, Exception) as e:
-                should_continue, error_chunks = await self.handle_mcp_retry_error(e, retry_count, max_retries)
-                async for ch in error_chunks:
-                    chunks.append(ch)
-                if not should_continue:
-                    success = False
-                    break
-
-            # Progressive backoff between retries
-            await asyncio.sleep(0.5 * retry_count)
-
-        async def _generator():
-            for ch in chunks:
-                yield ch
-
-        return success, _generator()
-
-    def get_active_mcp_sessions(self, convert_to_permission_sessions: bool = True) -> List[Any]:
-        """Return active MCP ClientSession objects, optionally wrapped with permission sessions."""
-        sessions: List[Any] = []
-        try:
-            if self._mcp_client:
-                sessions = self._mcp_client.get_active_sessions()
-        except Exception:
-            sessions = []
-
-        if convert_to_permission_sessions and sessions and self.filesystem_manager:
-            try:
-                return convert_sessions_to_permission_sessions(sessions, self.filesystem_manager.path_permission_manager)
-            except Exception as e:
-                logger.error(f"[GeminiMCPManager] Failed to convert sessions to permission sessions: {e}")
-                return sessions
-        return sessions
-
-    def should_block_mcp_tools_in_planning_mode(self, is_planning_mode: bool, available_tools: List[str]) -> bool:
-        """Return True to block MCP tools if planning mode is enabled; logs details."""
-        if is_planning_mode:
-            log_backend_activity(
-                "gemini",
-                "MCP tools blocked in planning mode",
-                {"blocked_tools": len(available_tools or []), "tools_preview": (available_tools or [])[:5]},
-                agent_id=self.agent_id,
-            )
-            return True
-        return False
-
-    async def cleanup_genai_resources(self, stream, client) -> None:
-        """Cleanup google-genai resources to avoid unclosed aiohttp sessions."""
-        # Close stream
-        try:
-            if stream is not None:
-                close_fn = getattr(stream, "aclose", None) or getattr(stream, "close", None)
-                if close_fn is not None:
-                    maybe = close_fn()
-                    if hasattr(maybe, "__await__"):
-                        await maybe
-        except Exception as e:
-            log_backend_activity(
-                "gemini",
-                "Stream cleanup failed",
-                {"error": str(e)},
-                agent_id=self.agent_id,
-            )
-        # Close internal aiohttp session held by google-genai BaseApiClient
-        try:
-            if client is not None:
-                base_client = getattr(client, "_api_client", None)
-                if base_client is not None:
-                    session = getattr(base_client, "_aiohttp_session", None)
-                    if session is not None and hasattr(session, "close"):
-                        if not session.closed:
-                            await session.close()
-                            log_backend_activity(
-                                "gemini",
-                                "Closed google-genai aiohttp session",
-                                {},
-                                agent_id=self.agent_id,
-                            )
-                        base_client._aiohttp_session = None
-                        # Yield control to allow connector cleanup
-                        await asyncio.sleep(0)
-        except Exception as e:
-            log_backend_activity(
-                "gemini",
-                "Failed to close google-genai aiohttp session",
-                {"error": str(e)},
-                agent_id=self.agent_id,
-            )
-        # Close internal async transport if exposed
-        try:
-            if client is not None and hasattr(client, "aio") and client.aio is not None:
-                aio_obj = client.aio
-                for method_name in ("close", "stop"):
-                    method = getattr(aio_obj, method_name, None)
-                    if method:
-                        maybe = method()
-                        if hasattr(maybe, "__await__"):
-                            await maybe
-                        break
-        except Exception as e:
-            log_backend_activity(
-                "gemini",
-                "Client AIO cleanup failed",
-                {"error": str(e)},
-                agent_id=self.agent_id,
-            )
-
-        # Close client
-        try:
-            if client is not None:
-                for method_name in ("aclose", "close"):
-                    method = getattr(client, method_name, None)
-                    if method:
-                        maybe = method()
-                        if hasattr(maybe, "__await__"):
-                            await maybe
-                        break
-        except Exception as e:
-            log_backend_activity(
-                "gemini",
-                "Client cleanup failed",
-                {"error": str(e)},
-                agent_id=self.agent_id,
-            )
-
-    def is_mcp_connected(self) -> bool:
-        try:
-            return bool(self._mcp_initialized and self._mcp_client and self._mcp_client.is_connected())
-        except Exception:
-            return False
-
-    def get_mcp_server_names(self) -> List[str]:
-        try:
-            if self._mcp_client:
-                return self._mcp_client.get_server_names()
-        except Exception:
-            pass
-        return []
-
-    def get_mcp_tools(self) -> Dict[str, Any]:
-        try:
-            if self._mcp_client and hasattr(self._mcp_client, "tools"):
-                return self._mcp_client.tools
-        except Exception:
-            pass
-        return {}