Implementation:Microsoft Autogen Studio MCP WSBridge

Metadata
Sources	python/packages/autogen-studio/autogenstudio/mcp/wsbridge.py
Domains	MCP, WebSocket, Bridge, Real_Time_Communication
Last Updated	2026-02-11 17:00 GMT

Overview

Description

The Studio MCP WSBridge module implements the MCPWebSocketBridge class, which serves as a bidirectional communication bridge between WebSocket connections and Model Context Protocol (MCP) operations. It implements the MCPEventHandler protocol to receive MCP events and translates them into WebSocket messages for real-time UI updates.

The bridge manages the complete lifecycle of MCP WebSocket sessions, including message routing, operation dispatching, session tracking, and user interaction handling (elicitation). It provides a clean abstraction that allows the UI to interact with MCP servers through a simple WebSocket interface.

Usage

This module is used to:

• Bridge WebSocket connections to MCP protocol operations
• Stream MCP events to the UI in real-time
• Handle user interactions like elicitation requests
• Manage session state including pending operations
• Dispatch operations to MCPClient asynchronously

The bridge is the central coordination point for MCP WebSocket endpoints in AutoGen Studio.

Code Reference

Source Location

Repository: https://github.com/microsoft/autogen
File Path: python/packages/autogen-studio/autogenstudio/mcp/wsbridge.py
Lines: 215

Class Signature

class MCPWebSocketBridge(MCPEventHandler):
    """Bridges WebSocket connections to MCP operations"""

    def __init__(self, websocket: WebSocket, session_id: str): ...
    async def send_message(self, message: Dict[str, Any]) -> None: ...

    # MCPEventHandler implementation
    async def on_initialized(self, session_id: str, capabilities: Any) -> None: ...
    async def on_operation_result(self, operation: str, data: Dict[str, Any]) -> None: ...
    async def on_operation_error(self, operation: str, error: str) -> None: ...
    async def on_mcp_activity(self, activity_type: str, message: str, details: Dict[str, Any]) -> None: ...
    async def on_elicitation_request(self, request_id: str, message: str, requested_schema: Any) -> None: ...

    def set_mcp_client(self, mcp_client: MCPClient) -> None: ...
    async def handle_websocket_message(self, message: Dict[str, Any]) -> None: ...
    async def run(self) -> None: ...
    def stop(self) -> None: ...

Import Statement

from autogenstudio.mcp.wsbridge import MCPWebSocketBridge

I/O Contract

Constructor Parameters

Parameter	Type	Description
websocket	WebSocket	FastAPI WebSocket connection instance
session_id	str	Unique identifier for this MCP session

Instance Attributes

Attribute	Type	Description
websocket	WebSocket	FastAPI WebSocket connection
session_id	str	Unique session identifier
mcp_client	Optional[MCPClient]	Associated MCP client (set after initialization)
pending_elicitations	Dict[str, asyncio.Future]	Map of request_id to Future for pending user input requests
_running	bool	Flag controlling the message loop

WebSocket Message Types (Incoming)

Message Type	Fields	Description
operation	operation, tool_name, arguments, uri, name, etc.	MCP operation request from UI
ping	-	Keepalive ping
elicitation_response	request_id, action, data	User response to elicitation request

WebSocket Message Types (Outgoing)

Message Type	Fields	Description
initialized	session_id, capabilities, timestamp	Session initialization complete
operation_result	operation, data, timestamp	Operation completed successfully
operation_error	operation, error, timestamp	Operation failed
mcp_activity	activity_type, message, details, session_id, timestamp	General MCP protocol activity
elicitation_request	request_id, message, requestedSchema, session_id, timestamp	Tool requesting user input
pong	timestamp	Response to ping
error	error, timestamp	General error message

Methods

Method	Parameters	Returns	Description
send_message	message: Dict[str, Any]	None	Send JSON message through WebSocket
set_mcp_client	mcp_client: MCPClient	None	Associate MCPClient with bridge
handle_websocket_message	message: Dict[str, Any]	None	Process incoming WebSocket message
run	-	None	Main message loop (blocks until connection closes)
stop	-	None	Stop the message loop

Usage Examples

Basic Bridge Setup

from fastapi import WebSocket
from autogenstudio.mcp.wsbridge import MCPWebSocketBridge
from autogenstudio.mcp.client import MCPClient

@app.websocket("/api/mcp/ws/{session_id}")
async def mcp_websocket(websocket: WebSocket, session_id: str):
    await websocket.accept()

    # Create bridge
    bridge = MCPWebSocketBridge(websocket, session_id)

    # Initialize MCP client
    async with ClientSession(read, write) as session:
        mcp_client = MCPClient(session, session_id, bridge)
        bridge.set_mcp_client(mcp_client)

        await mcp_client.initialize()

        # Run message loop
        await bridge.run()

Handling Operations from UI

# UI sends operation request:
{
    "type": "operation",
    "operation": "list_tools"
}

# Bridge receives and dispatches:
await bridge.handle_websocket_message(message)

# Operation runs in background via:
asyncio.create_task(mcp_client.handle_operation(message))

# Result sent back to UI via:
await bridge.on_operation_result("list_tools", {"tools": [...]})

User Elicitation Flow

# 1. Tool requests user input
# Bridge receives callback from MCP:
await bridge.on_elicitation_request(
    request_id="elicit-123",
    message="Please enter your API key",
    requested_schema={"type": "string"}
)

# 2. UI sends response:
{
    "type": "elicitation_response",
    "request_id": "elicit-123",
    "action": "accept",
    "data": {"value": "sk-..."}
}

# 3. Bridge completes Future
await bridge.handle_websocket_message(response_message)
# Future is resolved, tool receives response

Session Tracking

from datetime import datetime, timezone

# Bridge updates session activity on each message
from autogenstudio.web.routes.mcp import active_sessions

# In handle_websocket_message:
if session_id in active_sessions:
    active_sessions[session_id]["last_activity"] = datetime.now(timezone.utc)

Error Handling

# Send error to UI
await bridge.send_message({
    "type": "error",
    "error": "MCP client not initialized",
    "timestamp": datetime.now(timezone.utc).isoformat()
})

Keepalive Ping/Pong

# UI sends:
{"type": "ping"}

# Bridge responds:
{"type": "pong", "timestamp": "2026-02-11T17:00:00Z"}

Complete Session Example

import asyncio
from mcp import ClientSession
from mcp.client.stdio import stdio_client
from autogenstudio.mcp.wsbridge import MCPWebSocketBridge
from autogenstudio.mcp.client import MCPClient
from autogenstudio.mcp.callbacks import (
    create_message_handler,
    create_sampling_callback,
    create_elicitation_callback
)

async def handle_mcp_session(websocket: WebSocket, session_id: str, server_params):
    # Create bridge
    bridge = MCPWebSocketBridge(websocket, session_id)

    # Create callbacks
    message_handler = create_message_handler(bridge)
    sampling_callback = create_sampling_callback(bridge)
    elicitation_callback, _ = create_elicitation_callback(bridge)

    try:
        # Connect to MCP server
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(
                read, write,
                message_handler=message_handler,
                sampling_callback=sampling_callback,
                elicitation_callback=elicitation_callback
            ) as session:
                # Create and set MCP client
                mcp_client = MCPClient(session, session_id, bridge)
                bridge.set_mcp_client(mcp_client)

                # Initialize
                await mcp_client.initialize()

                # Run message loop
                await bridge.run()

    except Exception as e:
        await bridge.send_message({
            "type": "error",
            "error": str(e),
            "timestamp": datetime.now(timezone.utc).isoformat()
        })
    finally:
        bridge.stop()

Implementation Details

Message Serialization

The send_message() method:

• Checks WebSocket connection state (CONNECTED)
• Serializes message via serialize_for_json()
• Handles WebSocket disconnect gracefully
• Logs errors without raising exceptions

Async Operation Dispatch

Critical design decision: Operations are dispatched using asyncio.create_task() to avoid blocking the message loop:

if message_type == "operation":
    # Run in background task to avoid blocking
    asyncio.create_task(self.mcp_client.handle_operation(message))

This preserves exact behavior from original implementation and prevents UI freezing.

Elicitation Response Processing

The _handle_elicitation_response() method:

1. Validates request_id presence
2. Looks up pending Future
3. Parses action: "accept", "decline", or "cancel"
4. Creates appropriate ElicitResult
5. Completes Future if not already done
6. Handles errors by setting ErrorData on Future

Message Loop

The run() method implements the main loop:

1. Continuously receives WebSocket text messages
2. Parses JSON
3. Dispatches to handle_websocket_message()
4. Handles JSON decode errors gracefully
5. Detects WebSocket disconnect via is_websocket_disconnect()
6. Logs non-disconnect errors and re-raises

Session Lifecycle

1. WebSocket connection established
2. Bridge created
3. MCP client created and associated
4. Session initialized
5. Message loop runs
6. On disconnect or error: bridge.stop()
7. Cleanup and close

Related Pages

• Microsoft_Autogen_Studio_MCP_Client - MCPClient that bridge dispatches to
• Microsoft_Autogen_Studio_MCP_Callbacks - Callback factories for MCP events
• Implementation:MCP_Utils - Utility functions used by bridge
• Implementation:MCP_Routes - WebSocket endpoint definitions
• Domain:WebSocket - WebSocket communication patterns
• Domain:MCP - Model Context Protocol domain