chat_web_app/llm-fe/src/llm-fe/contexts/MessageContext.tsx

import { createContext, ReactNode, useContext, useEffect, useRef, useState } from "react";
import { WebSocketContext } from "./WebSocketContext";
import { AccountContext } from "./AccountContext";
import { ConversationContext } from "./ConversationContext";
import { ConversationPrompt, ConversationPromptType } from "../data";
import { axiosInstance } from "../../axiosApi";
import { AxiosResponse } from "axios";
import { AnalyticsEvents, trackEvent } from "../utils/analytics";

type MessageProviderProps ={
    children? : ReactNode;
}

type IMessageContext = {
    stateMessage: string;
    setStateMessage: (message: string) => void;
    conversationDetails: ConversationPrompt[];
    setConversationDetails: (conversationPrompts: ConversationPrompt[]) => void;
    isGeneratingMessage: boolean;
    /** True when a stream was interrupted by a socket drop; cleared on refetch/retry. */
    streamInterrupted: boolean;
    clearStreamInterrupted: () => void;
}

const initialValues = {
    stateMessage: '',
    setStateMessage: () => {},
    conversationDetails: [],
    setConversationDetails: () => {},
    isGeneratingMessage: false,
    streamInterrupted: false,
    clearStreamInterrupted: () => {},
}

const MessageContext = createContext<IMessageContext>(initialValues);

const STREAM_INTERRUPTED_MESSAGE =
    '_Response interrupted — connection lost. Re-send your message if needed._';

const MessageProvider = ( {children}: MessageProviderProps) => {
    const [subscribe, unsubscribe, , , , , reconnectGeneration] = useContext(WebSocketContext)
    const { account } = useContext(AccountContext)
    const {selectedConversation, setSelectedConversation} = useContext(ConversationContext);

    const [stateMessage, setStateMessage] = useState<string>('')
    const [conversationDetails, setConversationDetails] = useState<ConversationPrompt[]>([])
    const [isGeneratingMessage, setIsGeneratingMessage] = useState<boolean>(false)
    const [streamInterrupted, setStreamInterrupted] = useState<boolean>(false)

    const messageRef = useRef('')
    const messageResponsePart = useRef(0);
    const conversationRef = useRef(conversationDetails)
    const selectedConversationRef = useRef<undefined | number>(undefined)
    const isGeneratingRef = useRef(false)
    const prevReconnectGenerationRef = useRef(0)

    const clearStreamInterrupted = () => setStreamInterrupted(false)

    async function GetConversationDetails(conversationId: number | undefined) {
        if (!conversationId) {
            setConversationDetails([])
            return
        }

        try {
            selectedConversationRef.current = conversationId
            const { data }: AxiosResponse<ConversationPromptType[]> =
                await axiosInstance.get(
                    `conversation_details?conversation_id=${conversationId}`,
                )

            const tempConversations: ConversationPrompt[] = data.map(
                (item) =>
                    new ConversationPrompt({
                        id: item.id,
                        message: item.message,
                        user_created: item.user_created,
                        created_timestamp: item.created_timestamp,
                        tokens_in: item.tokens_in ?? null,
                        tokens_out: item.tokens_out ?? null,
                    }),
            )
            if (tempConversations.length === 1) {
                // first message still awaiting assistant reply placeholder
                tempConversations.push(
                    new ConversationPrompt({ message: '', user_created: false }),
                )
            }
            conversationRef.current = tempConversations
            setConversationDetails(tempConversations)
        } catch (err) {
            console.log('Failed to refetch conversation after reconnect', err)
        }
    }

    useEffect(() => {
        GetConversationDetails(selectedConversation)
    }, [selectedConversation])

    // Streaming recovery: on reconnect after a drop mid-stream, refetch conversation
    // and mark the in-flight assistant turn interrupted (no resume protocol on backend).
    useEffect(() => {
        if (reconnectGeneration === prevReconnectGenerationRef.current) {
            return
        }
        prevReconnectGenerationRef.current = reconnectGeneration
        if (reconnectGeneration === 0) {
            return
        }

        const wasGenerating = isGeneratingRef.current || messageResponsePart.current !== 0
        if (!wasGenerating) {
            return
        }

        const partial = messageRef.current
        messageRef.current = ''
        messageResponsePart.current = 0
        isGeneratingRef.current = false
        setIsGeneratingMessage(false)
        setStateMessage('')
        setStreamInterrupted(true)

        const details = [...conversationRef.current]
        if (details.length > 0) {
            const last = details[details.length - 1]
            if (last && !last.user_created && (!last.message || last.message.length === 0)) {
                details[details.length - 1] = new ConversationPrompt({
                    message: partial
                        ? `${partial}\n\n${STREAM_INTERRUPTED_MESSAGE}`
                        : STREAM_INTERRUPTED_MESSAGE,
                    user_created: false,
                })
                conversationRef.current = details
                setConversationDetails(details)
            }
        }

        if (selectedConversationRef.current) {
            GetConversationDetails(selectedConversationRef.current)
        }
    }, [reconnectGeneration])

    useEffect(() => {
        /* register a consistent channel name for identifing this chat messages */
        const channelName = `ACCOUNT_ID_${account?.email}`
  
        /* subscribe to channel and register callback */        
        subscribe(channelName, (message: string) => {
            /* when a message is received just add it to the UI */
            
            if (message === 'END_OF_THE_STREAM_ENDER_GAME_42'){
                messageResponsePart.current = 0
            
                conversationRef.current.pop()
            
                //handleAssistantPrompt({prompt: messageRef.current})
                setConversationDetails([...conversationRef.current, new ConversationPrompt({message: `${messageRef.current}`, user_created:false})])
                console.log([...conversationRef.current, new ConversationPrompt({message: `${messageRef.current}`, user_created:false})])
                messageRef.current = ''
                setStateMessage('')
                isGeneratingRef.current = false
                setIsGeneratingMessage(false)
                setStreamInterrupted(false)
            }
            else if (message === 'START_OF_THE_STREAM_ENDER_GAME_42'){
                conversationRef.current = conversationDetails
                isGeneratingRef.current = true
                setIsGeneratingMessage(true)
                setStreamInterrupted(false)
                messageResponsePart.current = 2
  
            }else if (message === 'CONVERSATION_ID'){
                isGeneratingRef.current = true
                setIsGeneratingMessage(true)
                setStreamInterrupted(false)
                messageResponsePart.current = 1
            }else{
                isGeneratingRef.current = true
                setIsGeneratingMessage(true)
                if (messageResponsePart.current === 1){
                    // this has to do with the conversation id
                    if(!selectedConversation){
                        const conversationId = Number(message);
                        setSelectedConversation(conversationId)
                        trackEvent(AnalyticsEvents.CONVERSATION_CREATED, {
                          conversationId,
                        });
                    }   
                }
                else if (messageResponsePart.current === 2){
                    messageRef.current += message
                    setStateMessage(messageRef.current)
  
                }
            }
            
        })
  
        return () => {
            /* unsubscribe from channel during cleanup */
            unsubscribe(channelName)
        }
    }, [account, subscribe, unsubscribe, conversationDetails, selectedConversation, setSelectedConversation])


    return(
        <MessageContext.Provider value={{
            stateMessage,
            setStateMessage,
            conversationDetails,
            setConversationDetails,
            isGeneratingMessage,
            streamInterrupted,
            clearStreamInterrupted,
        }}>
            {children}
        </MessageContext.Provider>
    )

}

export { MessageContext, MessageProvider}