Building Real-Time Chat in Astro with Cloudflare Durable Objects and WebSocket Hibernation

import { defineConfig } from 'astro/config';
import cloudflare from '@astrojs/cloudflare';

export default defineConfig({
  output: 'server',
  adapter: cloudflare({
    workerEntryPoint: {
      path: './src/worker.ts',
      namedExports: ['ChatRoom']
    }
  })
});

import { handle } from '@astrojs/cloudflare/handler';
import type { SSRManifest } from 'astro';
import { App } from 'astro/app';
import { DurableObject } from 'cloudflare:workers';

interface SessionAttachment {
  id: string
  userId: string
  username: string
}

interface ChatMessage {
  id: string
  type: 'message' | 'join' | 'leave' | 'presence'
  userId: string
  username: string
  content?: string
  timestamp: number
}

class ChatRoom extends DurableObject<ENV> {
  // Map of WebSocket -> session data
  // When the DO hibernates, this gets reconstructed in the constructor
  private sessions: Map<WebSocket, SessionAttachment>
  private messageHistory: ChatMessage[]

  constructor(ctx: DurableObjectState, env: ENV) {
    super(ctx, env)

    this.sessions = new Map()
    this.messageHistory = []

    // Restore hibernating WebSocket connections
    // When the DO wakes up, we need to restore session data from attachments
    this.ctx.getWebSockets().forEach((ws) => {
      const attachment = ws.deserializeAttachment() as SessionAttachment | null
      if (attachment) {
        this.sessions.set(ws, attachment)
      }
    })

    // Set up automatic ping/pong responses
    // This keeps connections alive without waking the DO
    this.ctx.setWebSocketAutoResponse(
      new WebSocketRequestResponsePair('ping', 'pong')
    )

    // Load message history from storage on initialization
    this.ctx.blockConcurrencyWhile(async () => {
      const stored = await this.ctx.storage.get<ChatMessage[]>('messages')
      if (stored) this.messageHistory = stored
    })
  }

  /**
   * Handle HTTP requests to this Durable Object
   * This is called when a client wants to establish a WebSocket connection
   */
  async fetch(request: Request): Promise<Response> {
    // Parse the URL to get query parameters (userId and username)
    const url = new URL(request.url)
    const userId = url.searchParams.get('userId')
    const username = url.searchParams.get('username')
    if (!userId || !username)
      return new Response('Missing userId or username', { status: 400 })
    // Expect a WebSocket upgrade request
    const upgradeHeader = request.headers.get('Upgrade')
    if (upgradeHeader !== 'websocket')
      return new Response('Expected WebSocket upgrade', { status: 426 })
    // Create a WebSocket pair (client and server)
    const pair = new WebSocketPair()
    const [client, server] = Object.values(pair)

    // Accept the WebSocket connection using the Hibernation API
    // Unlike server.accept(), this allows the DO to hibernate while
    // keeping the WebSocket connection open
    this.ctx.acceptWebSocket(server)

    // Generate a unique session ID
    const id = crypto.randomUUID()

    // Create session attachment data
    const attachment: SessionAttachment = { id, userId, username }

    // Serialize the attachment to the WebSocket
    // This data persists across hibernation cycles
    server.serializeAttachment(attachment)

    // Add to active sessions
    this.sessions.set(server, attachment)

    // Send message history to the newly connected client
    this.sendMessageHistory(server)

    // Broadcast join notification to all clients
    this.broadcast({
      id: crypto.randomUUID(),
      type: 'join',
      userId,
      username,
      timestamp: Date.now(),
    })

    // Send current presence to the new user
    this.sendPresence(server)

    // Return the client WebSocket in the response
    return new Response(null, {
      status: 101,
      webSocket: client,
    })
  }

  /**
   * Handle incoming WebSocket messages (Hibernation API)
   * Called when a message is received, even after hibernation
   */
  async webSocketMessage(ws: WebSocket, message: ArrayBuffer | string): Promise<void> {
    // Get session data from the map (or deserialize if just woken)
    let session = this.sessions.get(ws)
    if (!session) {
      session = ws.deserializeAttachment() as SessionAttachment
      if (session) this.sessions.set(ws, session)
    }
    if (!session) return

    try {
      const parsed = JSON.parse(message as string)
      if (parsed.type === 'message' && parsed.content) {
        // Create a chat message
        const chatMessage: ChatMessage = {
          id: crypto.randomUUID(),
          type: 'message',
          userId: session.userId,
          username: session.username,
          content: parsed.content,
          timestamp: Date.now(),
        }

        // Add to history
        this.messageHistory.push(chatMessage)
        // Persist to storage (limit to last 100 messages)
        if (this.messageHistory.length > 100)
          this.messageHistory = this.messageHistory.slice(-100)
        this.ctx.storage.put('messages', this.messageHistory)

        // Broadcast to all connected clients
        this.broadcast(chatMessage)
      }
    } catch (error) {
      console.error('Error handling message:', error)
    }
  }

  /**
   * Handle WebSocket close events (Hibernation API)
   * Called when a client disconnects
   */
  async webSocketClose(
    ws: WebSocket,
    code: number,
    reason: string,
    wasClean: boolean
  ): Promise<void> {
    // Get session data
    const session = this.sessions.get(ws) ||
      (ws.deserializeAttachment() as SessionAttachment | null)

    // Remove from sessions
    this.sessions.delete(ws)

    // Broadcast leave notification
    if (session) {
      this.broadcast({
        id: crypto.randomUUID(),
        type: 'leave',
        userId: session.userId,
        username: session.username,
        timestamp: Date.now(),
      })
    }

    // Close the WebSocket
    ws.close(code, 'Durable Object is closing WebSocket')
  }

  /**
   * Handle WebSocket errors (Hibernation API)
   */
  async webSocketError(ws: WebSocket, error: unknown): Promise<void> {
    console.error('WebSocket error:', error)
    // Treat errors as disconnections
    await this.webSocketClose(ws, 1011, 'WebSocket error', false)
  }

  /**
   * Broadcast a message to all connected clients
   */
  private broadcast(message: ChatMessage): void {
    const messageStr = JSON.stringify(message)

    // Send to all active sessions
    this.sessions.forEach((_, ws) => {
      try {
        ws.send(messageStr)
      } catch (error) {
        // Connection might be closed, will be cleaned up by close handler
        console.error('Error broadcasting to session:', error)
      }
    })
  }

  /**
   * Send message history to a specific WebSocket
   */
  private sendMessageHistory(ws: WebSocket): void {
    const historyMessage = JSON.stringify({
      type: 'history',
      messages: this.messageHistory,
    })
    try {
      ws.send(historyMessage)
    } catch (error) {
      console.error('Error sending history:', error)
    }
  }

  /**
   * Send current user presence to a specific WebSocket
   */
  private sendPresence(ws: WebSocket): void {
    const users = Array.from(this.sessions.values()).map((s) => ({
      userId: s.userId,
      username: s.username,
    }))
    const presenceMessage = JSON.stringify({
      type: 'presence',
      users,
    })

    try {
      ws.send(presenceMessage)
    } catch (error) {
      console.error('Error sending presence:', error)
    }
  }

  /**
   * Handle Durable Object alarm for cleanup tasks
   */
  async alarm(): Promise<void> {
    // Clean up old messages (older than 24 hours)
    const oneDayAgo = Date.now() - 24 * 60 * 60 * 1000
    this.messageHistory = this.messageHistory.filter((msg) => msg.timestamp > oneDayAgo)
    await this.ctx.storage.put('messages', this.messageHistory)
    // Schedule next cleanup in 1 hour
    await this.ctx.storage.setAlarm(Date.now() + 60 * 60 * 1000)
  }
}

export function createExports(manifest: SSRManifest) {
  const app = new App(manifest);
  return {
    default: {
      async fetch(request, env, ctx) {
        // @ts-expect-error - request is not typed correctly
        return handle(manifest, app, request, env, ctx);
      }
    } satisfies ExportedHandler<ENV>,
    ChatRoom,
  }
}

import type { APIRoute } from 'astro'

/**
 * GET /api/chat/:roomId
 *
 * Handles WebSocket upgrade requests and connects clients to the chat room
 */
export const GET: APIRoute = async ({ params, request, locals }) => {
  const { roomId } = params
  if (!roomId)
    return new Response('Room ID is required', { status: 400 })
  // Get the ChatRoom Durable Object namespace
  const ChatRoomNamespace = locals.runtime.env.ChatRoom
  if (!ChatRoomNamespace)
    return new Response('Chat room binding not found', { status: 500 })
  // Get a Durable Object ID for this room
  // idFromName ensures the same room ID always maps to the same Durable Object
  const durableObjectId = ChatRoomNamespace.idFromName(roomId)
  // Get a stub (reference) to the Durable Object
  const durableObjectStub = ChatRoomNamespace.get(durableObjectId)
  // Forward the request to the Durable Object
  // The Durable Object's fetch method will handle the WebSocket upgrade
  return durableObjectStub.fetch(request)
}

---
const { roomId } = Astro.params;

if (!roomId) return Astro.redirect("/");

// Generate a random user ID and username (in production, use authentication)
const userId = crypto.randomUUID();
const defaultUsername = `User${Math.floor(Math.random() * 10000)}`;
---

<!doctype html>
<html lang="en">
  <head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>Chat Room: {roomId}</title>
  </head>
  <body>
    <div id="app">
      <h1>Chat Room: {roomId}</h1>

      <div id="connection-status">
        <span id="status-indicator">●</span>
        <span id="status-text">Connecting...</span>
      </div>

      <div id="presence-container">
        <strong>Online Users:</strong>
        <div id="user-list"></div>
      </div>

      <div id="messages-container"></div>

      <form id="message-form">
        <input
          type="text"
          id="message-input"
          placeholder="Type a message..."
          autocomplete="off"
          required
        />
        <button type="submit">Send</button>
      </form>
    </div>

    <script define:vars={{ roomId, userId, defaultUsername }}>
      // WebSocket connection
      let ws = null;
      let reconnectAttempts = 0;
      const maxReconnectAttempts = 5;
      let username = defaultUsername;

      // DOM elements
      const statusIndicator = document.getElementById("status-indicator");
      const statusText = document.getElementById("status-text");
      const messagesContainer =
        document.getElementById("messages-container");
      const messageForm = document.getElementById("message-form");
      const messageInput = document.getElementById("message-input");
      const userList = document.getElementById("user-list");

      /**
       * Connect to the WebSocket server
       */
      function connect() {
        // Determine WebSocket protocol (ws:// or wss://)
        const protocol =
          window.location.protocol === "https:" ? "wss:" : "ws:";
        const host = window.location.host;

        // Build WebSocket URL
        const wsUrl = `${protocol}//${host}/api/chat/${roomId}?userId=${userId}&username=${encodeURIComponent(username)}`;

        // Create WebSocket connection
        ws = new WebSocket(wsUrl);

        // Connection opened
        ws.addEventListener("open", () => {
          console.log("WebSocket connected");
          reconnectAttempts = 0;
          updateConnectionStatus("connected");
        });

        // Listen for messages
        ws.addEventListener("message", (event) => {
          handleMessage(event.data);
        });

        // Connection closed
        ws.addEventListener("close", () => {
          console.log("WebSocket closed");
          updateConnectionStatus("disconnected");
          attemptReconnect();
        });

        // Connection error
        ws.addEventListener("error", (error) => {
          console.error("WebSocket error:", error);
          updateConnectionStatus("error");
        });
      }

      /**
       * Handle incoming WebSocket messages
       */
      function handleMessage(data) {
        try {
          const message = JSON.parse(data);

          switch (message.type) {
            case "history":
              // Received message history
              message.messages.forEach((msg) =>
                displayMessage(msg),
              );
              break;

            case "message":
              // Received a new chat message
              displayMessage(message);
              break;

            case "join":
              // User joined
              displaySystemMessage(
                `${message.username} joined the room`,
              );
              break;

            case "leave":
              // User left
              displaySystemMessage(
                `${message.username} left the room`,
              );
              break;

            case "presence":
              // Update user list
              updateUserList(message.users);
              break;

            default:
              console.warn("Unknown message type:", message.type);
          }
        } catch (error) {
          console.error("Error parsing message:", error);
        }
      }

      /**
       * Display a chat message in the UI
       */
      function displayMessage(message) {
        const messageEl = document.createElement("div");
        messageEl.className = "message";

        const isOwnMessage = message.userId === userId;
        if (isOwnMessage) {
          messageEl.classList.add("own-message");
        }

        const timestamp = new Date(
          message.timestamp,
        ).toLocaleTimeString();

        messageEl.innerHTML = `
          <div class="message-header">
            <strong>${escapeHtml(message.username)}</strong>
            <span class="message-time">${timestamp}</span>
          </div>
          <div class="message-content">${escapeHtml(message.content)}</div>
        `;

        messagesContainer.appendChild(messageEl);
        messagesContainer.scrollTop = messagesContainer.scrollHeight;
      }

      /**
       * Display a system message (join/leave notifications)
       */
      function displaySystemMessage(text) {
        const messageEl = document.createElement("div");
        messageEl.className = "system-message";
        messageEl.textContent = text;
        messagesContainer.appendChild(messageEl);
        messagesContainer.scrollTop = messagesContainer.scrollHeight;
      }

      /**
       * Update the connection status indicator
       */
      function updateConnectionStatus(status) {
        if (status === "connected") {
          statusIndicator.style.color = "#10b981";
          statusText.textContent = "Connected";
        } else if (status === "disconnected") {
          statusIndicator.style.color = "#ef4444";
          statusText.textContent = "Disconnected";
        } else if (status === "error") {
          statusIndicator.style.color = "#f59e0b";
          statusText.textContent = "Connection error";
        } else {
          statusIndicator.style.color = "#6b7280";
          statusText.textContent = "Connecting...";
        }
      }

      /**
       * Update the user list
       */
      function updateUserList(users) {
        userList.innerHTML = "";
        users.forEach((user) => {
          const userEl = document.createElement("span");
          userEl.className = "user-badge";
          userEl.textContent = user.username;
          userList.appendChild(userEl);
        });
      }

      /**
       * Attempt to reconnect if connection is lost
       */
      function attemptReconnect() {
        if (reconnectAttempts >= maxReconnectAttempts) {
          displaySystemMessage(
            "Failed to reconnect. Please refresh the page.",
          );
          return;
        }

        reconnectAttempts++;
        const delay = Math.min(
          1000 * Math.pow(2, reconnectAttempts),
          10000,
        );

        displaySystemMessage(
          `Reconnecting in ${delay / 1000} seconds... (attempt ${reconnectAttempts}/${maxReconnectAttempts})`,
        );

        setTimeout(() => {
          connect();
        }, delay);
      }

      /**
       * Send a message
       */
      function sendMessage(content) {
        if (!ws || ws.readyState !== WebSocket.OPEN) {
          displaySystemMessage("Cannot send message: not connected");
          return;
        }

        ws.send(
          JSON.stringify({
            type: "message",
            content: content.trim(),
          }),
        );
      }

      /**
       * Escape HTML to prevent XSS
       */
      function escapeHtml(text) {
        const div = document.createElement("div");
        div.textContent = text;
        return div.innerHTML;
      }

      // Handle form submission
      messageForm.addEventListener("submit", (e) => {
        e.preventDefault();

        const content = messageInput.value.trim();
        if (content) {
          sendMessage(content);
          messageInput.value = "";
        }
      });

      // Prompt for username on page load
      const savedUsername = localStorage.getItem("chat-username");
      if (savedUsername) {
        username = savedUsername;
      } else {
        const promptedUsername = prompt(
          "Enter your username:",
          defaultUsername,
        );
        if (promptedUsername && promptedUsername.trim()) {
          username = promptedUsername.trim();
          localStorage.setItem("chat-username", username);
        }
      }

      // Connect on page load
      connect();

      // Clean up on page unload
      window.addEventListener("beforeunload", () => {
        if (ws) {
          ws.close();
        }
      });
    </script>
  </body>
</html>

// In the constructor, load message history from storage
this.ctx.blockConcurrencyWhile(async () => {
  const stored = await this.ctx.storage.get<ChatMessage[]>('messages')
  if (stored) {
    this.messageHistory = stored
  }
})

// In webSocketMessage(), when a new message arrives, persist it
this.messageHistory.push(chatMessage)

// Limit to last 100 messages
if (this.messageHistory.length > 100) {
  this.messageHistory = this.messageHistory.slice(-100)
}

// Save to storage
this.ctx.storage.put('messages', this.messageHistory)

// When a user connects (in fetch method)
this.sessions.set(server, attachment)

// Send current presence to the new user
const users = Array.from(this.sessions.values()).map((s) => ({
  userId: s.userId,
  username: s.username,
}))

// When a user disconnects (in webSocketClose method)
this.sessions.delete(ws)

// Broadcast leave notification
this.broadcast({
  type: 'leave',
  userId: session.userId,
  username: session.username,
  timestamp: Date.now(),
})

function attemptReconnect() {
  if (reconnectAttempts >= maxReconnectAttempts) {
    displaySystemMessage('Failed to reconnect. Please refresh the page.')
    return
  }

  reconnectAttempts++
  // Exponential backoff: 1s, 2s, 4s, 8s, 10s (max)
  const delay = Math.min(1000 * Math.pow(2, reconnectAttempts), 10000)

  setTimeout(() => {
    connect()
  }, delay)
}