@futurity/chat-react 0.0.1

package/src/WebSocketConnection.ts

@@ -0,0 +1,284 @@
+export type ConnectionState = "disconnected" | "connecting" | "connected";
+
+export type WebSocketConnectionOptions = {
+  /** WebSocket URL (can be relative, will be converted to ws/wss) */
+  url: string;
+  /** Called when a parsed message is received */
+  onMessage: (message: unknown) => void;
+  /** Called when connection state changes */
+  onConnectionChange?: (state: ConnectionState) => void;
+  /** Called when an error occurs */
+  onError?: (error: Event) => void;
+  /** Heartbeat interval in ms (default: 5000) */
+  heartbeatInterval?: number;
+  /** How long to wait for pong before considering connection dead (default: 10000) */
+  heartbeatTimeout?: number;
+  /** Initial reconnection delay in ms (default: 1000) */
+  initialReconnectDelay?: number;
+  /** Maximum reconnection delay in ms (default: 30000) */
+  maxReconnectDelay?: number;
+  /** Log prefix for debugging */
+  debugPrefix?: string;
+};
+
+const DEFAULT_HEARTBEAT_INTERVAL = 5_000;
+const DEFAULT_HEARTBEAT_TIMEOUT = 10_000;
+const DEFAULT_INITIAL_RECONNECT_DELAY = 1_000;
+const DEFAULT_MAX_RECONNECT_DELAY = 30_000;
+
+/**
+ * Framework-agnostic WebSocket connection with automatic reconnection and heartbeat.
+ *
+ * Handles:
+ * - Connection lifecycle (connect / disconnect / ensureConnected)
+ * - Reconnection with exponential backoff
+ * - Heartbeat ping/pong to detect stale connections
+ * - Message queue while disconnected
+ * - JSON serialization on send, JSON parsing on receive
+ */
+export class WebSocketConnection {
+  private ws: WebSocket | null = null;
+  private reconnectTimeout: ReturnType<typeof setTimeout> | null = null;
+  private heartbeatIntervalTimer: ReturnType<typeof setInterval> | null = null;
+  private heartbeatTimeoutTimer: ReturnType<typeof setTimeout> | null = null;
+  private reconnectDelay: number;
+  private pendingMessages: unknown[] = [];
+  private intentionalClose = false;
+  private awaitingPong = false;
+  private connectPromise: {
+    resolve: (connected: boolean) => void;
+  } | null = null;
+
+  private _state: ConnectionState = "disconnected";
+
+  private readonly url: string;
+  private readonly heartbeatInterval: number;
+  private readonly heartbeatTimeout: number;
+  private readonly initialReconnectDelay: number;
+  private readonly maxReconnectDelay: number;
+  private readonly debugPrefix: string;
+
+  onMessage: (message: unknown) => void;
+  onConnectionChange?: (state: ConnectionState) => void;
+  onError?: (error: Event) => void;
+
+  get state(): ConnectionState {
+    return this._state;
+  }
+
+  get isConnected(): boolean {
+    return this._state === "connected";
+  }
+
+  constructor(options: WebSocketConnectionOptions) {
+    this.url = options.url;
+    this.onMessage = options.onMessage;
+    this.onConnectionChange = options.onConnectionChange;
+    this.onError = options.onError;
+    this.heartbeatInterval =
+      options.heartbeatInterval ?? DEFAULT_HEARTBEAT_INTERVAL;
+    this.heartbeatTimeout =
+      options.heartbeatTimeout ?? DEFAULT_HEARTBEAT_TIMEOUT;
+    this.initialReconnectDelay =
+      options.initialReconnectDelay ?? DEFAULT_INITIAL_RECONNECT_DELAY;
+    this.maxReconnectDelay =
+      options.maxReconnectDelay ?? DEFAULT_MAX_RECONNECT_DELAY;
+    this.debugPrefix = options.debugPrefix ?? "[WS]";
+    this.reconnectDelay = this.initialReconnectDelay;
+  }
+
+  private updateState(state: ConnectionState): void {
+    this._state = state;
+    this.onConnectionChange?.(state);
+  }
+
+  private clearTimers(): void {
+    if (this.reconnectTimeout) {
+      clearTimeout(this.reconnectTimeout);
+      this.reconnectTimeout = null;
+    }
+    if (this.heartbeatIntervalTimer) {
+      clearInterval(this.heartbeatIntervalTimer);
+      this.heartbeatIntervalTimer = null;
+    }
+    if (this.heartbeatTimeoutTimer) {
+      clearTimeout(this.heartbeatTimeoutTimer);
+      this.heartbeatTimeoutTimer = null;
+    }
+  }
+
+  private sendPing(): void {
+    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) return;
+
+    if (this.awaitingPong) {
+      console.warn(`${this.debugPrefix} No pong received, connection stale`);
+      this.ws.close();
+      return;
+    }
+
+    this.awaitingPong = true;
+    this.ws.send(JSON.stringify({ type: "__ping__" }));
+
+    this.heartbeatTimeoutTimer = setTimeout(() => {
+      if (this.awaitingPong) {
+        console.warn(`${this.debugPrefix} Pong timeout, closing connection`);
+        this.ws?.close();
+      }
+    }, this.heartbeatTimeout);
+  }
+
+  private startHeartbeat(): void {
+    if (this.heartbeatIntervalTimer) {
+      clearInterval(this.heartbeatIntervalTimer);
+    }
+    this.heartbeatIntervalTimer = setInterval(
+      () => this.sendPing(),
+      this.heartbeatInterval,
+    );
+  }
+
+  connect(): void {
+    if (
+      this.ws?.readyState === WebSocket.OPEN ||
+      this.ws?.readyState === WebSocket.CONNECTING
+    ) {
+      return;
+    }
+
+    this.clearTimers();
+    this.intentionalClose = false;
+
+    const wsUrl = this.url.replace(/^http/, "ws");
+    console.log(`${this.debugPrefix} Connecting to:`, wsUrl);
+
+    this.updateState("connecting");
+    const ws = new WebSocket(wsUrl);
+
+    ws.onopen = () => {
+      console.log(`${this.debugPrefix} Connected`);
+      this.updateState("connected");
+      this.reconnectDelay = this.initialReconnectDelay;
+      this.awaitingPong = false;
+
+      for (const msg of this.pendingMessages) {
+        ws.send(JSON.stringify(msg));
+      }
+      this.pendingMessages = [];
+
+      this.startHeartbeat();
+
+      this.connectPromise?.resolve(true);
+      this.connectPromise = null;
+    };
+
+    ws.onmessage = (event) => {
+      try {
+        const data: unknown = JSON.parse(event.data);
+
+        if (
+          typeof data === "object" &&
+          data !== null &&
+          "type" in data &&
+          (data as { type: unknown }).type === "__pong__"
+        ) {
+          this.awaitingPong = false;
+          if (this.heartbeatTimeoutTimer) {
+            clearTimeout(this.heartbeatTimeoutTimer);
+            this.heartbeatTimeoutTimer = null;
+          }
+          return;
+        }
+
+        this.onMessage(data);
+      } catch (error) {
+        console.error(`${this.debugPrefix} Error parsing message:`, error);
+      }
+    };
+
+    ws.onclose = () => {
+      console.log(`${this.debugPrefix} Disconnected`);
+      this.updateState("disconnected");
+      this.clearTimers();
+
+      this.connectPromise?.resolve(false);
+      this.connectPromise = null;
+
+      if (this.intentionalClose) {
+        return;
+      }
+
+      const delay = this.reconnectDelay;
+      console.log(`${this.debugPrefix} Reconnecting in ${delay}ms...`);
+
+      this.reconnectTimeout = setTimeout(() => {
+        if (this.ws === ws) {
+          this.reconnectDelay = Math.min(
+            this.reconnectDelay * 2,
+            this.maxReconnectDelay,
+          );
+          this.connect();
+        }
+      }, delay);
+    };
+
+    ws.onerror = (error) => {
+      if (ws.readyState !== WebSocket.CLOSED) {
+        console.error(`${this.debugPrefix} WebSocket error`);
+        this.onError?.(error);
+      }
+    };
+
+    this.ws = ws;
+  }
+
+  disconnect(): void {
+    this.intentionalClose = true;
+    this.clearTimers();
+
+    if (this.ws) {
+      this.ws.onclose = null;
+      this.ws.onerror = null;
+      this.ws.close();
+      this.ws = null;
+    }
+
+    this.updateState("disconnected");
+  }
+
+  send(message: unknown): void {
+    if (this.ws?.readyState === WebSocket.OPEN) {
+      this.ws.send(JSON.stringify(message));
+    } else {
+      this.pendingMessages.push(message);
+
+      if (this.ws?.readyState !== WebSocket.CONNECTING) {
+        this.reconnectDelay = this.initialReconnectDelay;
+        this.connect();
+      }
+    }
+  }
+
+  ensureConnected(): Promise<boolean> {
+    if (this.ws?.readyState === WebSocket.OPEN) {
+      return Promise.resolve(true);
+    }
+
+    if (this.ws?.readyState === WebSocket.CONNECTING) {
+      return new Promise((resolve) => {
+        this.connectPromise = { resolve };
+      });
+    }
+
+    return new Promise((resolve) => {
+      this.connectPromise = { resolve };
+      this.reconnectDelay = this.initialReconnectDelay;
+      this.connect();
+    });
+  }
+
+  reconnect(): void {
+    this.disconnect();
+    this.reconnectDelay = this.initialReconnectDelay;
+    setTimeout(() => this.connect(), 100);
+  }
+}

package/src/chat-protocol.ts

@@ -0,0 +1,22 @@
+/**
+ * Chat WebSocket protocol helpers.
+ *
+ * All protocol schemas live in @futurity/chat-protocol;
+ * this module provides a lightweight parse helper for the SDK.
+ */
+import {
+  type WsClientCommand,
+  type WsServerMessage,
+  wsServerMessageSchema,
+} from "@futurity/chat-protocol";
+
+export type { WsClientCommand as ClientCommand };
+
+export function parseServerMessage(data: unknown): WsServerMessage | null {
+  const result = wsServerMessageSchema.safeParse(data);
+  if (!result.success) {
+    console.error("Failed to parse server message:", result.error, data);
+    return null;
+  }
+  return result.data;
+}

package/src/index.ts

@@ -0,0 +1,39 @@
+// Core WebSocket class (framework-agnostic)
+
+// Hooks
+export type { ClientCommand } from "./chat-protocol";
+// Protocol helpers
+export { parseServerMessage } from "./chat-protocol";
+// Tree utilities
+export { buildTree, findLatestPath, MessageNode } from "./tree-builder";
+// Types
+export type {
+  ChatMessage,
+  ChatStatus,
+  ClarifyData,
+  MessageMetadata,
+  MessagePart,
+  SendMessageFn,
+  SendMessagePayload,
+  StreamDelta,
+  WsClarifyQuestion,
+  WsClientCommand,
+  WsErrorMessage,
+  WsServerMessage,
+  WsStreamMessage,
+  WsVaultItem,
+} from "./types";
+export { messageMetadataSchema, Z_ChatMessage } from "./types";
+export type {
+  ConnectionState,
+  WebSocketOptions,
+  WebSocketResult,
+} from "./useReconnectingWebSocket";
+export { useReconnectingWebSocket } from "./useReconnectingWebSocket";
+export type {
+  UseStreamChatOptions,
+  UseStreamChatReturn,
+} from "./useStreamChat";
+export { useStreamChat } from "./useStreamChat";
+export type { WebSocketConnectionOptions } from "./WebSocketConnection";
+export { WebSocketConnection } from "./WebSocketConnection";

package/src/tree-builder.ts

@@ -0,0 +1,116 @@
+import type { ChatMessage } from "./types";
+
+export class MessageNode {
+  public readonly id: string;
+  public readonly parent_id: string | null;
+  public readonly message: ChatMessage;
+  private children: MessageNode[] = [];
+
+  constructor(message: ChatMessage) {
+    this.id = message.id;
+    this.message = message;
+    this.parent_id = message.metadata?.parent_id ?? null;
+  }
+
+  addChild(child: MessageNode): void {
+    this.children.push(child);
+  }
+
+  getChildren(): readonly MessageNode[] {
+    return this.children;
+  }
+}
+
+/**
+ * Builds a tree of MessageNodes from a flat list of messages using a two-pass approach.
+ * This ensures that children are connected even if the parent appears later in the array.
+ * The input messages SHOULD BE SORTED by `created_at` ASC for optimal performance.
+ * @param messages - An array of ChatMessage objects.
+ * @param byId - An empty Map that will be populated with all nodes by their ID.
+ * @returns An array of root MessageNode(s).
+ */
+export function buildTree(
+  messages: ChatMessage[],
+  byId: Map<string, MessageNode>,
+): MessageNode[] {
+  const roots: MessageNode[] = [];
+  const orphans: MessageNode[] = [];
+
+  for (const message of messages) {
+    const node = new MessageNode(message);
+    byId.set(node.id, node);
+
+    if (node.parent_id === null) {
+      roots.push(node);
+    } else {
+      const parent = byId.get(node.parent_id);
+      if (parent) {
+        parent.addChild(node);
+      } else {
+        orphans.push(node);
+      }
+    }
+  }
+
+  for (const node of orphans) {
+    const parent = byId.get(node.parent_id ?? "");
+    if (parent) {
+      parent.addChild(node);
+    } else {
+      console.error(
+        `Could not find parent ${node.parent_id} for orphan node ${node.id}. Treating as root.`,
+      );
+      roots.push(node);
+    }
+  }
+
+  return roots;
+}
+
+/**
+ * Finds the latest path in the conversation tree by following
+ * the leaf with the most recent `createdAt` timestamp.
+ */
+export function findLatestPath(
+  tree: MessageNode[],
+  byId: Map<string, MessageNode>,
+): ChatMessage[] {
+  const leaves: MessageNode[] = [];
+
+  const stack: MessageNode[] = [...tree];
+  while (stack.length > 0) {
+    const node = stack.pop();
+    if (!node) continue;
+    const children = node.getChildren();
+    if (children.length === 0) {
+      leaves.push(node);
+    } else {
+      for (const child of children) {
+        stack.push(child);
+      }
+    }
+  }
+
+  if (leaves.length === 0) {
+    return [];
+  }
+
+  let latestLeaf = leaves[0];
+  for (const leaf of leaves) {
+    if (
+      leaf.message.createdAt &&
+      latestLeaf.message.createdAt &&
+      new Date(leaf.message.createdAt) > new Date(latestLeaf.message.createdAt)
+    ) {
+      latestLeaf = leaf;
+    }
+  }
+
+  const path: ChatMessage[] = [];
+  let currentNode: MessageNode | undefined = latestLeaf;
+  while (currentNode) {
+    path.unshift(currentNode.message);
+    currentNode = byId.get(currentNode.parent_id ?? "");
+  }
+  return path;
+}

package/src/types.ts

@@ -0,0 +1,63 @@
+import type {
+  WsClarifyQuestion,
+  WsClientCommand,
+  WsErrorMessage,
+  WsServerMessage,
+  WsStreamMessage,
+  WsVaultItem,
+} from "@futurity/chat-protocol";
+import { type MessagePart, Z_MessagePart } from "@futurity/chat-protocol";
+import { z } from "zod";
+
+export type { MessagePart };
+
+export const messageMetadataSchema = z.object({
+  parent_id: z.uuid().nullable(),
+});
+
+export type MessageMetadata = z.infer<typeof messageMetadataSchema>;
+
+/** A chat message suitable for rendering in a UI. */
+export type ChatMessage = {
+  id: string;
+  role: "user" | "assistant" | "system";
+  parts: MessagePart[];
+  createdAt?: Date;
+  metadata?: MessageMetadata;
+};
+
+export const Z_ChatMessage = z.looseObject({
+  id: z.string(),
+  role: z.enum(["user", "assistant", "system"]),
+  parts: z.array(Z_MessagePart),
+  createdAt: z.preprocess((val) => {
+    if (!val || typeof val !== "string") return undefined;
+    return new Date(val);
+  }, z.date().optional()),
+  metadata: messageMetadataSchema.optional(),
+});
+
+export type SendMessagePayload = {
+  parts: MessagePart[];
+  metadata?: MessageMetadata;
+  /** Optional vault items to attach to the message. */
+  vaultItems?: WsVaultItem[];
+};
+
+export type SendMessageFn = (payload: SendMessagePayload) => Promise<void>;
+
+export type StreamDelta = WsStreamMessage["delta"];
+
+export type ChatStatus = "ready" | "submitted" | "streaming" | "error";
+
+export type ClarifyData = WsClarifyQuestion[];
+
+/** Re-export protocol types for consumers. */
+export type {
+  WsClientCommand,
+  WsServerMessage,
+  WsErrorMessage,
+  WsStreamMessage,
+  WsClarifyQuestion,
+  WsVaultItem,
+};

package/src/useReconnectingWebSocket.ts

@@ -0,0 +1,126 @@
+import { useCallback, useEffect, useRef, useState } from "react";
+import {
+  type ConnectionState,
+  WebSocketConnection,
+} from "./WebSocketConnection";
+
+export type { ConnectionState } from "./WebSocketConnection";
+
+export type WebSocketOptions<TServerMessage> = {
+  /** WebSocket URL (can be relative, will be converted to ws/wss) */
+  url: string;
+  /** Called when a message is received */
+  onMessage: (message: TServerMessage) => void;
+  /** Called when connection state changes */
+  onConnectionChange?: (state: ConnectionState) => void;
+  /** Called when an error occurs */
+  onError?: (error: Event) => void;
+  /** Whether the WebSocket should be enabled (default: true) */
+  enabled?: boolean;
+  /** Heartbeat interval in ms (default: 5000 = 5s) */
+  heartbeatInterval?: number;
+  /** How long to wait for pong before considering connection dead (default: 10000 = 10s) */
+  heartbeatTimeout?: number;
+  /** Initial reconnection delay in ms (default: 1000) */
+  initialReconnectDelay?: number;
+  /** Maximum reconnection delay in ms (default: 30000) */
+  maxReconnectDelay?: number;
+  /** Log prefix for debugging */
+  debugPrefix?: string;
+};
+
+export type WebSocketResult<TClientCommand> = {
+  /** Current connection state */
+  connectionState: ConnectionState;
+  /** Whether the WebSocket is currently connected */
+  isConnected: boolean;
+  /** Send a message, auto-reconnecting if needed */
+  send: (message: TClientCommand) => void;
+  /** Ensure connection before performing an action */
+  ensureConnected: () => Promise<boolean>;
+  /** Force reconnection */
+  reconnect: () => void;
+};
+
+/**
+ * A React hook wrapping {@link WebSocketConnection} with lifecycle management.
+ *
+ * - Instantiates a `WebSocketConnection` in a ref
+ * - Manages connect/disconnect on mount/unmount and when `enabled` changes
+ * - Bridges connection events to React state and callback props
+ * - Exposes `send`, `ensureConnected`, `reconnect` with stable references
+ */
+export function useReconnectingWebSocket<TServerMessage, TClientCommand>({
+  url,
+  onMessage,
+  onConnectionChange,
+  onError,
+  enabled = true,
+  heartbeatInterval,
+  heartbeatTimeout,
+  initialReconnectDelay,
+  maxReconnectDelay,
+  debugPrefix,
+}: WebSocketOptions<TServerMessage>): WebSocketResult<TClientCommand> {
+  const [connectionState, setConnectionState] =
+    useState<ConnectionState>("disconnected");
+
+  const onMessageRef = useRef(onMessage);
+  const onConnectionChangeRef = useRef(onConnectionChange);
+  const onErrorRef = useRef(onError);
+
+  useEffect(() => {
+    onMessageRef.current = onMessage;
+    onConnectionChangeRef.current = onConnectionChange;
+    onErrorRef.current = onError;
+  });
+
+  const connRef = useRef<WebSocketConnection | null>(null);
+
+  // Lazy-init the connection instance once
+  if (!connRef.current) {
+    connRef.current = new WebSocketConnection({
+      url,
+      onMessage: (data) => onMessageRef.current(data as TServerMessage),
+      onConnectionChange: (state) => {
+        setConnectionState(state);
+        onConnectionChangeRef.current?.(state);
+      },
+      onError: (error) => onErrorRef.current?.(error),
+      heartbeatInterval,
+      heartbeatTimeout,
+      initialReconnectDelay,
+      maxReconnectDelay,
+      debugPrefix,
+    });
+  }
+
+  useEffect(() => {
+    if (enabled) {
+      connRef.current?.connect();
+    }
+    return () => {
+      connRef.current?.disconnect();
+    };
+  }, [enabled]);
+
+  const send = useCallback((message: TClientCommand) => {
+    connRef.current?.send(message);
+  }, []);
+
+  const ensureConnected = useCallback((): Promise<boolean> => {
+    return connRef.current?.ensureConnected() ?? Promise.resolve(false);
+  }, []);
+
+  const reconnect = useCallback(() => {
+    connRef.current?.reconnect();
+  }, []);
+
+  return {
+    connectionState,
+    isConnected: connectionState === "connected",
+    send,
+    ensureConnected,
+    reconnect,
+  };
+}