zclaw_openfang/desktop/src/hooks/useOptimisticMessages.ts

import { useCallback, useRef } from 'react';
import { useChatStore, type Message } from '../store/chatStore';
import { createLogger } from '../lib/logger';

const log = createLogger('OptimisticMessages');

/**
 * Represents a file attached to an optimistic message,
 * tracking its upload lifecycle. Extends MessageFile with a status field.
 */
interface OptimisticFile {
  name: string;
  size: number;
  status: 'uploading' | 'uploaded' | 'error';
  url?: string;
}

/**
 * 3-phase optimistic message merging hook (inspired by DeerFlow useThreadStream).
 *
 * Phase 1: Instant local echo -- creates a synthetic user message with `optimistic: true`
 * Phase 2: Server confirmation -- removes optimistic message when real message arrives
 * Phase 3: File status transition -- updates file status from uploading -> uploaded | error
 *
 * This hook provides standalone utilities for components that need fine-grained
 * control over optimistic rendering outside the main chat flow.
 */
export function useOptimisticMessages() {
  const optimisticIdCounter = useRef(0);

  const generateOptimisticId = useCallback(() => {
    optimisticIdCounter.current += 1;
    return `opt-user-${Date.now()}-${optimisticIdCounter.current}`;
  }, []);

  /**
   * Phase 1: Create and insert an optimistic user message into the store.
   * Returns the optimistic ID for later correlation.
   */
  const addOptimistic = useCallback((content: string, files?: File[]) => {
    const id = generateOptimisticId();
    const optimisticFiles: OptimisticFile[] | undefined = files?.map(f => ({
      name: f.name,
      size: f.size,
      status: 'uploading' as const,
    }));

    const optimisticMessage: Message = {
      id,
      role: 'user',
      content,
      timestamp: new Date(),
      optimistic: true,
      // Cast through unknown because OptimisticFile extends MessageFile with status
      files: optimisticFiles as Message['files'],
    };

    log.debug('Adding optimistic message', { id, content: content.slice(0, 50) });

    useChatStore.setState(state => ({
      messages: [...state.messages, optimisticMessage],
    }));

    return id;
  }, [generateOptimisticId]);

  /**
   * Phase 2: Remove an optimistic message when the server confirms
   * by sending back the real message.
   */
  const clearOnConfirm = useCallback((optimisticId: string) => {
    log.debug('Clearing optimistic message on confirm', { optimisticId });

    useChatStore.setState(state => ({
      messages: state.messages.filter(m => m.id !== optimisticId),
    }));
  }, []);

  /**
   * Phase 3: Transition file attachment status for an optimistic message.
   */
  const updateFileStatus = useCallback((optimisticId: string, status: 'uploaded' | 'error') => {
    log.debug('Updating file status', { optimisticId, status });

    useChatStore.setState(state => ({
      messages: state.messages.map(m => {
        if (m.id === optimisticId && m.files) {
          return {
            ...m,
            files: m.files.map(f => ({
              ...f,
              status,
            })),
          };
        }
        return m;
      }),
    }));
  }, []);

  return { addOptimistic, clearOnConfirm, updateFileStatus };
}