@futurity/chat-react 0.0.2 → 0.1.0

package/src/useStreamChat.ts

@@ -4,13 +4,14 @@ import type {
 } from "@futurity/chat-protocol";
 import { useCallback, useEffect, useRef, useState } from "react";
 import { type ClientCommand, parseServerMessage } from "./chat-protocol";
+import {
+  createPreprocessorState,
+  incrementalPreprocess,
+  type PreprocessorState,
+  type ProcessedPart,
+} from "./stream-accumulator";
 import { buildTree, findLatestPath, type MessageNode } from "./tree-builder";
-import type {
-  ChatMessage,
-  ChatStatus,
-  SendMessagePayload,
-  StreamDelta,
-} from "./types";
+import type { ChatMessage, ChatStatus, SendMessagePayload } from "./types";
 import { Z_ChatMessage } from "./types";
 import { useReconnectingWebSocket } from "./useReconnectingWebSocket";
 
@@ -22,12 +23,49 @@ type TransformedHistory = {
   activeMessageId?: string;
 };
 
-function transformChatHistory(rawMessages: unknown[]): TransformedHistory {
+/** Produce simple 1:1 processedParts for a non-assistant message. */
+function simpleProcessedParts(parts: MessagePart[]): ProcessedPart[] {
+  return parts.map((part, i) => ({
+    type: "regular" as const,
+    part,
+    originalIndex: i,
+  }));
+}
+
+/** Run the preprocessor on a message and return it with processedParts populated. */
+function preprocessMessage(
+  msg: Omit<ChatMessage, "processedParts"> & {
+    processedParts?: ProcessedPart[];
+  },
+  statesMap: Map<string, PreprocessorState>,
+): ChatMessage {
+  if (msg.role !== "assistant") {
+    return { ...msg, processedParts: simpleProcessedParts(msg.parts) };
+  }
+  let state = statesMap.get(msg.id);
+  if (!state || state.scannedLength > (msg.parts?.length ?? 0)) {
+    // Reset if message changed (parts shrunk) or no state yet
+    state = createPreprocessorState(msg.id);
+    statesMap.set(msg.id, state);
+  }
+  return {
+    ...msg,
+    processedParts: incrementalPreprocess(state, msg.parts ?? []),
+  };
+}
+
+function transformChatHistory(
+  rawMessages: unknown[],
+  statesMap: Map<string, PreprocessorState>,
+): TransformedHistory {
   const parsedMessages = Z_ChatMessage.array().safeParse(rawMessages);
   if (!parsedMessages.success) {
     throw new Error("Invalid chat history");
   }
-  const messages = parsedMessages.data;
+  // Parsed messages don't have processedParts yet — add them
+  const messages = parsedMessages.data.map((msg) =>
+    preprocessMessage(msg, statesMap),
+  );
   const byId = new Map<string, MessageNode>();
   const tree = buildTree(messages, byId);
   const initialPath = findLatestPath(tree, byId);
@@ -47,10 +85,6 @@ export type UseStreamChatOptions = {
   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. */
@@ -86,8 +120,6 @@ export function useStreamChat({
   chatId,
   wsUrl,
   onStart,
-  onDelta,
-  onResume,
   onFinish,
   onError,
   onHistory,
@@ -107,9 +139,10 @@ export function useStreamChat({
     WsClarifyRequestMessage["data"] | null
   >(null);
 
+  // Per-message preprocessor states, keyed by message ID
+  const preprocessorStatesRef = useRef(new Map<string, PreprocessorState>());
+
   const onStartRef = useRef(onStart);
-  const onDeltaRef = useRef(onDelta);
-  const onResumeRef = useRef(onResume);
   const onFinishRef = useRef(onFinish);
   const onErrorRef = useRef(onError);
   const onHistoryRef = useRef(onHistory);
@@ -117,8 +150,6 @@ export function useStreamChat({
 
   useEffect(() => {
     onStartRef.current = onStart;
-    onDeltaRef.current = onDelta;
-    onResumeRef.current = onResume;
     onFinishRef.current = onFinish;
     onErrorRef.current = onError;
     onHistoryRef.current = onHistory;
@@ -155,16 +186,90 @@ export function useStreamChat({
 
         case "stream": {
           const messageId = message.messageId;
+          const delta = message.delta;
+          const consolidated = message.consolidated;
 
           setStatus("streaming");
-          onDeltaRef.current?.(messageId, message.delta, message.consolidated);
+          setMessages((prev) => {
+            const states = preprocessorStatesRef.current;
+            const exists = prev.some((m) => m.id === messageId);
+            let updated: ChatMessage[];
+            if (!exists) {
+              // Unknown message ID (e.g. new assistant message from inject) — create it
+              const newMsg: ChatMessage = {
+                id: messageId,
+                role: "assistant",
+                parts: [delta],
+                processedParts: [],
+              };
+              updated = [...prev, newMsg];
+            } else {
+              updated = prev.map((m) => {
+                if (m.id !== messageId) return m;
+                const parts = m.parts ?? [];
+                const nextParts = consolidated
+                  ? [...parts.slice(0, -1), delta]
+                  : [...parts, delta];
+                return { ...m, parts: nextParts };
+              });
+            }
+            // Reprocess the affected message
+            return updated.map((m) =>
+              m.id === messageId ? preprocessMessage(m, states) : m,
+            );
+          });
           break;
         }
 
         case "stream_resume": {
           const messageId = message.messageId;
           setStatus("streaming");
-          onResumeRef.current?.(messageId, message.parts);
+          setMessages((prev) => {
+            const states = preprocessorStatesRef.current;
+            // Reset preprocessor state for this message since parts are being replaced
+            states.delete(messageId);
+            return prev.map((m) => {
+              if (m.id !== messageId) return m;
+              const updated = { ...m, parts: message.parts };
+              return preprocessMessage(updated, states);
+            });
+          });
+          break;
+        }
+
+        case "inject_ack": {
+          const { new_assistant_id } = message.data;
+          setMessages((prev) => {
+            const exists = prev.some((m) => m.id === new_assistant_id);
+            if (exists) return prev;
+            return [
+              ...prev,
+              {
+                id: new_assistant_id,
+                role: "assistant" as const,
+                parts: [],
+                processedParts: [],
+              },
+            ];
+          });
+          break;
+        }
+
+        case "inject_split": {
+          const { new_message_id } = message.data;
+          setMessages((prev) => {
+            const exists = prev.some((m) => m.id === new_message_id);
+            if (exists) return prev;
+            return [
+              ...prev,
+              {
+                id: new_message_id,
+                role: "assistant" as const,
+                parts: [],
+                processedParts: [],
+              },
+            ];
+          });
           break;
         }
 
@@ -182,7 +287,10 @@ export function useStreamChat({
         }
 
         case "chat_history": {
-          const result = transformChatHistory(message.data.messages);
+          const result = transformChatHistory(
+            message.data.messages,
+            preprocessorStatesRef.current,
+          );
           result.activeMessageId = message.data.activeMessageId;
           onHistoryRef.current?.(result);
           break;
@@ -252,6 +360,7 @@ export function useStreamChat({
         id: crypto.randomUUID(),
         role: "user",
         parts: payload.parts,
+        processedParts: simpleProcessedParts(payload.parts),
         metadata: {
           parent_id: parent_id,
         },
@@ -292,14 +401,16 @@ export function useStreamChat({
       if (!connected) return;
 
       const message_id = crypto.randomUUID();
+      const parts = [{ type: "text" as const, text }];
 
       // Optimistically add the user message to state
       setMessages((prev) => [
         ...prev,
         {
           id: message_id,
-          role: "user",
-          parts: [{ type: "text", text }],
+          role: "user" as const,
+          parts,
+          processedParts: simpleProcessedParts(parts),
         },
       ]);
 
@@ -337,6 +448,7 @@ export function useStreamChat({
     setJob("");
     setPendingClarify(null);
     hasRequestedChat.current = false;
+    preprocessorStatesRef.current.clear();
   }, []);
 
   return {