Real-Time AI Apps With WebSockets and LLM Streaming

HTTP Streaming vs WebSockets for AI

HTTP streaming (server-sent events or streaming fetch) is unidirectional: the server streams, the client receives. This covers 90% of AI use cases — chat, generation, summarization.

WebSockets are bidirectional: both client and server can send messages at any time. Use them for: collaborative AI editing (multiple users see the same stream), voice AI with interruption (user can interrupt mid-response), real-time agent status updates, and multi-user AI experiences.

Setting Up a WebSocket Server

Next.js App Router doesn't natively support WebSockets. Use a standalone WebSocket server alongside your Next.js app:

// server/ws.ts — runs as a separate process
import { WebSocketServer, WebSocket } from 'ws';
import { streamText } from 'ai';
import { openai } from '@ai-sdk/openai';

const wss = new WebSocketServer({ port: 8080 });

interface Client {
  ws: WebSocket;
  messages: Array<{ role: string; content: string }>;
}

const clients = new Map<string, Client>();

wss.on('connection', (ws) => {
  const clientId = crypto.randomUUID();
  clients.set(clientId, { ws, messages: [] });
  
  ws.on('message', async (data) => {
    const { type, content } = JSON.parse(data.toString());
    const client = clients.get(clientId)!;
    
    if (type === 'message') {
      client.messages.push({ role: 'user', content });
      
      // Stream LLM response back over WebSocket
      const result = streamText({
        model: openai('gpt-4o-mini'),
        messages: client.messages,
      });
      
      for await (const chunk of result.textStream) {
        if (ws.readyState === WebSocket.OPEN) {
          ws.send(JSON.stringify({ type: 'chunk', content: chunk }));
        }
      }
      
      const fullText = await result.text;
      client.messages.push({ role: 'assistant', content: fullText });
      ws.send(JSON.stringify({ type: 'done' }));
    }
  });
  
  ws.on('close', () => clients.delete(clientId));
});

console.log('WebSocket server running on port 8080');

React Hook for WebSocket AI Chat

// hooks/useWebSocketChat.ts
import { useState, useEffect, useRef, useCallback } from 'react';

export function useWebSocketChat() {
  const [messages, setMessages] = useState<Array<{role: string; content: string}>>([]);
  const [streaming, setStreaming] = useState('');
  const [connected, setConnected] = useState(false);
  const wsRef = useRef<WebSocket | null>(null);

  useEffect(() => {
    const ws = new WebSocket('ws://localhost:8080');
    wsRef.current = ws;
    
    ws.onopen = () => setConnected(true);
    ws.onclose = () => setConnected(false);
    
    ws.onmessage = (event) => {
      const { type, content } = JSON.parse(event.data);
      
      if (type === 'chunk') {
        setStreaming(prev => prev + content);
      } else if (type === 'done') {
        setMessages(prev => [
          ...prev,
          { role: 'assistant', content: streaming }
        ]);
        setStreaming('');
      }
    };
    
    return () => ws.close();
  }, []);

  const sendMessage = useCallback((content: string) => {
    setMessages(prev => [...prev, { role: 'user', content }]);
    wsRef.current?.send(JSON.stringify({ type: 'message', content }));
  }, []);

  return { messages, streaming, connected, sendMessage };
}

Deployment Considerations

WebSocket servers need sticky sessions in load-balanced deployments — all messages from a client must hit the same server instance (which holds the conversation state). Options:

• Single instance deployments (Railway, Fly.io): no problem, no stickiness needed
• Multi-instance: use Redis pub/sub to share state across instances
• Managed WebSocket services: Ably, Pusher, or Liveblocks handle this for you

When to Use Each Approach

For most AI chat: use HTTP streaming with the AI SDK — simpler, works everywhere, scales easily. Reach for WebSockets only when you genuinely need bidirectional communication or multi-user synchronization. The added complexity is real and should be justified by the use case.