ColabSessionProxy and ColabProxyClient Reference

The session proxy layer is responsible for wiring together the local WebSocket server, the FastMCP proxy, and the middleware chain so that an AI agent’s tool calls are transparently forwarded to a live Google Colab browser session. Three classes share this responsibility: ColabSessionProxy acts as the top-level orchestrator, ColabProxyClient manages the lifecycle of the FastMCP Client that connects to Colab, and ColabTransport provides the ClientTransport implementation that routes MCP messages over ColabWebSocketServer’s memory streams.

ColabSessionProxy

Module: colab_mcp.session ColabSessionProxy is the entry point for enabling the proxy subsystem. It creates and wires together a ColabWebSocketServer, a ColabProxyClient, a FastMCPProxy, and the two middleware layers. All resources are tracked through an AsyncExitStack so that a single cleanup() call tears everything down cleanly.

Startup Pattern

The following pattern, taken directly from __init__.py, shows how ColabSessionProxy is used to extend the main FastMCP server:

session_mcp = ColabSessionProxy()
await session_mcp.start_proxy_server()
mcp.mount(session_mcp.proxy_server)
for middleware in session_mcp.middleware:
    mcp.add_middleware(middleware)

Methods

`start_proxy_server()`

async def start_proxy_server() -> None

Starts all proxy subsystem components in order:

Enters ColabWebSocketServer as an async context manager and assigns it to self.wss.
Enters ColabProxyClient(self.wss) as an async context manager (which kicks off the background task that waits for the WebSocket connection).
Constructs a FastMCPProxy with client_factory=proxy_client.client_factory and instructions="Connects to a user's Google Colab session in a browser and allows for interactions with their Google Colab notebook", and assigns it to self.proxy_server.
Appends ColabProxyMiddleware(proxy_client) then ToolInjectionMiddleware(tools=[check_session_proxy_tool]) to self.middleware.

After this method returns, self.proxy_server and self.middleware are fully initialized and ready to be mounted on the main MCP server.

`cleanup()`

async def cleanup() -> None

Closes all resources managed by the internal AsyncExitStack. This shuts down the ColabWebSocketServer (closing the WebSocket listener and both memory streams) and the ColabProxyClient (cancelling any pending background tasks and closing the proxy Client).

Attributes

proxy_server

FastMCPProxy | None

The FastMCPProxy instance to mount on the main MCP server with mcp.mount(). None before start_proxy_server() is called. The proxy uses proxy_client.client_factory to obtain either the live Colab client or a stubbed empty client on each request.

middleware

list[Middleware]

Ordered list of middleware to register with mcp.add_middleware(). After start_proxy_server() this list contains exactly two entries:

ColabProxyMiddleware — sets connection state on every message and handles the open_colab_browser_connection wait logic.
ToolInjectionMiddleware — injects the open_colab_browser_connection tool into the tool list.

wss

ColabWebSocketServer | None

The ColabWebSocketServer instance. None before start_proxy_server() is called.

Middleware order in session_mcp.middleware matters. ColabProxyMiddleware must come before ToolInjectionMiddleware because it sets the fe_connected, proxy_token, and proxy_port context state keys on every incoming message. ToolInjectionMiddleware (and the check_session_proxy_tool_fn it calls) reads those keys. If the order were reversed, the state would not be populated when the injected tool function executes.

ColabProxyClient

Module: colab_mcp.session ColabProxyClient manages the FastMCP Client that communicates with the Colab browser session via ColabTransport. It uses an async background task to establish the Client connection (which blocks until the WebSocket is connected), so the rest of the application can continue initializing without blocking.

Usage

ColabProxyClient is used as an async context manager:

async with ColabProxyClient(wss) as proxy_client:
    # background task is running; will connect when Colab tab connects
    await proxy_client.await_proxy_connection()
    if proxy_client.is_connected():
        client = proxy_client.client_factory()

Constructor

ColabProxyClient(wss: ColabWebSocketServer)

wss

ColabWebSocketServer

required

The ColabWebSocketServer whose connection_live event and memory streams will be used to detect and serve the active connection.

Methods

`is_connected()`

def is_connected() -> bool

Returns True only when both conditions are met:

wss.connection_live is set (the WebSocket client is actively connected), and
proxy_mcp_client is not None (the Client context has been fully entered).

Both conditions must be true simultaneously because connection_live is set as soon as the WebSocket handshake completes, but proxy_mcp_client is only assigned once the Client async context manager’s __aenter__ finishes.

`client_factory()`

def client_factory() -> Client

Returns the appropriate FastMCP Client for the current state:

If is_connected() is True: returns proxy_mcp_client — the live client connected to the Colab session.
Otherwise: returns stubbed_mcp_client — a Client connected to a local empty FastMCP() instance that returns empty responses for all requests.

This factory is passed directly to FastMCPProxy, which calls it on every proxied request.

`await_proxy_connection()`

async def await_proxy_connection() -> None

Waits up to UI_CONNECTION_TIMEOUT (60.0 seconds) for both wss.connection_live to be set and the background _start_proxy_client task to complete. Uses asyncio.wait_for with asyncio.gather internally; a TimeoutError is suppressed so this method always returns without raising.

ColabTransport

Module: colab_mcp.session ColabTransport implements the ClientTransport interface expected by FastMCP’s Client, routing all MCP session I/O through ColabWebSocketServer’s in-process memory streams rather than a real network connection.

Constructor

ColabTransport(wss: ColabWebSocketServer)

wss

ColabWebSocketServer

required

The ColabWebSocketServer instance whose read_stream and write_stream will be used as the transport’s I/O channels.

Methods

`connect_session(**session_kwargs)`

@contextlib.asynccontextmanager
async def connect_session(**session_kwargs) -> AsyncIterator[ClientSession]

Async context manager that yields a ClientSession backed by wss.read_stream and wss.write_stream. Any keyword arguments are forwarded directly to the ClientSession constructor. This is the method called by the FastMCP Client internals when it needs to establish a session.

transport = ColabTransport(wss)
async with transport.connect_session(foo="bar") as session:
    # session is an mcp.client.session.ClientSession
    ...

API Reference

ColabSessionProxy and ColabProxyClient Reference