WebSocket Patterns for AI Applications

This guide covers advanced WebSocket patterns for building real-time, bidirectional AI applications with Claude Code, including collaborative features, multi-agent coordination, and interactive AI sessions.

Overview

WebSockets provide full-duplex communication channels over a single TCP connection, making them ideal for:

• Real-time collaborative AI coding
• Interactive AI chat sessions with context
• Multi-agent coordination
• Live debugging and pair programming
• Bidirectional data streaming

Core WebSocket Implementation

Server Setup with TypeScript

import { WebSocketServer, WebSocket } from 'ws';
import { Anthropic } from '@anthropic-ai/sdk';
import { v4 as uuidv4 } from 'uuid';
 
interface Session {
  id: string;
  userId: string;
  ws: WebSocket;
  context: Array<{ role: string; content: string }>;
  activeStream?: AbortController;
  metadata: {
    connectedAt: Date;
    lastActivity: Date;
    model: string;
  };
}
 
class AIWebSocketServer {
  private wss: WebSocketServer;
  private sessions: Map<string, Session> = new Map();
  private anthropic: Anthropic;
  private heartbeatInterval: NodeJS.Timeout;
 
  constructor(port: number) {
    this.wss = new WebSocketServer({ port });
    this.anthropic = new Anthropic();
    this.setupServer();
    this.startHeartbeat();
  }
 
  private setupServer() {
    this.wss.on('connection', (ws: WebSocket, request) => {
      const sessionId = uuidv4();
      const session = this.createSession(sessionId, ws);
      
      this.sessions.set(sessionId, session);
      
      // Send welcome message
      this.sendMessage(ws, {
        type: 'connection',
        sessionId,
        timestamp: new Date().toISOString()
      });
 
      // Set up event handlers
      ws.on('message', (data) => this.handleMessage(session, data));
      ws.on('close', () => this.handleDisconnect(sessionId));
      ws.on('error', (error) => this.handleError(session, error));
      ws.on('pong', () => this.handlePong(session));
    });
  }
 
  private createSession(id: string, ws: WebSocket): Session {
    return {
      id,
      userId: '', // Set after authentication
      ws,
      context: [],
      metadata: {
        connectedAt: new Date(),
        lastActivity: new Date(),
        model: 'claude-opus-4-20250514'
      }
    };
  }
 
  private async handleMessage(session: Session, data: WebSocket.Data) {
    try {
      const message = JSON.parse(data.toString());
      session.metadata.lastActivity = new Date();
 
      switch (message.type) {
        case 'auth':
          await this.handleAuth(session, message);
          break;
        case 'query':
          await this.handleQuery(session, message);
          break;
        case 'stop':
          this.handleStop(session);
          break;
        case 'context':
          this.handleContextUpdate(session, message);
          break;
        case 'collaborate':
          await this.handleCollaboration(session, message);
          break;
        default:
          this.sendError(session, `Unknown message type: ${message.type}`);
      }
    } catch (error) {
      this.sendError(session, error.message);
    }
  }
 
  private async handleQuery(session: Session, message: any) {
    // Abort any active stream
    if (session.activeStream) {
      session.activeStream.abort();
    }
 
    const controller = new AbortController();
    session.activeStream = controller;
 
    try {
      const stream = await this.anthropic.messages.create({
        model: session.metadata.model,
        messages: [
          ...session.context,
          { role: 'user', content: message.content }
        ],
        max_tokens: message.maxTokens || 1024,
        stream: true
      }, {
        signal: controller.signal
      });
 
      let fullResponse = '';
      
      for await (const event of stream) {
        if (event.type === 'content_block_delta') {
          fullResponse += event.delta.text;
          
          this.sendMessage(session.ws, {
            type: 'stream_delta',
            delta: event.delta.text,
            index: event.index
          });
        } else if (event.type === 'message_stop') {
          // Update context
          session.context.push(
            { role: 'user', content: message.content },
            { role: 'assistant', content: fullResponse }
          );
          
          this.sendMessage(session.ws, {
            type: 'stream_complete',
            fullResponse,
            usage: event.usage
          });
        }
      }
    } catch (error) {
      if (error.name !== 'AbortError') {
        this.sendError(session, error.message);
      }
    } finally {
      session.activeStream = undefined;
    }
  }
 
  private sendMessage(ws: WebSocket, message: any) {
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify(message));
    }
  }
 
  private sendError(session: Session, error: string) {
    this.sendMessage(session.ws, {
      type: 'error',
      error,
      timestamp: new Date().toISOString()
    });
  }
 
  private startHeartbeat() {
    this.heartbeatInterval = setInterval(() => {
      this.sessions.forEach((session) => {
        if (session.ws.readyState === WebSocket.OPEN) {
          session.ws.ping();
        }
      });
    }, 30000); // 30 seconds
  }
 
  shutdown() {
    clearInterval(this.heartbeatInterval);
    this.sessions.forEach(session => {
      session.ws.close();
    });
    this.wss.close();
  }
}

Robust Client Implementation

class AIWebSocketClient {
  private ws: WebSocket | null = null;
  private url: string;
  private sessionId: string | null = null;
  private reconnectAttempts = 0;
  private maxReconnectAttempts = 5;
  private reconnectDelay = 1000;
  private pingInterval: number | null = null;
  private messageQueue: any[] = [];
  private isConnected = false;
  
  constructor(url: string) {
    this.url = url;
  }
 
  connect(options: {
    onOpen?: () => void;
    onMessage?: (data: any) => void;
    onError?: (error: any) => void;
    onClose?: () => void;
  }) {
    this.ws = new WebSocket(this.url);
    
    this.ws.onopen = () => {
      this.isConnected = true;
      this.reconnectAttempts = 0;
      this.flushMessageQueue();
      this.startPingInterval();
      options.onOpen?.();
    };
    
    this.ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data);
        
        if (data.type === 'connection') {
          this.sessionId = data.sessionId;
        }
        
        options.onMessage?.(data);
      } catch (error) {
        console.error('Failed to parse message:', error);
      }
    };
    
    this.ws.onerror = (error) => {
      options.onError?.(error);
    };
    
    this.ws.onclose = () => {
      this.isConnected = false;
      this.stopPingInterval();
      options.onClose?.();
      this.handleReconnect(options);
    };
  }
 
  private handleReconnect(options: any) {
    if (this.reconnectAttempts >= this.maxReconnectAttempts) {
      console.error('Max reconnection attempts reached');
      return;
    }
 
    this.reconnectAttempts++;
    const delay = this.reconnectDelay * Math.pow(2, this.reconnectAttempts - 1);
    
    console.log(`Reconnecting in ${delay}ms (attempt ${this.reconnectAttempts})`);
    
    setTimeout(() => {
      this.connect(options);
    }, delay);
  }
 
  send(message: any) {
    const msg = typeof message === 'string' ? message : JSON.stringify(message);
    
    if (this.isConnected && this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(msg);
    } else {
      // Queue message for later delivery
      this.messageQueue.push(msg);
    }
  }
 
  private flushMessageQueue() {
    while (this.messageQueue.length > 0 && this.isConnected) {
      const msg = this.messageQueue.shift();
      this.ws?.send(msg);
    }
  }
 
  private startPingInterval() {
    this.pingInterval = window.setInterval(() => {
      if (this.ws?.readyState === WebSocket.OPEN) {
        this.ws.send(JSON.stringify({ type: 'ping' }));
      }
    }, 25000); // 25 seconds
  }
 
  private stopPingInterval() {
    if (this.pingInterval) {
      clearInterval(this.pingInterval);
      this.pingInterval = null;
    }
  }
 
  disconnect() {
    this.reconnectAttempts = this.maxReconnectAttempts; // Prevent reconnection
    this.stopPingInterval();
    
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
  }
}

Advanced Patterns

1. Collaborative AI Coding Session

interface CollaborativeSession {
  id: string;
  participants: Map<string, Participant>;
  sharedContext: SharedContext;
  aiAssistant: AIAssistant;
}
 
interface Participant {
  id: string;
  name: string;
  ws: WebSocket;
  cursor: CursorPosition;
  permissions: Set<Permission>;
}
 
class CollaborativeAIServer {
  private sessions: Map<string, CollaborativeSession> = new Map();
  
  async createCollaborativeSession(hostId: string): Promise<string> {
    const sessionId = uuidv4();
    const session: CollaborativeSession = {
      id: sessionId,
      participants: new Map(),
      sharedContext: new SharedContext(),
      aiAssistant: new AIAssistant()
    };
    
    this.sessions.set(sessionId, session);
    
    // Set up AI context observer
    session.sharedContext.on('change', async (change) => {
      await this.handleContextChange(session, change);
    });
    
    return sessionId;
  }
 
  async handleCollaborativeAction(
    sessionId: string,
    participantId: string,
    action: any
  ) {
    const session = this.sessions.get(sessionId);
    if (!session) return;
    
    const participant = session.participants.get(participantId);
    if (!participant) return;
 
    switch (action.type) {
      case 'code_edit':
        await this.handleCodeEdit(session, participant, action);
        break;
      
      case 'ai_request':
        await this.handleAIRequest(session, participant, action);
        break;
      
      case 'cursor_move':
        this.handleCursorMove(session, participant, action);
        break;
      
      case 'selection_change':
        this.handleSelectionChange(session, participant, action);
        break;
    }
  }
 
  private async handleAIRequest(
    session: CollaborativeSession,
    participant: Participant,
    action: any
  ) {
    // Check permissions
    if (!participant.permissions.has('ai_request')) {
      this.sendToParticipant(participant, {
        type: 'error',
        error: 'Insufficient permissions'
      });
      return;
    }
 
    // Get relevant context
    const context = await session.sharedContext.getRelevantContext(action.query);
    
    // Stream AI response to all participants
    const stream = await session.aiAssistant.streamResponse({
      query: action.query,
      context,
      requestedBy: participant.name
    });
 
    for await (const chunk of stream) {
      this.broadcastToSession(session, {
        type: 'ai_stream',
        chunk,
        requestedBy: participant.id
      });
    }
  }
 
  private broadcastToSession(
    session: CollaborativeSession,
    message: any,
    excludeParticipant?: string
  ) {
    session.participants.forEach((participant, id) => {
      if (id !== excludeParticipant) {
        this.sendToParticipant(participant, message);
      }
    });
  }
}
 
// Client-side collaborative features
class CollaborativeClient {
  private ws: AIWebSocketClient;
  private editor: CodeEditor;
  private cursors: Map<string, RemoteCursor> = new Map();
  
  constructor(editorElement: HTMLElement) {
    this.editor = new CodeEditor(editorElement);
    this.setupCollaborativeFeatures();
  }
 
  private setupCollaborativeFeatures() {
    // Track local cursor movements
    this.editor.on('cursorMove', (position) => {
      this.ws.send({
        type: 'cursor_move',
        position
      });
    });
 
    // Track code changes
    this.editor.on('change', (change) => {
      this.ws.send({
        type: 'code_edit',
        change
      });
    });
 
    // Handle remote updates
    this.ws = new AIWebSocketClient('wss://collab.example.com');
    this.ws.connect({
      onMessage: (data) => {
        switch (data.type) {
          case 'cursor_move':
            this.updateRemoteCursor(data.participantId, data.position);
            break;
          
          case 'code_edit':
            this.applyRemoteEdit(data.change);
            break;
          
          case 'ai_stream':
            this.handleAIStream(data);
            break;
        }
      }
    });
  }
 
  requestAIAssistance(query: string) {
    this.ws.send({
      type: 'ai_request',
      query
    });
  }
}

2. Multi-Agent Coordination Pattern

interface AgentMessage {
  from: string;
  to: string;
  type: 'request' | 'response' | 'broadcast';
  payload: any;
  correlationId: string;
}
 
class MultiAgentCoordinator {
  private agents: Map<string, Agent> = new Map();
  private messageRouter: MessageRouter;
  
  constructor() {
    this.messageRouter = new MessageRouter();
  }
 
  registerAgent(agent: Agent) {
    this.agents.set(agent.id, agent);
    
    agent.on('message', (message: AgentMessage) => {
      this.routeMessage(message);
    });
  }
 
  private async routeMessage(message: AgentMessage) {
    if (message.type === 'broadcast') {
      // Broadcast to all agents except sender
      this.agents.forEach((agent, id) => {
        if (id !== message.from) {
          agent.receive(message);
        }
      });
    } else {
      // Direct message
      const targetAgent = this.agents.get(message.to);
      if (targetAgent) {
        targetAgent.receive(message);
      }
    }
    
    // Log for debugging and auditing
    await this.messageRouter.log(message);
  }
 
  async coordinateTask(task: ComplexTask) {
    // Decompose task into subtasks
    const subtasks = await this.decomposeTask(task);
    
    // Create execution plan
    const plan = this.createExecutionPlan(subtasks);
    
    // Execute plan with agent coordination
    const results = await this.executePlan(plan);
    
    return this.aggregateResults(results);
  }
 
  private async executePlan(plan: ExecutionPlan) {
    const results: Map<string, any> = new Map();
    
    for (const step of plan.steps) {
      const agent = this.selectAgent(step.requirements);
      
      const result = await this.executeWithAgent(agent, step, (progress) => {
        // Stream progress updates
        this.broadcastProgress({
          step: step.id,
          agent: agent.id,
          progress
        });
      });
      
      results.set(step.id, result);
      
      // Update context for dependent steps
      this.updateExecutionContext(plan, step.id, result);
    }
    
    return results;
  }
}
 
// Specialized agents
class CodeAnalysisAgent extends Agent {
  async analyze(code: string, options: AnalysisOptions) {
    const ws = this.getWebSocket();
    
    return new Promise((resolve, reject) => {
      const correlationId = uuidv4();
      
      ws.send({
        type: 'analyze_code',
        code,
        options,
        correlationId
      });
      
      this.once(`response:${correlationId}`, (result) => {
        resolve(result);
      });
      
      setTimeout(() => {
        reject(new Error('Analysis timeout'));
      }, 30000);
    });
  }
}

3. Interactive Debugging Session

class InteractiveDebugSession {
  private ws: WebSocket;
  private debugger: AIDebugger;
  private breakpoints: Set<Breakpoint> = new Set();
  private watchExpressions: Map<string, WatchExpression> = new Map();
  
  async startDebugSession(code: string, error: Error) {
    // Initialize AI debugger
    this.debugger = new AIDebugger({
      model: 'claude-opus-4-20250514',
      context: { code, error }
    });
    
    // Connect to debug server
    this.ws = new WebSocket('wss://debug.example.com');
    
    this.ws.on('message', async (data) => {
      const message = JSON.parse(data.toString());
      
      switch (message.type) {
        case 'breakpoint_hit':
          await this.handleBreakpoint(message);
          break;
        
        case 'expression_evaluated':
          this.updateWatchExpression(message);
          break;
        
        case 'ai_suggestion':
          this.displayAISuggestion(message);
          break;
      }
    });
    
    // Start interactive session
    await this.initializeSession();
  }
 
  private async handleBreakpoint(message: any) {
    const { location, variables, callStack } = message;
    
    // Get AI analysis of current state
    const analysis = await this.debugger.analyzeState({
      location,
      variables,
      callStack
    });
    
    // Stream analysis to client
    this.ws.send(JSON.stringify({
      type: 'ai_analysis',
      analysis,
      suggestions: analysis.suggestions
    }));
    
    // Set up interactive prompt
    this.enableInteractiveMode({
      onQuery: async (query) => {
        const response = await this.debugger.interactiveQuery({
          query,
          currentState: { location, variables, callStack }
        });
        
        return response;
      }
    });
  }
 
  async stepThrough() {
    // AI-guided step-through debugging
    const nextStep = await this.debugger.suggestNextStep();
    
    this.ws.send(JSON.stringify({
      type: 'step',
      action: nextStep.action,
      reasoning: nextStep.reasoning
    }));
  }
}

Performance Optimization

1. Message Compression

import pako from 'pako';
 
class CompressedWebSocket {
  private ws: WebSocket;
  private compressionThreshold = 1024; // 1KB
  
  send(data: any) {
    const message = JSON.stringify(data);
    
    if (message.length > this.compressionThreshold) {
      const compressed = pako.deflate(message);
      
      this.ws.send(new Blob([compressed], {
        type: 'application/octet-stream'
      }));
    } else {
      this.ws.send(message);
    }
  }
  
  private setupDecompression() {
    this.ws.onmessage = async (event) => {
      let data;
      
      if (event.data instanceof Blob) {
        const arrayBuffer = await event.data.arrayBuffer();
        const decompressed = pako.inflate(new Uint8Array(arrayBuffer), {
          to: 'string'
        });
        data = JSON.parse(decompressed);
      } else {
        data = JSON.parse(event.data);
      }
      
      this.handleMessage(data);
    };
  }
}

2. Connection Pooling

class WebSocketPool {
  private pool: WebSocket[] = [];
  private activeConnections: Map<string, WebSocket> = new Map();
  private maxPoolSize = 10;
  private url: string;
  
  constructor(url: string, poolSize: number = 5) {
    this.url = url;
    this.maxPoolSize = poolSize;
    this.initializePool();
  }
  
  private initializePool() {
    for (let i = 0; i < this.maxPoolSize; i++) {
      this.createConnection();
    }
  }
  
  private createConnection(): WebSocket {
    const ws = new WebSocket(this.url);
    
    ws.onopen = () => {
      this.pool.push(ws);
    };
    
    ws.onclose = () => {
      this.removeFromPool(ws);
      // Replace closed connection
      if (this.pool.length < this.maxPoolSize) {
        this.createConnection();
      }
    };
    
    return ws;
  }
  
  async getConnection(): Promise<WebSocket> {
    // Wait for available connection
    while (this.pool.length === 0) {
      await new Promise(resolve => setTimeout(resolve, 100));
    }
    
    return this.pool.shift()!;
  }
  
  releaseConnection(ws: WebSocket) {
    if (ws.readyState === WebSocket.OPEN) {
      this.pool.push(ws);
    } else {
      this.createConnection();
    }
  }
}

Security Best Practices

1. Authentication and Authorization

class SecureWebSocketServer {
  private authenticate = async (ws: WebSocket, token: string): Promise<User> => {
    try {
      const user = await this.verifyToken(token);
      
      if (!user) {
        ws.send(JSON.stringify({
          type: 'auth_error',
          error: 'Invalid token'
        }));
        ws.close(1008, 'Invalid token');
        throw new Error('Authentication failed');
      }
      
      return user;
    } catch (error) {
      ws.close(1008, 'Authentication failed');
      throw error;
    }
  };
  
  private authorizeAction = (user: User, action: string): boolean => {
    const permissions = this.getPermissions(user);
    return permissions.has(action);
  };
  
  handleConnection(ws: WebSocket, request: IncomingMessage) {
    const token = this.extractToken(request);
    
    this.authenticate(ws, token).then(user => {
      const session = this.createAuthenticatedSession(user, ws);
      
      ws.on('message', (data) => {
        const message = JSON.parse(data.toString());
        
        if (!this.authorizeAction(user, message.type)) {
          ws.send(JSON.stringify({
            type: 'error',
            error: 'Unauthorized action'
          }));
          return;
        }
        
        this.handleAuthorizedMessage(session, message);
      });
    }).catch(error => {
      console.error('Auth error:', error);
    });
  }
}

2. Rate Limiting

class RateLimitedWebSocket {
  private rateLimiter: Map<string, RateLimit> = new Map();
  
  private checkRateLimit(sessionId: string): boolean {
    const limit = this.rateLimiter.get(sessionId) || this.createRateLimit();
    
    const now = Date.now();
    const windowStart = now - limit.windowMs;
    
    // Remove old entries
    limit.requests = limit.requests.filter(time => time > windowStart);
    
    if (limit.requests.length >= limit.maxRequests) {
      return false;
    }
    
    limit.requests.push(now);
    this.rateLimiter.set(sessionId, limit);
    
    return true;
  }
  
  private createRateLimit(): RateLimit {
    return {
      requests: [],
      maxRequests: 100,
      windowMs: 60000 // 1 minute
    };
  }
}

Testing WebSocket Applications

// WebSocket test utilities
class WebSocketTestClient {
  private messages: any[] = [];
  private ws: WebSocket;
  
  async connect(url: string): Promise<void> {
    return new Promise((resolve, reject) => {
      this.ws = new WebSocket(url);
      
      this.ws.onopen = () => resolve();
      this.ws.onerror = (error) => reject(error);
      this.ws.onmessage = (event) => {
        this.messages.push(JSON.parse(event.data));
      };
    });
  }
  
  async sendAndWait(message: any, expectedType: string, timeout = 5000): Promise<any> {
    this.ws.send(JSON.stringify(message));
    
    const start = Date.now();
    
    while (Date.now() - start < timeout) {
      const response = this.messages.find(m => m.type === expectedType);
      if (response) {
        return response;
      }
      await new Promise(resolve => setTimeout(resolve, 100));
    }
    
    throw new Error(`Timeout waiting for ${expectedType}`);
  }
  
  disconnect() {
    this.ws.close();
  }
}
 
// Example test
describe('WebSocket AI Server', () => {
  let server: AIWebSocketServer;
  let client: WebSocketTestClient;
  
  beforeEach(async () => {
    server = new AIWebSocketServer(8080);
    client = new WebSocketTestClient();
    await client.connect('ws://localhost:8080');
  });
  
  afterEach(() => {
    client.disconnect();
    server.shutdown();
  });
  
  it('should handle AI queries', async () => {
    const response = await client.sendAndWait({
      type: 'query',
      content: 'Hello, Claude!'
    }, 'stream_complete');
    
    expect(response.fullResponse).toBeTruthy();
    expect(response.usage).toBeDefined();
  });
});

🔗 Related Resources

• Streaming Deep Dive
• SSE Implementation Guide
• Streaming Performance
• Real-Time Collaboration

📚 External References

• MDN: WebSocket API
• WebSocket Protocol RFC
• ws: WebSocket library for Node.js