learning_ai_common_plat/packages/ai-ui/src/useChat.ts

import { useCallback, useMemo, useRef, useState } from 'react';
import type { Message, UseChatHelpers, UseChatOptions } from './types.js';
import { streamText } from './useStreamingText.js';

let __id = 0;
const defaultGenerateId = (): string => `m-${Date.now().toString(36)}-${++__id}`;

/**
 * `useChat` — primary hook for chat UIs.
 *
 * Adopts the Vercel AI SDK `useChat` return shape (per
 * `learning_ai_common_plat/docs/ROADMAP_2026_DECISIONS.md` §10) so any
 * React engineer familiar with the AI SDK can drop straight in. The
 * transport is **pluggable** — products supply `endpoint`,
 * `streamProtocol`, `fetcher`, or `headers` to talk to OpenAI directly,
 * Anthropic, an LysnrAI service, or any backend that streams text.
 *
 * @example
 * ```tsx
 * const { messages, input, handleInputChange, handleSubmit, isLoading } =
 *   useChat({ endpoint: '/api/chat' });
 *
 * return (
 *   <>
 *     {messages.map(m => <MessageBubble key={m.id} message={m} />)}
 *     <PromptComposer value={input} onChange={handleInputChange}
 *       onSubmit={handleSubmit} disabled={isLoading} />
 *   </>
 * );
 * ```
 */
export function useChat(options: UseChatOptions = {}): UseChatHelpers {
  const {
    initialMessages = [],
    initialInput = '',
    endpoint = '/api/chat',
    streamProtocol = 'text',
    fetcher,
    headers,
    onFinish,
    onError,
    generateId = defaultGenerateId,
  } = options;

  const [messages, setMessages] = useState<Message[]>(initialMessages);
  const [input, setInput] = useState(initialInput);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<Error | undefined>(undefined);
  const abortRef = useRef<AbortController | null>(null);

  const handleInputChange = useCallback(
    (e: { target: { value: string } }) => setInput(e.target.value),
    [],
  );

  const stop = useCallback(() => {
    abortRef.current?.abort();
    abortRef.current = null;
    // Drop the trailing isStreaming flag so the bubble settles into
    // its final state even though we stopped mid-stream.
    setMessages(prev =>
      prev.map(m => (m.isStreaming ? { ...m, isStreaming: false } : m)),
    );
    setIsLoading(false);
  }, []);

  /**
   * Core send loop — appends `userMessage`, POSTs the conversation,
   * streams the response into a new assistant message, and updates
   * loading/error state along the way.
   */
  const sendMessages = useCallback(
    async (next: Message[]) => {
      // Cancel any in-flight request before starting a new one.
      abortRef.current?.abort();
      const controller = new AbortController();
      abortRef.current = controller;

      const assistantId = generateId();
      setMessages([
        ...next,
        {
          id: assistantId,
          role: 'assistant',
          content: '',
          createdAt: new Date(),
          isStreaming: true,
        },
      ]);
      setIsLoading(true);
      setError(undefined);

      const doFetch = fetcher ?? globalThis.fetch.bind(globalThis);
      let finalContent = '';

      try {
        const response = await doFetch(endpoint, {
          method: 'POST',
          signal: controller.signal,
          headers: { 'Content-Type': 'application/json', ...headers },
          body: JSON.stringify({ messages: next }),
        });

        if (!response.ok) {
          throw new Error(`Chat endpoint returned ${response.status}`);
        }

        for await (const delta of streamText(
          response,
          streamProtocol,
          controller.signal,
        )) {
          finalContent += delta;
          // Functional update so concurrent setMessages calls don't lose work.
          setMessages(prev =>
            prev.map(m =>
              m.id === assistantId ? { ...m, content: m.content + delta } : m,
            ),
          );
        }

        // Settle isStreaming flag.
        setMessages(prev =>
          prev.map(m =>
            m.id === assistantId ? { ...m, isStreaming: false } : m,
          ),
        );

        if (onFinish && !controller.signal.aborted) {
          onFinish({
            id: assistantId,
            role: 'assistant',
            content: finalContent,
            isStreaming: false,
          });
        }
      } catch (err) {
        if ((err as Error).name === 'AbortError') {
          // User-initiated stop — already handled in `stop()`.
          return;
        }
        const e = err instanceof Error ? err : new Error(String(err));
        setError(e);
        // Mark the partial assistant message as no longer streaming so
        // the bubble's "thinking" indicator clears.
        setMessages(prev =>
          prev.map(m =>
            m.id === assistantId ? { ...m, isStreaming: false } : m,
          ),
        );
        onError?.(e);
      } finally {
        if (abortRef.current === controller) abortRef.current = null;
        setIsLoading(false);
      }
    },
    [endpoint, fetcher, generateId, headers, onError, onFinish, streamProtocol],
  );

  const append = useCallback<UseChatHelpers['append']>(
    async incoming => {
      const message: Message = {
        id: incoming.id ?? generateId(),
        role: incoming.role,
        content: incoming.content,
        createdAt: new Date(),
      };
      const next = [...messages, message];
      setMessages(next);
      await sendMessages(next);
    },
    [generateId, messages, sendMessages],
  );

  const handleSubmit = useCallback<UseChatHelpers['handleSubmit']>(
    async e => {
      e?.preventDefault?.();
      const trimmed = input.trim();
      if (!trimmed || isLoading) return;
      setInput('');
      await append({ role: 'user', content: trimmed });
    },
    [append, input, isLoading],
  );

  const reload = useCallback<UseChatHelpers['reload']>(async () => {
    // Find the last user message; drop everything after it.
    const lastUserIndex = [...messages]
      .reverse()
      .findIndex(m => m.role === 'user');
    if (lastUserIndex === -1) return;
    const cutoff = messages.length - lastUserIndex; // exclusive
    const truncated = messages.slice(0, cutoff);
    setMessages(truncated);
    await sendMessages(truncated);
  }, [messages, sendMessages]);

  return useMemo(
    () => ({
      messages,
      input,
      setInput,
      handleInputChange,
      handleSubmit,
      append,
      isLoading,
      error,
      stop,
      reload,
      setMessages,
    }),
    [
      messages,
      input,
      handleInputChange,
      handleSubmit,
      append,
      isLoading,
      error,
      stop,
      reload,
    ],
  );
}