diff --git a/marimo/_server/ai/mcp/transport.py b/marimo/_server/ai/mcp/transport.py index 724247ffd5d..a27a638ba7f 100644 --- a/marimo/_server/ai/mcp/transport.py +++ b/marimo/_server/ai/mcp/transport.py @@ -106,9 +106,12 @@ class MCPTransportRegistry: """Registry for MCP transport connectors.""" def __init__(self) -> None: + # Moved local names outside of dict literal to avoid repeated attribute lookups + stdio = MCPTransportType.STDIO + http = MCPTransportType.STREAMABLE_HTTP self._connectors: dict[MCPTransportType, MCPTransportConnector] = { - MCPTransportType.STDIO: StdioTransportConnector(), - MCPTransportType.STREAMABLE_HTTP: StreamableHTTPTransportConnector(), + stdio: StdioTransportConnector(), + http: StreamableHTTPTransportConnector(), } def get_connector( @@ -125,6 +128,7 @@ def get_connector( Raises: ValueError: If transport type is not supported """ - if transport_type not in self._connectors: + try: + return self._connectors[transport_type] + except KeyError: raise ValueError(f"Unsupported transport type: {transport_type}") - return self._connectors[transport_type]