Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/googlecolab/colab-mcp/llms.txt

Use this file to discover all available pages before exploring further.

The session proxy layer is the glue between the local FastMCP server and the live Google Colab notebook in your browser. It is responsible for starting the WebSocket server, managing the FastMCP client connection to Colab, creating a proxy server that dynamically reflects Colab’s tools, and running the middleware chain that keeps the MCP client aware of connection-state changes. Together, these components ensure that the MCP client always has a working tool list — even before Colab has connected — and that it instantly discovers new notebook tools the moment a browser session comes online.

​
Components

ColabSessionProxy

The top-level orchestrator, instantiated once at startup. Its start_proxy_server() method:
  1. Starts a ColabWebSocketServer as an async context manager.
  2. Starts a ColabProxyClient backed by that WebSocket server.
  3. Creates a FastMCPProxy whose client_factory is proxy_client.client_factory — so it always routes tool calls to whichever client is currently live.
  4. Builds the ordered middleware list and returns it so __init__.py can register it on the top-level FastMCP instance.
ColabSessionProxy also holds the AsyncExitStack that owns the lifetimes of the WebSocket server and proxy client; calling cleanup() tears them all down in reverse order.

ColabProxyClient

Manages the fastmcp.Client connection to Colab through a ColabTransport. On startup it fires an asyncio.Task (_start_proxy_client) that blocks — via ColabTransport.connect_session — until the WebSocket server receives its first authenticated connection. Once the Client context is entered, proxy_mcp_client is set.is_connected() requires both conditions to be true: wss.connection_live.is_set() (the WebSocket is currently open) and proxy_mcp_client is not None (the MCP client session has fully initialized). Either condition alone is not sufficient.The key method is client_factory():
def client_factory(self):
    if self.is_connected():
        return self.proxy_mcp_client
    # return a client mapped to a stubbed mcp server if there is no session proxy
    return self.stubbed_mcp_client
When Colab is connected it returns proxy_mcp_client, which has the full set of Colab notebook tools. When Colab is disconnected it returns stubbed_mcp_client — a Client wrapping an empty FastMCP() — so the proxy server never crashes on an absent backend.

ColabTransport

A thin ClientTransport implementation that bridges ColabWebSocketServer’s anyio memory streams to the mcp library’s ClientSession. Its connect_session context manager wires wss.read_stream and wss.write_stream directly into a ClientSession, giving ColabProxyClient a standard MCP session interface over the local WebSocket.
@contextlib.asynccontextmanager
async def connect_session(self, **session_kwargs) -> AsyncIterator[ClientSession]:
    async with ClientSession(
        self.wss.read_stream, self.wss.write_stream, **session_kwargs
    ) as session:
        yield session

ColabProxyMiddleware

Intercepts every MCP message processed by the FastMCP server. On each message it:
  1. Reads the current connection state from ColabProxyClient.is_connected().
  2. Stores fe_connected, proxy_token, and proxy_port in the FastMCP context so downstream tools can read them.
  3. After the message is processed, checks whether the connection state has changed since the last message.
  4. If it has changed, calls send_tool_list_changed() to emit a notifications/tools/list_changed notification to the MCP client.
on_call_tool is also overridden to handle the open_colab_browser_connection tool call specifically — see the tool flow section below.

ToolInjectionMiddleware

A standard FastMCP ToolInjectionMiddleware configured with the open_colab_browser_connection tool object. It intercepts every tools/list response and appends the injected tool, ensuring the MCP client always sees it regardless of whether a Colab session is connected.
Middleware order matters. ColabProxyMiddleware must be registered before ToolInjectionMiddleware in the middleware list. FastMCP middleware runs in registration order, so ColabProxyMiddleware must execute first on every message — it is responsible for setting fe_connected, proxy_token, and proxy_port in the request context. If ToolInjectionMiddleware ran first, the injected tool’s handler would find those state values unset when it tries to read them.

​
The open_colab_browser_connection Tool Flow

The injected open_colab_browser_connection tool is the entry point for establishing a Colab session. Its behavior depends on whether a connection is already live.
1

Tool is always present

Because ToolInjectionMiddleware appends it to every tools/list response, the MCP client sees open_colab_browser_connection from the very first handshake — before any browser session has connected.
2

Called while already connected

If fe_connected is true in the current context, check_session_proxy_tool_fn returns True immediately. No browser window is opened.
3

Called while not connected — open the browser

If fe_connected is false, the function calls webbrowser.open_new(...) with the Colab scratch notebook URL, embedding the token and port from context:
webbrowser.open_new(
    f"{COLAB}{SCRATCH_PATH}#mcpProxyToken={token}&mcpProxyPort={port}"
)
This opens https://colab.research.google.com/notebooks/empty.ipynb#mcpProxyToken=<token>&mcpProxyPort=<port> in the system browser. The underlying tool function returns False at this point, but ColabProxyMiddleware.on_call_tool intercepts the result before it is returned to the client and replaces it with the final outcome after waiting for the connection (see next step).
4

ColabProxyMiddleware waits for connection

After call_next completes (the browser has been opened), on_call_tool in ColabProxyMiddleware checks whether the proxy client is now connected. If not yet connected, it reports progress to the MCP client and waits:
  • 1/3 — “The user is not connected to the Colab UI”
  • 2/3 — “Waiting for user to connect in Colab - will wait for 60s”
It then calls proxy_client.await_proxy_connection(), which uses asyncio.wait_for with a UI_CONNECTION_TIMEOUT of 60.0 seconds to await both wss.connection_live being set and the _start_proxy_client task completing.
5

Connection confirmed or timed out

After the wait resolves, middleware checks is_connected() once more:
  • Connected — reports progress 3/3 (“The Colab UI is successfully connected!”) and returns ToolResult with {"result": True}.
  • Timed out — reports progress 3/3 (“Timeout while waiting for the user to connect.”) and returns ToolResult with {"result": False}.

​
client_factory Fallback

FastMCPProxy calls client_factory() each time it needs to forward a tool call. When Colab is disconnected, ColabProxyClient.client_factory() returns a Client pointing at a stubbed empty FastMCP() server instead of raising an error or returning None. 
def client_factory(self):
    if self.is_connected():
        return self.proxy_mcp_client
    # return a client mapped to a stubbed mcp server if there is no session proxy
    return self.stubbed_mcp_client
This means the proxy server always has a valid backend to route to. Tool calls issued while Colab is disconnected will simply return empty results from the stub rather than crashing the MCP session. Once the user connects via open_colab_browser_connection, the factory begins returning proxy_mcp_client and all subsequent calls go to the live notebook.

Build docs developers (and LLMs) love

Get started for freeTalk to us