otherwise-cli 0.1.0

package/src/ui/hooks/useChatState.js

@@ -0,0 +1,675 @@
+/**
+ * Chat state management hook for Ink UI
+ * Manages chat messages, streaming content, and tool executions
+ */
+
+import { useState, useCallback, useReducer, useRef, useEffect } from "react";
+
+/**
+ * Message types
+ */
+export const MessageRole = {
+  USER: "user",
+  ASSISTANT: "assistant",
+  SYSTEM: "system",
+};
+
+/**
+ * Generation states
+ */
+export const GenerationState = {
+  IDLE: "idle",
+  THINKING: "thinking",
+  GENERATING: "generating",
+  TOOL_EXECUTING: "tool_executing",
+  STOPPED: "stopped",
+  ERROR: "error",
+};
+
+/**
+ * Tool execution states
+ */
+export const ToolState = {
+  PREPARING: "preparing",
+  RUNNING: "running",
+  STREAMING: "streaming",
+  COMPLETE: "complete",
+  ERROR: "error",
+};
+
+/**
+ * Tool reducer actions
+ */
+const toolActions = {
+  START: "tool/start",
+  UPDATE: "tool/update",
+  STREAMING: "tool/streaming",
+  COMPLETE: "tool/complete",
+  ERROR: "tool/error",
+  CLEAR: "tool/clear",
+};
+
+/**
+ * Tool state reducer
+ */
+function toolReducer(state, action) {
+  switch (action.type) {
+    case toolActions.START: {
+      // Clean up any existing streaming entry for this tool name (streaming → running transition)
+      const cleaned = { ...state };
+      for (const [id, tool] of Object.entries(cleaned)) {
+        if (
+          id !== action.id &&
+          tool.name === action.name &&
+          (tool.status === ToolState.STREAMING || tool.status === ToolState.PREPARING)
+        ) {
+          delete cleaned[id];
+        }
+      }
+      return {
+        ...cleaned,
+        [action.id]: {
+          id: action.id,
+          name: action.name,
+          args: action.args,
+          status: ToolState.RUNNING,
+          startTime: state[action.id]?.startTime || Date.now(),
+          streamingContent: state[action.id]?.streamingContent || "",
+          result: null,
+          error: null,
+        },
+      };
+    }
+
+    case toolActions.STREAMING: {
+      const existing = state[action.id];
+      if (existing) {
+        return {
+          ...state,
+          [action.id]: {
+            ...existing,
+            status: ToolState.STREAMING,
+            args: { ...existing.args, ...action.args },
+            streamingContent: action.content || "",
+          },
+        };
+      }
+      // Create entry when it doesn't exist yet (tool_streaming arrives before tool_start)
+      return {
+        ...state,
+        [action.id]: {
+          id: action.id,
+          name: action.name,
+          args: action.args || {},
+          status: ToolState.STREAMING,
+          startTime: Date.now(),
+          streamingContent: action.content || "",
+          result: null,
+          error: null,
+        },
+      };
+    }
+
+    case toolActions.COMPLETE:
+      if (!state[action.id]) return state;
+      return {
+        ...state,
+        [action.id]: {
+          ...state[action.id],
+          status: ToolState.COMPLETE,
+          result: action.result,
+          endTime: Date.now(),
+        },
+      };
+
+    case toolActions.ERROR:
+      if (!state[action.id]) return state;
+      return {
+        ...state,
+        [action.id]: {
+          ...state[action.id],
+          status: ToolState.ERROR,
+          error: action.error,
+          endTime: Date.now(),
+        },
+      };
+
+    case toolActions.CLEAR:
+      return {};
+
+    default:
+      return state;
+  }
+}
+
+/**
+ * Custom hook for chat state management
+ * @param {object} options - Configuration options
+ * @returns {object} - Chat state and methods
+ */
+export function useChatState(options = {}) {
+  const { initialChatId = null, initialModel = "claude-sonnet-4-20250514" } =
+    options;
+
+  // Chat metadata
+  const [currentChatId, setCurrentChatId] = useState(initialChatId);
+  const currentChatIdRef = useRef(initialChatId);
+  const [currentChatTitle, setCurrentChatTitle] = useState(null);
+  const [currentModel, setCurrentModel] = useState(initialModel);
+
+  // Messages (completed)
+  const [messages, setMessages] = useState([]);
+
+  // Generation state
+  const [generationState, setGenerationState] = useState(GenerationState.IDLE);
+  const [streamingContent, setStreamingContent] = useState("");
+  const [thinkingContent, setThinkingContent] = useState("");
+  const [thinkingStartTime, setThinkingStartTime] = useState(null);
+
+  // Ref to track current streaming content (avoids stale closure issues)
+  const streamingContentRef = useRef("");
+
+  // Ref to track current tools for message history
+  const toolsRef = useRef({});
+
+  // Tool executions
+  const [tools, dispatchTool] = useReducer(toolReducer, {});
+
+  // Stats from last generation
+  const [lastStats, setLastStats] = useState(null);
+
+  // Error state
+  const [error, setError] = useState(null);
+
+  // Keep refs in sync (must be after state declarations)
+  useEffect(() => {
+    streamingContentRef.current = streamingContent;
+  }, [streamingContent]);
+
+  // Throttled flush so CLI re-renders during streaming (fixes "text only on frontend, full message at end on CLI")
+  // Without this, many WebSocket 'text' chunks in the same tick batch into one React update
+  const isStreaming =
+    generationState === GenerationState.THINKING ||
+    generationState === GenerationState.GENERATING ||
+    generationState === GenerationState.TOOL_EXECUTING;
+  useEffect(() => {
+    if (!isStreaming) return;
+    const interval = setInterval(() => {
+      const current = streamingContentRef.current;
+      setStreamingContent((prev) => (prev === current ? prev : current));
+    }, 50);
+    return () => clearInterval(interval);
+  }, [isStreaming]);
+
+  useEffect(() => {
+    toolsRef.current = tools;
+  }, [tools]);
+
+  /**
+   * Start a new generation
+   */
+  const startGeneration = useCallback(() => {
+    setGenerationState(GenerationState.THINKING);
+    setStreamingContent("");
+    streamingContentRef.current = "";
+    setThinkingContent("");
+    setThinkingStartTime(Date.now());
+    setError(null);
+    dispatchTool({ type: toolActions.CLEAR });
+  }, []);
+
+  /**
+   * Handle thinking content
+   * @param {string} content - Thinking text chunk
+   * @param {number} totalChars - Total character count from server
+   */
+  const handleThinking = useCallback(
+    (content, totalChars) => {
+      setThinkingContent((prev) => prev + content);
+      if (generationState === GenerationState.IDLE) {
+        setGenerationState(GenerationState.THINKING);
+        setThinkingStartTime(Date.now());
+      }
+    },
+    [generationState],
+  );
+
+  /**
+   * Handle text content
+   * @param {string} content - Text chunk
+   */
+  const handleText = useCallback((content) => {
+    if (!content) return;
+
+    // Skip empty whitespace at start
+    if (streamingContentRef.current.length === 0 && !content.trim()) {
+      return;
+    }
+
+    setStreamingContent((prev) => {
+      const newContent = prev + content;
+      streamingContentRef.current = newContent;
+      return newContent;
+    });
+    setGenerationState((prev) =>
+      prev === GenerationState.THINKING ? GenerationState.GENERATING : prev,
+    );
+  }, []);
+
+  /**
+   * Handle tool start
+   * Transitions existing streaming entry to running, or creates a new one
+   * @param {object} data - Tool start data
+   */
+  const handleToolStart = useCallback((data) => {
+    if (data.name === "set_title") return;
+
+    setGenerationState(GenerationState.TOOL_EXECUTING);
+    // START action in the reducer automatically cleans up streaming entries
+    // with the same tool name, so no explicit removal needed here
+    dispatchTool({
+      type: toolActions.START,
+      id: data.callId,
+      name: data.name,
+      args: data.args,
+    });
+  }, []);
+
+  /**
+   * Handle tool streaming update
+   * @param {object} data - Tool streaming data (has streamingId, not callId)
+   */
+  const handleToolStreaming = useCallback((data) => {
+    if (data.name === "set_title" || data.name === "_preparing") return;
+
+    setGenerationState(GenerationState.TOOL_EXECUTING);
+    dispatchTool({
+      type: toolActions.STREAMING,
+      id: data.streamingId || data.callId || `streaming-${data.name}`,
+      name: data.name,
+      args: data.args,
+      content: data.streamingContent,
+    });
+  }, []);
+
+  /**
+   * Handle tool result
+   * @param {object} data - Tool result data
+   */
+  const handleToolResult = useCallback((data) => {
+    dispatchTool({
+      type: toolActions.COMPLETE,
+      id: data.callId,
+      result: data.result,
+    });
+  }, []);
+
+  /**
+   * Handle tool error
+   * @param {object} data - Tool error data
+   */
+  const handleToolError = useCallback((data) => {
+    dispatchTool({
+      type: toolActions.ERROR,
+      id: data.callId,
+      error: data.error,
+    });
+  }, []);
+
+  /**
+   * Handle generation done
+   * @param {object} data - Done data with stats
+   */
+  const handleDone = useCallback(
+    (data) => {
+      // Use refs to get current values (avoids stale closure)
+      const currentContent = streamingContentRef.current;
+      const currentTools = toolsRef.current;
+
+      // Convert tools to array for storage, filtering out hidden tools
+      const toolsArray = Object.values(currentTools).filter(
+        (t) => t.name !== "set_title" && t.name !== "_preparing",
+      );
+
+      // Add assistant message to history (even if empty, to show tools)
+      if (currentContent || toolsArray.length > 0) {
+        setMessages((prev) => [
+          ...prev,
+          {
+            id: `assistant-${Date.now()}`,
+            role: MessageRole.ASSISTANT,
+            content: currentContent || "",
+            model: data.message?.model || currentModel,
+            timestamp: new Date().toISOString(),
+            tools: toolsArray.length > 0 ? toolsArray : undefined,
+          },
+        ]);
+      }
+
+      setLastStats({
+        tps: data.tps,
+        numTokens: data.numTokens,
+        model: data.message?.model || currentModel,
+        finishReason: data.finishReason,
+      });
+
+      setGenerationState(GenerationState.IDLE);
+      setStreamingContent("");
+      streamingContentRef.current = "";
+      setThinkingContent("");
+      setThinkingStartTime(null);
+      dispatchTool({ type: toolActions.CLEAR });
+      toolsRef.current = {};
+    },
+    [currentModel],
+  );
+
+  /**
+   * Handle generation error
+   * @param {string} message - Error message
+   */
+  const handleError = useCallback((message) => {
+    setGenerationState(GenerationState.ERROR);
+    setError(message);
+  }, []);
+
+  /**
+   * Handle generation stopped
+   * Preserves any partial content that was streamed
+   */
+  const handleStopped = useCallback(() => {
+    // Get current content before clearing
+    const currentContent = streamingContentRef.current;
+    const currentTools = toolsRef.current;
+
+    // Convert tools to array for storage
+    const toolsArray = Object.values(currentTools).filter(
+      (t) => t.name !== "set_title" && t.name !== "_preparing",
+    );
+
+    // If we have partial content, save it with a [stopped] indicator
+    if (currentContent || toolsArray.length > 0) {
+      const stoppedContent = currentContent
+        ? currentContent + "\n\n*[Generation stopped]*"
+        : "*[Generation stopped]*";
+
+      setMessages((prev) => [
+        ...prev,
+        {
+          id: `assistant-${Date.now()}`,
+          role: MessageRole.ASSISTANT,
+          content: stoppedContent,
+          timestamp: new Date().toISOString(),
+          tools: toolsArray.length > 0 ? toolsArray : undefined,
+          stopped: true,
+        },
+      ]);
+    }
+
+    setGenerationState(GenerationState.STOPPED);
+    setStreamingContent("");
+    streamingContentRef.current = "";
+    setThinkingContent("");
+    setThinkingStartTime(null);
+    dispatchTool({ type: toolActions.CLEAR });
+    toolsRef.current = {};
+  }, []);
+
+  /**
+   * Handle chat title update
+   * @param {string} title - New title
+   */
+  const handleTitle = useCallback((title) => {
+    setCurrentChatTitle(title);
+  }, []);
+
+  /**
+   * Handle chat created event
+   * @param {number} chatId - New chat ID
+   */
+  const handleChatCreated = useCallback((chatId) => {
+    currentChatIdRef.current = chatId;
+    setCurrentChatId(chatId);
+  }, []);
+
+  /**
+   * Handle user message from another client (e.g., frontend)
+   * @param {object} data - Message data with chatId and content
+   */
+  const handleUserMessage = useCallback(
+    (data) => {
+      const { chatId, content } = data;
+      const activeChatId = currentChatIdRef.current;
+
+      // Accept message if it's for the current chat, or if no chat is active yet
+      if (!activeChatId || chatId === activeChatId) {
+        if (!activeChatId && chatId) {
+          currentChatIdRef.current = chatId;
+          setCurrentChatId(chatId);
+        }
+
+        setMessages((prev) => [
+          ...prev,
+          {
+            id: `user-${Date.now()}`,
+            role: MessageRole.USER,
+            content,
+            source: "web",
+            timestamp: new Date().toISOString(),
+          },
+        ]);
+
+        setThinkingContent("");
+        setStreamingContent("");
+        streamingContentRef.current = "";
+        dispatchTool({ type: toolActions.CLEAR });
+        toolsRef.current = {};
+        setGenerationState(GenerationState.THINKING);
+        setThinkingStartTime(Date.now());
+      }
+    },
+    [],
+  );
+
+  /**
+   * Handle chat selection from another client
+   * @param {number} chatId - Selected chat ID
+   */
+  const handleChatSelected = useCallback(
+    (chatId) => {
+      if (chatId !== currentChatIdRef.current) {
+        currentChatIdRef.current = chatId;
+        setCurrentChatId(chatId);
+        setMessages([]);
+        setCurrentChatTitle(null);
+        setGenerationState(GenerationState.IDLE);
+      }
+    },
+    [],
+  );
+
+  /**
+   * Handle model change from another client
+   * @param {string} model - New model ID
+   */
+  const handleModelChanged = useCallback((model) => {
+    setCurrentModel(model);
+  }, []);
+
+  /**
+   * Add a user message
+   * @param {string} content - Message content
+   */
+  const addUserMessage = useCallback((content) => {
+    setMessages((prev) => [
+      ...prev,
+      {
+        id: `user-${Date.now()}`,
+        role: MessageRole.USER,
+        content,
+        timestamp: new Date().toISOString(),
+      },
+    ]);
+  }, []);
+
+  /**
+   * Start a new chat
+   */
+  const newChat = useCallback(() => {
+    currentChatIdRef.current = null;
+    setCurrentChatId(null);
+    setCurrentChatTitle(null);
+    setMessages([]);
+    setStreamingContent("");
+    setThinkingContent("");
+    setThinkingStartTime(null);
+    setGenerationState(GenerationState.IDLE);
+    setError(null);
+    dispatchTool({ type: toolActions.CLEAR });
+  }, []);
+
+  /**
+   * Load a chat with messages
+   * @param {object} chat - Chat object with messages
+   */
+  const loadChat = useCallback((chat) => {
+    currentChatIdRef.current = chat.id;
+    setCurrentChatId(chat.id);
+    setCurrentChatTitle(chat.title);
+    setMessages(chat.messages || []);
+    setGenerationState(GenerationState.IDLE);
+    setStreamingContent("");
+    setThinkingContent("");
+    setError(null);
+    dispatchTool({ type: toolActions.CLEAR });
+  }, []);
+
+  /**
+   * Process a WebSocket message
+   * @param {object} data - Message data
+   */
+  const processMessage = useCallback(
+    (data) => {
+      switch (data.type) {
+        // Chat lifecycle events
+        case "chat_created":
+          handleChatCreated(data.chatId);
+          break;
+        case "chat_selected":
+          handleChatSelected(data.chatId);
+          break;
+
+        // User message from another client (frontend)
+        case "user_message":
+          handleUserMessage(data);
+          break;
+
+        // Model change from another client
+        case "model_changed":
+          handleModelChanged(data.model);
+          break;
+
+        // Generation streaming events
+        case "thinking":
+          handleThinking(data.content, data.totalChars);
+          break;
+        case "text":
+          handleText(data.content);
+          break;
+
+        // Tool events
+        case "tool_start":
+          handleToolStart(data);
+          break;
+        case "tool_streaming":
+          handleToolStreaming(data);
+          break;
+        case "tool_result":
+          handleToolResult(data);
+          break;
+        case "tool_error":
+          handleToolError(data);
+          break;
+
+        // Chat metadata
+        case "title":
+          handleTitle(data.content);
+          break;
+
+        // Generation completion
+        case "done":
+          handleDone(data);
+          break;
+        case "error":
+          handleError(data.message);
+          break;
+        case "stopped":
+          handleStopped();
+          break;
+      }
+    },
+    [
+      handleChatCreated,
+      handleChatSelected,
+      handleUserMessage,
+      handleModelChanged,
+      handleThinking,
+      handleText,
+      handleToolStart,
+      handleToolStreaming,
+      handleToolResult,
+      handleToolError,
+      handleTitle,
+      handleDone,
+      handleError,
+      handleStopped,
+    ],
+  );
+
+  return {
+    // Chat metadata
+    currentChatId,
+    currentChatTitle,
+    currentModel,
+    setCurrentModel,
+
+    // Messages
+    messages,
+    addUserMessage,
+
+    // Generation state
+    generationState,
+    isGenerating:
+      generationState !== GenerationState.IDLE &&
+      generationState !== GenerationState.STOPPED &&
+      generationState !== GenerationState.ERROR,
+    isThinking: generationState === GenerationState.THINKING,
+    streamingContent,
+    thinkingContent,
+    thinkingStartTime,
+
+    // Tools
+    tools,
+    activeTools: Object.values(tools).filter(
+      (t) => t.status === ToolState.RUNNING || t.status === ToolState.STREAMING,
+    ),
+
+    // Stats
+    lastStats,
+
+    // Error
+    error,
+
+    // Actions
+    startGeneration,
+    newChat,
+    loadChat,
+    processMessage,
+
+    // For external state updates
+    setCurrentChatId,
+    setCurrentChatTitle,
+  };
+}
+
+export default useChatState;