Spaces:
Running
Running
| """ | |
| StdIO connection management for MCP implementations. | |
| This module provides a connection manager for stdio-based MCP connections | |
| that ensures proper task isolation and resource cleanup. | |
| """ | |
| import asyncio | |
| import io | |
| import logging | |
| import sys | |
| from typing import Any, TextIO, Tuple | |
| from mcp import StdioServerParameters | |
| from mcp.client.stdio import stdio_client | |
| from openspace.utils.logging import Logger | |
| from openspace.grounding.core.transport.task_managers import ( | |
| AsyncContextConnectionManager, | |
| ) | |
| logger = Logger.get_logger(__name__) | |
| class FilteredStderrWrapper(io.TextIOBase): | |
| """Wrapper for stderr that filters out harmless MCP server shutdown messages. | |
| This wrapper suppresses error messages from MCP servers during shutdown | |
| that are harmless but create noise in the logs. | |
| """ | |
| def __init__(self, wrapped_stream: TextIO): | |
| """Initialize the wrapper. | |
| Args: | |
| wrapped_stream: The underlying stderr stream | |
| """ | |
| self._stream = wrapped_stream | |
| self._buffer = "" | |
| self._in_traceback = False | |
| self._traceback_lines = [] | |
| self._in_rich_traceback = False # Track rich-formatted tracebacks | |
| self._rich_traceback_needs_error_line = False # After ╰, need one more line | |
| def write(self, s: str) -> int: | |
| """Write to stderr, filtering out harmless error messages. | |
| Args: | |
| s: The string to write | |
| Returns: | |
| Number of characters written | |
| """ | |
| # Buffer the input for line-by-line processing | |
| self._buffer += s | |
| # Process complete lines | |
| while '\n' in self._buffer: | |
| line, self._buffer = self._buffer.split('\n', 1) | |
| self._process_line(line + '\n') | |
| return len(s) | |
| def _process_line(self, line: str): | |
| """Process a single line and decide whether to output it.""" | |
| # Detect start of traceback or exception group | |
| if line.lstrip().startswith(("╭", "┏")): | |
| self._in_traceback = True | |
| self._in_rich_traceback = True | |
| self._rich_traceback_needs_error_line = False | |
| self._traceback_lines = [line] | |
| return | |
| if (line.strip().startswith('Traceback (most recent call last)') or | |
| line.strip().startswith('Exception Group Traceback (most recent call last)') or | |
| line.strip().startswith('BaseExceptionGroup:') or | |
| line.strip().startswith('ExceptionGroup:')): | |
| self._in_traceback = True | |
| self._traceback_lines = [line] | |
| self._in_rich_traceback = False | |
| self._rich_traceback_needs_error_line = False | |
| return | |
| # Collect traceback lines | |
| if self._in_traceback: | |
| self._traceback_lines.append(line) | |
| # If not in rich traceback mode, but current line contains rich border characters, switch to rich mode | |
| if not self._in_rich_traceback and any(ch in line for ch in ("╭", "┏")): | |
| self._in_rich_traceback = True | |
| # Check for end of rich-formatted traceback (line with ╰) | |
| if self._in_rich_traceback and '╰' in line: | |
| # Rich traceback box ended, but we need to collect the error line that follows | |
| self._rich_traceback_needs_error_line = True | |
| return | |
| # If we just ended a rich traceback, this should be the error line | |
| if self._rich_traceback_needs_error_line: | |
| # Now we have the complete rich traceback including the error line | |
| if self._is_harmless_error(): | |
| logger.debug(f"Suppressed harmless rich-formatted MCP server error") | |
| else: | |
| # Output the full traceback | |
| for tb_line in self._traceback_lines: | |
| self._stream.write(tb_line) | |
| self._stream.flush() | |
| # Reset traceback collection | |
| self._in_traceback = False | |
| self._in_rich_traceback = False | |
| self._rich_traceback_needs_error_line = False | |
| self._traceback_lines = [] | |
| return | |
| # For exception groups, we need to collect more lines | |
| # Check if we've collected enough to determine if it's harmless | |
| if len(self._traceback_lines) > 5 and not self._in_rich_traceback: | |
| # Check periodically if this is a harmless error | |
| if self._is_harmless_error(): | |
| # Suppress this traceback | |
| logger.debug(f"Suppressed harmless MCP server shutdown error") | |
| self._in_traceback = False | |
| self._in_rich_traceback = False | |
| self._rich_traceback_needs_error_line = False | |
| self._traceback_lines = [] | |
| return | |
| # Check if this is the error line (last line of regular traceback) | |
| # But not for rich tracebacks which use box characters | |
| # A final traceback line is typically unindented and contains "ErrorType: message" | |
| if not self._in_rich_traceback and line and not line[0].isspace() and ':' in line: | |
| # Check if this is a harmless cleanup error | |
| if self._is_harmless_error(): | |
| # Suppress this traceback | |
| logger.debug(f"Suppressed harmless MCP server shutdown error") | |
| else: | |
| # Output the full traceback | |
| for tb_line in self._traceback_lines: | |
| self._stream.write(tb_line) | |
| self._stream.flush() | |
| # Reset traceback collection | |
| self._in_traceback = False | |
| self._in_rich_traceback = False | |
| self._rich_traceback_needs_error_line = False | |
| self._traceback_lines = [] | |
| return | |
| # If we've collected too many lines without finding the end, output and reset | |
| if len(self._traceback_lines) > 100: | |
| # Output what we have | |
| for tb_line in self._traceback_lines: | |
| self._stream.write(tb_line) | |
| self._stream.flush() | |
| self._in_traceback = False | |
| self._in_rich_traceback = False | |
| self._rich_traceback_needs_error_line = False | |
| self._traceback_lines = [] | |
| return | |
| else: | |
| # Normal line - check if it's a harmless error log | |
| line_lower = line.lower() | |
| harmless_log_patterns = [ | |
| 'an error occurred during closing of asynchronous generator', | |
| 'asyncgen:', | |
| 'service stopped.', | |
| ] | |
| # Check if this is a harmless log line | |
| is_harmless_log = any(pattern in line_lower for pattern in harmless_log_patterns) | |
| if not is_harmless_log: | |
| # Output normal lines | |
| self._stream.write(line) | |
| self._stream.flush() | |
| else: | |
| # Suppress harmless log messages | |
| logger.debug(f"Suppressed harmless log line: {line.strip()}") | |
| def _is_harmless_error(self) -> bool: | |
| """Check if the collected traceback is a harmless error.""" | |
| traceback_text = ''.join(self._traceback_lines).lower() | |
| # List of harmless error patterns (case-insensitive) | |
| harmless_patterns = [ | |
| 'valueerror: i/o operation on closed file', | |
| 'oserror: [errno 9] bad file descriptor', | |
| 'brokenpipeerror', | |
| 'runtimeerror: attempted to exit cancel scope in a different task', | |
| 'baseexceptiongroup: unhandled errors in a taskgroup', | |
| 'generatorexit', | |
| 'an error occurred during closing of asynchronous generator', | |
| ] | |
| # Check if any pattern matches and it's related to shutdown | |
| for pattern in harmless_patterns: | |
| if pattern in traceback_text: | |
| # Also check if it's related to shutdown/cleanup | |
| shutdown_keywords = ['finally:', 'stopped', 'cleanup', '__exit__', '__aexit__', 'stdio_client', 'service stopped'] | |
| if any(keyword in traceback_text for keyword in shutdown_keywords): | |
| return True | |
| return False | |
| def flush(self): | |
| """Flush any remaining buffered content and the underlying stream.""" | |
| if self._buffer: | |
| self._process_line(self._buffer) | |
| self._buffer = "" | |
| if self._traceback_lines: | |
| # Flush incomplete traceback | |
| for line in self._traceback_lines: | |
| self._stream.write(line) | |
| self._traceback_lines = [] | |
| self._stream.flush() | |
| def fileno(self) -> int: | |
| """Return the file descriptor of the underlying stream.""" | |
| if hasattr(self._stream, 'fileno'): | |
| return self._stream.fileno() | |
| return -1 | |
| def closed(self) -> bool: | |
| """Check if the stream is closed.""" | |
| return self._stream.closed | |
| class StdioConnectionManager(AsyncContextConnectionManager[Tuple[Any, Any], ...]): | |
| """Connection manager for stdio-based MCP connections. | |
| This class handles the proper task isolation for stdio_client context managers | |
| to prevent the "cancel scope in different task" error. It runs the stdio_client | |
| in a dedicated task and manages its lifecycle. | |
| Note: Error handling during cleanup (e.g., I/O operations on closed files) is | |
| handled by the parent AsyncContextConnectionManager class in _close_connection(). | |
| """ | |
| def __init__( | |
| self, | |
| server_params: StdioServerParameters, | |
| errlog: TextIO | None = None, | |
| ): | |
| """Initialize a new stdio connection manager. | |
| Args: | |
| server_params: The parameters for the stdio server | |
| errlog: The error log stream (defaults to filtered sys.stderr) | |
| """ | |
| # Wrap stderr to filter out harmless shutdown errors | |
| if errlog is None: | |
| errlog = FilteredStderrWrapper(sys.stderr) | |
| elif not isinstance(errlog, FilteredStderrWrapper): | |
| errlog = FilteredStderrWrapper(errlog) | |
| super().__init__(stdio_client, server_params, errlog) | |
| self.server_params = server_params | |
| self.errlog = errlog | |
| self._mcp_logger_filter = None | |
| self._stop_event: asyncio.Event | None = None # Signal for background task | |
| self._runner_task: asyncio.Task | None = None # Background runner task | |
| self._conn_future: asyncio.Future | None = None # Future for the established connection | |
| logger.debug("StdioConnectionManager init with params=%s", server_params) | |
| async def _establish_connection(self) -> Tuple[Any, Any]: | |
| """Establish connection in a dedicated task to avoid cancel-scope issues.""" | |
| # Suppress MCP SDK's noisy JSON parse errors **before** starting the runner | |
| self._suppress_mcp_json_errors() | |
| # Lazily create primitives the first time we connect | |
| if self._stop_event is None: | |
| self._stop_event = asyncio.Event() | |
| if self._conn_future is None or self._conn_future.done(): | |
| self._conn_future = asyncio.get_event_loop().create_future() | |
| async def _runner(): # Runs in its *own* task (same task for enter/exit) | |
| try: | |
| async with stdio_client(self.server_params, self.errlog) as conn: | |
| # Pass connection back to the caller | |
| if not self._conn_future.done(): | |
| self._conn_future.set_result(conn) | |
| # Wait until close is requested | |
| await self._stop_event.wait() | |
| finally: | |
| # Make sure the future is set even on error so awaiters don’t hang | |
| if not self._conn_future.done(): | |
| self._conn_future.set_exception(RuntimeError("Connection failed")) | |
| # Start background runner if not already active | |
| if self._runner_task is None or self._runner_task.done(): | |
| self._runner_task = asyncio.create_task(_runner(), name="stdio_client_runner") | |
| # Wait for the connection tuple from the future | |
| conn: Tuple[Any, Any] = await self._conn_future # type: ignore | |
| return conn | |
| async def _close_connection(self) -> None: | |
| """Request the background task to exit its context and wait for it.""" | |
| try: | |
| # Restore original logging configuration *before* shutdown | |
| self._restore_mcp_logging() | |
| # Signal the runner to exit its context manager | |
| if self._stop_event and not self._stop_event.is_set(): | |
| self._stop_event.set() | |
| # Await the runner task so that __aexit__ executes in *its* task | |
| if self._runner_task: | |
| try: | |
| await asyncio.wait_for(self._runner_task, timeout=2.0) | |
| except asyncio.TimeoutError: | |
| logger.warning("Timeout while waiting for stdio_client to shut down") | |
| finally: | |
| # Clean up helpers so next connect() creates new ones | |
| self._runner_task = None | |
| self._stop_event = None | |
| self._conn_future = None | |
| def _suppress_mcp_json_errors(self): | |
| """Suppress MCP SDK's JSON parsing error logs. | |
| The MCP SDK logs errors when it receives non-JSON messages from servers. | |
| These are harmless (the SDK continues working), so we filter them out. | |
| """ | |
| mcp_logger = logging.getLogger("mcp.client.stdio") | |
| class JSONErrorFilter(logging.Filter): | |
| """Filter out JSON parsing errors from MCP SDK.""" | |
| def filter(self, record): | |
| # Suppress "Failed to parse JSONRPC message" errors | |
| if "Failed to parse JSONRPC message" in str(record.msg): | |
| return False | |
| return True | |
| self._mcp_logger_filter = JSONErrorFilter() | |
| mcp_logger.addFilter(self._mcp_logger_filter) | |
| def _restore_mcp_logging(self): | |
| """Restore MCP SDK logging to normal.""" | |
| if self._mcp_logger_filter: | |
| mcp_logger = logging.getLogger("mcp.client.stdio") | |
| mcp_logger.removeFilter(self._mcp_logger_filter) | |
| self._mcp_logger_filter = None | |
| if not isinstance(sys.stderr, FilteredStderrWrapper): | |
| sys.stderr = FilteredStderrWrapper(sys.stderr) | |
| logger.debug("Applied global FilteredStderrWrapper to sys.stderr") |