WebSocket MCP协议实现

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

// 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 }
    }
  })
});

// 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"
}

// 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
    });
  }
}

// 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()
  }));
}

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 });