tkmind_go/ui/desktop/src/hooks/useChatStream.ts

import { useState, useCallback, useRef } from 'react';
import { ChatState } from '../types/chatState';
import { Message } from '../api';
import { getApiUrl } from '../config';

const TextDecoder = globalThis.TextDecoder;

interface UseChatStreamProps {
  sessionId: string;
  messages: Message[];
  setMessages: (messages: Message[]) => void;
  onStreamFinish?: () => void;
}

function pushMessage(currentMessages: Message[], incomingMsg: Message): Message[] {
  const lastMsg = currentMessages[currentMessages.length - 1];

  if (lastMsg?.id && lastMsg.id === incomingMsg.id) {
    const lastContent = lastMsg.content[lastMsg.content.length - 1];
    const newContent = incomingMsg.content[incomingMsg.content.length - 1];

    if (
      lastContent?.type === 'text' &&
      newContent?.type === 'text' &&
      incomingMsg.content.length === 1
    ) {
      lastContent.text += newContent.text;
    } else {
      lastMsg.content.push(...incomingMsg.content);
    }
    return [...currentMessages];
  } else {
    return [...currentMessages, incomingMsg];
  }
}

export function useChatStream({
  sessionId,
  messages,
  setMessages,
  onStreamFinish,
}: UseChatStreamProps) {
  const [chatState, setChatState] = useState<ChatState>(ChatState.Idle);
  const abortControllerRef = useRef<AbortController | null>(null);

  const handleSubmit = useCallback(
    async (userMessage: string) => {
      const newMessage: Message = {
        role: 'user',
        content: [{ type: 'text', text: userMessage }],
        created: Date.now(),
      };

      let currentMessages = [...messages, newMessage];
      setMessages(currentMessages);
      setChatState(ChatState.Streaming);

      abortControllerRef.current = new AbortController();

      try {
        // TODO(Douwe): this side steps our API. heyapi does support streaming though which should make
        // this all nice & typed
        const response = await fetch(getApiUrl('/reply'), {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
            'X-Secret-Key': await window.electron.getSecretKey(),
          },
          body: JSON.stringify({
            session_id: sessionId,
            messages: currentMessages,
          }),
          signal: abortControllerRef.current.signal,
        });

        if (!response.ok) throw new Error(`HTTP ${response.status}`);
        if (!response.body) throw new Error('No response body');

        const reader = response.body.getReader();
        const decoder = new TextDecoder();

        while (true) {
          const { done, value } = await reader.read();
          if (done) break;

          const chunk = decoder.decode(value);
          const lines = chunk.split('\n');

          for (const line of lines) {
            if (!line.startsWith('data: ')) continue;

            const data = line.slice(6);
            if (data === '[DONE]') continue;

            try {
              const event = JSON.parse(data);

              if (event.message) {
                const msg = event.message as Message;
                currentMessages = pushMessage(currentMessages, msg);
                setMessages(currentMessages);
              }

              if (event.error) {
                console.error('Stream error:', event.error);
                setChatState(ChatState.Idle);
                return;
              }

              if (event.finish) {
                setChatState(ChatState.Idle);
                onStreamFinish?.();
                return;
              }
            } catch (e) {
              console.error('Failed to parse SSE:', e);
            }
          }
        }
      } catch (error) {
        if (error instanceof Error && error.name !== 'AbortError') {
          console.error('Stream error:', error);
        }
        setChatState(ChatState.Idle);
      }
    },
    [sessionId, messages, setMessages, onStreamFinish]
  );

  const stopStreaming = useCallback(() => {
    abortControllerRef.current?.abort();
    setChatState(ChatState.Idle);
  }, []);

  return {
    chatState,
    handleSubmit,
    stopStreaming,
  };
}