@futurity/chat-react 0.0.1

package/src/useStreamChat.ts

@@ -0,0 +1,354 @@
+import type {
+  MessagePart,
+  WsClarifyRequestMessage,
+} from "@futurity/chat-protocol";
+import { useCallback, useEffect, useRef, useState } from "react";
+import { type ClientCommand, parseServerMessage } from "./chat-protocol";
+import { buildTree, findLatestPath, type MessageNode } from "./tree-builder";
+import type {
+  ChatMessage,
+  ChatStatus,
+  SendMessagePayload,
+  StreamDelta,
+} from "./types";
+import { Z_ChatMessage } from "./types";
+import { useReconnectingWebSocket } from "./useReconnectingWebSocket";
+
+type TransformedHistory = {
+  messages: ChatMessage[];
+  tree: MessageNode[];
+  byId: Map<string, MessageNode>;
+  initialPath: ChatMessage[];
+  activeMessageId?: string;
+};
+
+function transformChatHistory(rawMessages: unknown[]): TransformedHistory {
+  const parsedMessages = Z_ChatMessage.array().safeParse(rawMessages);
+  if (!parsedMessages.success) {
+    throw new Error("Invalid chat history");
+  }
+  const messages = parsedMessages.data;
+  const byId = new Map<string, MessageNode>();
+  const tree = buildTree(messages, byId);
+  const initialPath = findLatestPath(tree, byId);
+
+  return {
+    messages,
+    tree,
+    byId,
+    initialPath,
+  };
+}
+
+export type UseStreamChatOptions = {
+  /** The chat ID to connect to. */
+  chatId: string;
+  /** Full WebSocket URL for the chat endpoint (e.g. `wss://api.example.com/api/v2/chat/<id>`). */
+  wsUrl: string;
+  /** Called when a new assistant message starts streaming. */
+  onStart?: (id: string) => void;
+  /** Called on each stream delta. */
+  onDelta?: (id: string, delta: StreamDelta, consolidated: boolean) => void;
+  /** Called when a stream resumes with accumulated parts. */
+  onResume?: (id: string, parts: MessagePart[]) => void;
+  /** Called when streaming finishes. */
+  onFinish?: () => void;
+  /** Called on a protocol error. */
+  onError?: (error: { error?: string; message?: string }) => void;
+  /** Called when chat history is received from the server. */
+  onHistory?: (history: TransformedHistory) => void;
+};
+
+export type UseStreamChatReturn = {
+  /** Current list of messages in the active branch. */
+  messages: ChatMessage[];
+  /** Replace the messages list (e.g. for branch switching). */
+  setMessages: React.Dispatch<React.SetStateAction<ChatMessage[]>>;
+  /** Send a new user message. */
+  sendMessage: (payload: SendMessagePayload) => Promise<void>;
+  /** Inject a user message into an active stream. */
+  injectMessage: (text: string) => Promise<void>;
+  /** Current chat status. */
+  status: ChatStatus;
+  /** Stop the current generation. */
+  stop: () => Promise<void>;
+  /** Reset the chat state. */
+  reset: () => void;
+  /** Whether the WebSocket is connected. */
+  isConnected: boolean;
+  /** Pending clarification request, if any. */
+  pendingClarify: WsClarifyRequestMessage["data"] | null;
+  /** Submit answers to a clarification request. */
+  sendClarifyResponse: (answers: Record<string, string>) => void;
+};
+
+export function useStreamChat({
+  chatId,
+  wsUrl,
+  onStart,
+  onDelta,
+  onResume,
+  onFinish,
+  onError,
+  onHistory,
+}: UseStreamChatOptions): UseStreamChatReturn {
+  if (!chatId) {
+    throw new Error("useStreamChat: chatId is required");
+  }
+  if (!wsUrl) {
+    throw new Error("useStreamChat: wsUrl is required");
+  }
+
+  const [messages, setMessages] = useState<ChatMessage[]>([]);
+  const [status, setStatus] = useState<ChatStatus>("ready");
+  const [job, setJob] = useState("");
+  const hasRequestedChat = useRef(false);
+  const [pendingClarify, setPendingClarify] = useState<
+    WsClarifyRequestMessage["data"] | null
+  >(null);
+
+  const onStartRef = useRef(onStart);
+  const onDeltaRef = useRef(onDelta);
+  const onResumeRef = useRef(onResume);
+  const onFinishRef = useRef(onFinish);
+  const onErrorRef = useRef(onError);
+  const onHistoryRef = useRef(onHistory);
+  const sendRef = useRef<(cmd: ClientCommand) => void>(() => {});
+
+  useEffect(() => {
+    onStartRef.current = onStart;
+    onDeltaRef.current = onDelta;
+    onResumeRef.current = onResume;
+    onFinishRef.current = onFinish;
+    onErrorRef.current = onError;
+    onHistoryRef.current = onHistory;
+  });
+
+  const handleMessage = useCallback(
+    (rawData: unknown) => {
+      const message = parseServerMessage(rawData);
+
+      if (!message) {
+        console.error("Invalid message received from server");
+        return;
+      }
+
+      switch (message.type) {
+        case "ready": {
+          if (!hasRequestedChat.current) {
+            hasRequestedChat.current = true;
+            sendRef.current({
+              type: "get_chat",
+              data: { chat_id: chatId },
+            });
+          }
+          break;
+        }
+
+        case "run": {
+          const { job_id, message_id } = message.data;
+          setStatus("submitted");
+          setJob(job_id);
+          onStartRef.current?.(message_id);
+          break;
+        }
+
+        case "stream": {
+          const messageId = message.messageId;
+
+          setStatus("streaming");
+          onDeltaRef.current?.(messageId, message.delta, message.consolidated);
+          break;
+        }
+
+        case "stream_resume": {
+          const messageId = message.messageId;
+          setStatus("streaming");
+          onResumeRef.current?.(messageId, message.parts);
+          break;
+        }
+
+        case "done": {
+          setStatus("ready");
+          setPendingClarify(null);
+          onFinishRef.current?.();
+          break;
+        }
+
+        case "error": {
+          setStatus("error");
+          onErrorRef.current?.(message);
+          break;
+        }
+
+        case "chat_history": {
+          const result = transformChatHistory(message.data.messages);
+          result.activeMessageId = message.data.activeMessageId;
+          onHistoryRef.current?.(result);
+          break;
+        }
+
+        case "clarify_request": {
+          setPendingClarify(message.data);
+          break;
+        }
+
+        case "cancel": {
+          setJob("");
+          setPendingClarify(null);
+          setTimeout(() => {
+            setStatus("ready");
+          }, 1000);
+          break;
+        }
+      }
+    },
+    [chatId],
+  );
+
+  const handleConnectionChange = useCallback((state: string) => {
+    if (state === "connected") {
+      setStatus("ready");
+      hasRequestedChat.current = false;
+    }
+  }, []);
+
+  const handleError = useCallback(() => {
+    setStatus("error");
+  }, []);
+
+  const { send, ensureConnected, isConnected } = useReconnectingWebSocket<
+    unknown,
+    ClientCommand
+  >({
+    url: wsUrl,
+    onMessage: handleMessage,
+    onConnectionChange: handleConnectionChange,
+    onError: handleError,
+    enabled: !!chatId,
+    debugPrefix: "[ChatWS]",
+  });
+
+  useEffect(() => {
+    sendRef.current = send;
+  }, [send]);
+
+  const sendMessage = useCallback(
+    async (payload: SendMessagePayload): Promise<void> => {
+      if (status === "streaming" || status === "submitted") return;
+
+      const connected = await ensureConnected();
+      if (!connected) {
+        console.error("[ChatWS] Failed to connect for sendMessage");
+        return;
+      }
+
+      const parent_id =
+        payload.metadata !== undefined
+          ? (payload.metadata.parent_id ?? null)
+          : (messages.at(-1)?.id ?? null);
+
+      const message: ChatMessage = {
+        id: crypto.randomUUID(),
+        role: "user",
+        parts: payload.parts,
+        metadata: {
+          parent_id: parent_id,
+        },
+      };
+
+      const timezone = Intl.DateTimeFormat().resolvedOptions().timeZone;
+      setMessages((prev) => [...prev, message]);
+      send({
+        type: "run",
+        data: {
+          id: chatId,
+          message: {
+            id: message.id,
+            role: message.role,
+            parts: message.parts,
+            metadata: message.metadata,
+          },
+          vaultItems: payload.vaultItems ?? [],
+          timezone,
+        },
+      });
+    },
+    [chatId, ensureConnected, messages, send, status],
+  );
+
+  const stop = useCallback(async () => {
+    if (!job) return;
+    const connected = await ensureConnected();
+    if (!connected) return;
+    send({ type: "cancel", data: { job_id: job } });
+  }, [job, ensureConnected, send]);
+
+  const injectMessage = useCallback(
+    async (text: string): Promise<void> => {
+      if (!job) return;
+
+      const connected = await ensureConnected();
+      if (!connected) return;
+
+      const message_id = crypto.randomUUID();
+
+      // Optimistically add the user message to state
+      setMessages((prev) => [
+        ...prev,
+        {
+          id: message_id,
+          role: "user",
+          parts: [{ type: "text", text }],
+        },
+      ]);
+
+      send({
+        type: "inject_message",
+        data: {
+          job_id: job,
+          text,
+          message_id,
+        },
+      });
+    },
+    [job, ensureConnected, send],
+  );
+
+  const sendClarifyResponse = useCallback(
+    (answers: Record<string, string>) => {
+      if (!pendingClarify) return;
+      send({
+        type: "clarify_response",
+        data: {
+          job_id: pendingClarify.job_id,
+          requestId: pendingClarify.requestId,
+          answers,
+        },
+      });
+      setPendingClarify(null);
+    },
+    [pendingClarify, send],
+  );
+
+  const reset = useCallback(() => {
+    setMessages([]);
+    setStatus("ready");
+    setJob("");
+    setPendingClarify(null);
+    hasRequestedChat.current = false;
+  }, []);
+
+  return {
+    messages,
+    setMessages,
+    sendMessage,
+    injectMessage,
+    status,
+    stop,
+    reset,
+    isConnected,
+    pendingClarify,
+    sendClarifyResponse,
+  };
+}