WebSockets MCP Math Demo

A reference implementation demonstrating the Model Context Protocol (MCP) over WebSockets using Cloudflare Workers and Durable Objects.

Overview

This repository provides a reference implementation of MCP over WebSockets. It showcases:

Complete MCP client-server architecture
Persistent stateful sessions via Durable Objects
Bidirectional real-time communication over WebSockets
Tool discovery and invocation
Deployment using Cloudflare Workers

Technical Overview

Architecture

This project demonstrates a full MCP implementation over WebSockets with both client and server components:

┌─────────────────┐                 ┌─────────────────┐
│                 │                 │                 │
│  MCP Client     │◄───WebSocket───►│  MCP Server     │
│  (CF Worker)    │                 │  (CF Worker)    │
│                 │      HTTP       │                 │
└─────────────────┘───────────────►└─────────────────┘
                                        │
                                        │ State Persistence
                                        ▼
                                  ┌─────────────────┐
                                  │  Durable Object │
                                  │  (MathAgent)    │
                                  │                 │
                                  └─────────────────┘

Client: A Cloudflare Worker that serves the HTML/JS client application
Server: A Cloudflare Worker that implements the MCP protocol with tool endpoints
Durable Objects: Maintains persistent state for each agent session

WebSocket Implementation

The implementation supports both HTTP and WebSocket transports:

Connection Establishment:
- Client creates an agent via HTTP POST
- Client establishes WebSocket connection to /agent/{agentId}/websocket
- Server maintains the connection in a Durable Object

Message Format:

// Client to Server
{
  "type": "mcp_request",
  "request": {
    "method": "add",
    "params": { "a": 5, "b": 3 }
  }
}

// Server to Client
{
  "type": "mcp_response",
  "result": {
    "result": 8,
    "operation": "add",
    "a": 5,
    "b": 3
  },
  "timestamp": "2023-05-01T12:34:56.789Z"
}

Connection Management:
- Ping/pong heartbeat mechanism
- Automatic reconnection
- Session tracking

Getting Started

Prerequisites

Node.js (v18 or later)
Wrangler (Cloudflare Workers CLI)
Cloudflare account

Installation

Clone this repository:

git clone https://github.com/your-username/mcp-websockets-demo.git
cd mcp-websockets-demo/math-mcp

Install dependencies:
```
npm install
```
Deploy the server:
```
cd server
wrangler deploy
```
Deploy the client:
```
cd ../client
wrangler deploy
```
Note the deployed URLs for both workers, you'll need them to use the application.

Usage

Web Interface

Open the client URL in your browser. The interface allows you to:
- Connect to the MCP server
- Run math operations
- View the WebSocket message log

Programmatic API

You can also use the MCP server programmatically:

HTTP Example:

// Create an agent
const agentResponse = await fetch('https://your-server.workers.dev/agent', {
  method: 'POST',
  headers: { 'Content-Type': 'application/json' },
  body: JSON.stringify({ name: 'MathAgent' })
});
const { agentId } = await agentResponse.json();

// Make an MCP request
const result = await fetch('https://your-server.workers.dev/mcp', {
  method: 'POST',
  headers: { 'Content-Type': 'application/json' },
  body: JSON.stringify({
    agentId,
    request: {
      method: 'add',
      params: { a: 5, b: 3 }
    }
  })
});

WebSocket Example:

// Create an agent first via HTTP (see above)

// Establish WebSocket connection
const ws = new WebSocket(`wss://your-server.workers.dev/agent/${agentId}/websocket`);

// Listen for messages
ws.addEventListener('message', (event) => {
  const message = JSON.parse(event.data);
  console.log('Received:', message);
});

// Send an MCP request
ws.addEventListener('open', () => {
  ws.send(JSON.stringify({
    type: 'mcp_request',
    request: {
      method: 'add',
      params: { a: 5, b: 3 }
    }
  }));
});

## WebSocket MCP Protocol Specification

This implementation proposes the following extensions to the MCP protocol for WebSocket support:

### 1. Transport Layer

The WebSocket transport extends MCP with these characteristics:

- **Bidirectional Communication**: Both client and server can initiate messages
- **Persistent Connection**: Long-lived connection reduces overhead
- **Real-time Updates**: Enables server-initiated notifications and streaming results
- **Reduced Latency**: Eliminates HTTP request overhead for frequent interactions

### 2. Message Envelope

All WebSocket messages are wrapped in an envelope with a `type` field:

```json
{
  "type": "message_type",
  "payload": { ... },
  "timestamp": "ISO-8601 timestamp"
}

Common message types include:

mcp_request: Client to server MCP method call
mcp_response: Server to client response
ping/pong: Connection health checks
error: Error notifications
notification: Server-initiated notifications

3. Connection Lifecycle

Initialization: Client creates an agent via HTTP before establishing WebSocket
Connection: Client connects to a WebSocket endpoint specific to the agent
Heartbeat: Client sends periodic pings to maintain the connection
Termination: Either side can close the connection

4. Implementation Considerations

When implementing WebSocket support for MCP:

State Management: Handle reconnection and state recovery
Message Ordering: Implement sequencing for reliable message ordering
Error Handling: Gracefully handle connection errors and message failures
Security: Apply same authentication mechanisms as HTTP transport

Key Code Components

Here are the key components for implementing WebSocket MCP:

Server-Side WebSocket Handling

// Handle WebSocket connections
async function handleWebSocketConnection(request, agentId, env) {
  // Get Durable Object stub for the agent
  const id = env.MATH_AGENT.idFromName(agentId);
  const stub = env.MATH_AGENT.get(id);
  
  // Forward the request to the Durable Object
  return await stub.fetch(request);
}

// Durable Object implementation
export class MathAgent {
  // Handle WebSocket connections
  async handleWebSocketConnection(request) {
    // Create a WebSocket pair
    const pair = new WebSocketPair();
    const [client, server] = Object.values(pair);
    
    // Accept the WebSocket connection
    server.accept();
    
    // Set up event handlers for the WebSocket
    server.addEventListener("message", async (event) => {
      const message = JSON.parse(event.data);
      
      // Handle different message types
      if (message.type === "mcp_request") {
        const result = await this.handleMcpRequest(message.request);
        server.send(JSON.stringify({
          type: "mcp_response",
          result,
          timestamp: new Date().toISOString()
        }));
      }
    });
    
    return new Response(null, {
      status: 101,
      webSocket: client
    });
  }
}

Client-Side WebSocket Usage

// Connect WebSocket
function connectWebSocket(agentId, serverUrl) {
  const ws = new WebSocket(`${serverUrl}/agent/${agentId}/websocket`);
  
  ws.onopen = () => {
    console.log('WebSocket connection established');
  };
  
  ws.onmessage = (event) => {
    const message = JSON.parse(event.data);
    
    // Handle different message types
    if (message.type === 'mcp_response') {
      handleMcpResponse(message);
    }
  };
  
  return ws;
}

// Send MCP request
function sendMcpRequest(ws, method, params) {
  ws.send(JSON.stringify({
    type: 'mcp_request',
    request: {
      method,
      params
    },
    timestamp: new Date().toISOString()
  }));
}

Integration with TypeScript SDK

This reference implementation can be used to extend the MCP TypeScript SDK with WebSocket support:

import { MCPClient } from '@modelcontextprotocol/typescript-sdk';

// Create WebSocket transport
class WebSocketTransport implements MCPTransport {
  private ws: WebSocket;
  private pendingRequests: Map<string, {resolve, reject}>;
  
  constructor(serverUrl: string, agentId: string) {
    this.ws = new WebSocket(`${serverUrl}/agent/${agentId}/websocket`);
    this.pendingRequests = new Map();
    
    this.ws.addEventListener('message', this.handleMessage.bind(this));
  }
  
  async send(method: string, params: any): Promise<any> {
    return new Promise((resolve, reject) => {
      const requestId = crypto.randomUUID();
      
      this.pendingRequests.set(requestId, { resolve, reject });
      
      this.ws.send(JSON.stringify({
        type: 'mcp_request',
        request: { method, params },
        requestId
      }));
    });
  }
  
  private handleMessage(event: MessageEvent) {
    const message = JSON.parse(event.data);
    
    if (message.type === 'mcp_response' && message.requestId) {
      const pending = this.pendingRequests.get(message.requestId);
      if (pending) {
        pending.resolve(message.result);
        this.pendingRequests.delete(message.requestId);
      }
    }
  }
}

// Use transport with MCP client
const transport = new WebSocketTransport('wss://example.com', 'agent-123');
const client = new MCPClient({ transport });

// Use MCP methods as usual
const result = await client.invoke('add', { a: 5, b: 3 });

Advantages of WebSocket MCP

Adding WebSocket support to MCP provides several advantages:

Lower Latency: Perfect for contexts requiring rapid interactions
- High-frequency trading
- Real-time collaborative environments
- Interactive agents requiring quick responses
Bidirectional Communication: Enables new interaction patterns
- Server can push updates without client polling
- Streaming large responses in chunks
- Push notifications for external events
Reduced Network Overhead: More efficient for frequent communications
- No HTTP header overhead for each request
- Connection setup cost amortized over multiple requests
- Especially helpful on mobile networks
Stateful Sessions: Simplifies maintaining conversation context
- Server can associate state with the WebSocket connection
- Client doesn't need to send full context with each request
- Easier to implement streaming responses and partial updates

Challenges and Solutions

WebSockets also introduce challenges that this implementation addresses:

Connection Management:
- Challenge: WebSockets can disconnect unexpectedly
- Solution: Heartbeat mechanism and automatic reconnection
Stateless Workers:
- Challenge: Cloudflare Workers are stateless by default
- Solution: Durable Objects maintain connection state
Request/Response Pairing:
- Challenge: Matching responses to requests over a shared channel
- Solution: Message ID tracking and correlation
Error Handling:
- Challenge: Managing connection failures gracefully
- Solution: Structured error responses and reconnection logic

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

License

This project is licensed under the MIT License - see the LICENSE file for details.