import { createContext, useContextSelector } from 'use-context-selector'; import { ReactNode, useMemo, useCallback, useRef, useEffect } from 'react'; import { useQuery, useQueryClient } from '@tanstack/react-query'; import { v4 as uuidv4 } from 'uuid'; import { API_BASE_URL, TARGET_ENV, ASSISTANT_ENABLED, PLUGIN_ASSISTANT_SETTINGS, PLUGIN_REST_ENDPOINTS, getRestNonce } from '../constants'; import { useState } from 'react'; import type { ChatActivity, ChatApprovalState, ChatStreamEvent, ChatToolSurface, DelegatedToolRequest, FinalizedThinkingSummary, ThinkingTimelineEntry, } from '../types/chatStream'; import { useSetIsRateLimitModalOpen } from './UIContext'; import { useOpenProposeChanges, useOpenReadApproval } from './UIContext'; const DEBUG_UI = false; const MAX_PARALLEL_DELEGATED_READS = 10; // Empty arrays for stable references const EMPTY_MESSAGES: Message[] = []; const EMPTY_CONVERSATIONS: Conversation[] = []; const EMPTY_ARRAY: any[] = []; // Message interface export interface Message { message_id: string; sent_at: number; role: 'user' | 'assistant' | 'assistant-streaming' | 'error'; content: string; userMessageId?: string; } // Conversation metadata interface export interface Conversation { conversation_id: string; title: string; created_at: number; updated_at: number; } // Client-only conversation state interface EphemeralConversationState { phase: 'IDLE' | 'THINKING' | 'STREAMING' | 'ERROR'; isPersisted: boolean; // On server statusLabel?: string; activities: ChatActivity[]; approvalState?: ChatApprovalState; reasoningText?: string; timeline: ThinkingTimelineEntry[]; hasStartedStreaming: boolean; startedAt?: number; thoughtDurationMs?: number; } // Chat provider props interface interface ChatProviderProps { children: ReactNode; } // Types for currently streaming message type PartialContentsMap = Record; type LatestChunksMap = Record; // Chat context interface interface ChatContextProps { currentConversationId: string | null; setCurrentConversationId: (conversationId: string | null) => void; conversations: Conversation[]; messages: Message[]; ephemeralMap: Record; finalizedThinkingMap: Record; sendMessage: (message: string) => Promise; retryError: (userMessageId: string) => Promise; dismissError: (userMessageId: string) => Promise; deleteConversation: (conversationId: string) => Promise; cloneSeoPage: (seoPageId: string) => Promise; setOnStreamChunk?: (cb?: (msgId: string, partialText: string, opts?: { isFinal?: boolean }) => void) => void; partialContents: PartialContentsMap; latestChunks: LatestChunksMap; } // Default context value (for typescript) const defaultContextValue: ChatContextProps = { currentConversationId: null, setCurrentConversationId: () => {}, conversations: [], messages: [], ephemeralMap: {}, finalizedThinkingMap: {}, sendMessage: async () => {}, retryError: async () => {}, dismissError: async () => {}, deleteConversation: async () => {}, cloneSeoPage: async () => {}, partialContents: {}, latestChunks: {}, }; // Chat context const ChatContext = createContext(defaultContextValue); const DEFAULT_EPHEMERAL_STATE: EphemeralConversationState = { phase: 'IDLE', isPersisted: false, activities: [], timeline: [], hasStartedStreaming: false, }; const TOOL_LABELS: Record = { read_site: 'Checking your site', propose_changes: 'Preparing changes', }; const REASONING_STATUS_LABEL = 'Analyzing results'; const TOOL_SURFACES: Record = { read_site: 'site', propose_changes: 'site', search_wordpress_documentation: 'web', search_wordpress_forums: 'web', search_specific_wordpress_plugin: 'web', search_wordpress_plugins: 'web', get_external_pages: 'web', get_atiba_marketing_material: 'web', }; const getStatusLabelForTool = (name?: string) => TOOL_LABELS[name || ''] || 'Thinking'; const getSurfaceForTool = (name?: string): ChatToolSurface => TOOL_SURFACES[name || ''] || 'internal'; const isVisibleToolActivity = (name?: string) => name === 'read_site' || name === 'propose_changes'; const getRunningVisibleTool = (activities: ChatActivity[]) => ( [...activities].reverse().find(activity => ( isVisibleToolActivity(activity.name) && activity.state === 'running' )) ); const getLiveStatusLabel = ( activities: ChatActivity[], approvalState?: ChatApprovalState, options?: { preferReasoning?: boolean } ) => { if (approvalState?.state === 'needed') return 'Waiting for approval'; const runningVisibleTool = getRunningVisibleTool(activities); if (runningVisibleTool) return getStatusLabelForTool(runningVisibleTool.name); if (options?.preferReasoning) return REASONING_STATUS_LABEL; return 'Thinking'; }; const buildActivityLabel = (name?: string, target?: string) => { if (target?.trim()) return target.trim(); if (name === 'propose_changes') return 'Proposed changes'; return getStatusLabelForTool(name); }; const buildConversationTurnKey = (conversationId: string, userTurnIndex: number) => `${conversationId}:turn:${userTurnIndex}`; const getApprovalTimelineLabel = (state: ChatApprovalState['state'], target?: string) => { if (state === 'granted') { return target ? `Approval granted: ${target}` : 'Approval granted'; } if (state === 'rejected') { return target ? `Approval rejected: ${target}` : 'Approval rejected'; } return target ? `Waiting for approval: ${target}` : 'Waiting for approval'; }; const sanitizeReasoningChunk = (text: string) => { return text .replace(/\*\*\s*\*\*/g, ' ') .replace(/__\s*__/g, ' ') .replace(/[ \t]+\n/g, '\n') .replace(/\n[ \t]+/g, '\n') .replace(/ {2,}/g, ' '); }; const startsWithReasoningHeading = (text: string) => ( /^\s*(\*\*[^*\n][^*\n]*\*\*|#{1,6}\s+)/.test(text) ); const appendReasoningText = (current: string | undefined, incoming: string) => { const safeCurrent = current || ''; const safeIncoming = incoming || ''; const separator = ( safeCurrent && safeIncoming && startsWithReasoningHeading(safeIncoming) && !/\n\s*$/.test(safeCurrent) ) ? '\n\n' : ''; const combined = `${safeCurrent}${separator}${safeIncoming}`; return sanitizeReasoningChunk(combined).trim(); }; // Chat provider component export const ChatProvider = ({ children }: ChatProviderProps) => { // React Query client const queryClient = useQueryClient(); // State variables const [currentConversationId, setCurrentConversationId] = useState(null); const [partialContents, setPartialContents] = useState({}); const [latestChunks, setLatestChunks] = useState({}); const [ephemeralMap, setEphemeralMap] = useState>({}); const [finalizedThinkingMap, setFinalizedThinkingMap] = useState>({}); const setIsRateLimitModalOpen = useSetIsRateLimitModalOpen(); const openProposeChanges = useOpenProposeChanges(); const openReadApproval = useOpenReadApproval(); // Reader and abort controller refs const readerRefs = useRef>({}); const abortRefs = useRef>({}); const ephemeralMapRef = useRef>({}); // Tracks conversations for which site context has already been attempted. const contextAttemptedForConversationRef = useRef>(new Set()); // Stream callback ref to update partial content const streamCallbackRef = useRef<(msgId: string, partialText: string, opts?: { isFinal?: boolean }) => void>(); // ================= // // REACT QUERY STUFF // // ================= // const assistantEnabled = ASSISTANT_ENABLED; const requireReadApproval = Boolean(PLUGIN_ASSISTANT_SETTINGS.requireReadApproval); const updateConversationState = useCallback(( conversationId: string, partial: Partial ): EphemeralConversationState => { let newState: EphemeralConversationState; setEphemeralMap(prev => { const oldState = prev[conversationId] ?? DEFAULT_EPHEMERAL_STATE; newState = { ...oldState, ...partial }; if ( oldState.phase === newState.phase && oldState.isPersisted === newState.isPersisted && oldState.statusLabel === newState.statusLabel && oldState.reasoningText === newState.reasoningText && oldState.timeline === newState.timeline && oldState.hasStartedStreaming === newState.hasStartedStreaming && oldState.startedAt === newState.startedAt && oldState.thoughtDurationMs === newState.thoughtDurationMs && oldState.activities === newState.activities && oldState.approvalState === newState.approvalState ) { return prev; } const nextMap = { ...prev, [conversationId]: newState }; ephemeralMapRef.current = nextMap; return nextMap; }); return newState!; }, []); // Function to fetch conversations const fetchConversations = useCallback(async (): Promise => { if (!assistantEnabled) { return []; } const response = await fetch(`${API_BASE_URL}/conversations`, { method: 'GET', credentials: 'include', }); if (!response.ok) { throw new Error('Failed to fetch conversations'); } // Update the ephemeral map with the new conversation data const data = await response.json(); data.forEach((conversation: Conversation) => { updateConversationState(conversation.conversation_id, { ...DEFAULT_EPHEMERAL_STATE, isPersisted: true, }); }); return data; }, [assistantEnabled, updateConversationState]); // Function to fetch messages for current conversation const fetchMessages = useCallback(async (): Promise => { if (!assistantEnabled || !currentConversationId) return []; const response = await fetch(`${API_BASE_URL}/conversations/${currentConversationId}`, { method: 'GET', credentials: 'include', }); if (!response.ok) { throw new Error('Failed to fetch messages'); } const data = await response.json(); // Map `message_id` to `id` const mappedMessages: Message[] = data.messages.map((msg: any) => ({ message_id: msg.message_id, sent_at: msg.sent_at, role: msg.role, content: msg.content, })); return mappedMessages; }, [assistantEnabled, currentConversationId]); // React Query hook to fetch conversations const { data: rawConversations = EMPTY_CONVERSATIONS, } = useQuery({ queryKey: ['conversations'], queryFn: fetchConversations, enabled: assistantEnabled && Object.values(ephemeralMap).every(state => state.phase == 'IDLE'), retry: false, staleTime: 30000, placeholderData: EMPTY_CONVERSATIONS, }); // Conversation state variable const conversations = rawConversations.length > 0 ? rawConversations : EMPTY_CONVERSATIONS; // React Query hook to fetch messages for current conversation const { data: rawMessages = EMPTY_MESSAGES, } = useQuery({ queryKey: ['messages', currentConversationId], queryFn: fetchMessages, enabled: (assistantEnabled && !!currentConversationId && ephemeralMap[currentConversationId]?.phase === 'IDLE' && ephemeralMap[currentConversationId]?.isPersisted), staleTime: 1000 * 60, // 1 minute gcTime: 1000 * 60, // 1 minute }); // ================ // // HELPER FUNCTIONS // // ================ // // Update the ephemeral state of a conversation // Add user message to an existing conversation const addNewMessage = ( role: 'user' | 'assistant' | 'error', messageContent: string, targetConversationId: string, userMessageId?: string ) => { // Create user message const newMessageId = uuidv4(); const newMessage: Message = { message_id: newMessageId, sent_at: Date.now(), role: role, content: messageContent, userMessageId: userMessageId, }; // Add user message to react query queryClient.setQueryData( ['messages', targetConversationId], (old = []) => [...old, newMessage] ); return newMessageId; }; // Add a new conversation to react query const addNewConversation = (messageContent: string) => { // Create new conversation const newConversationId = uuidv4(); const newConversation: Conversation = { conversation_id: newConversationId, title: messageContent, created_at: Date.now(), updated_at: Date.now(), }; // Add it to react query queryClient.setQueryData(['conversations'], (old = []) => [ newConversation, ...old, ]); return newConversationId; }; // Send a chat request to the server const sendChatRequest = async ( endpoint: string, messageContent: string, conversationId: string, newConversation: boolean, signal: AbortSignal, extraBody?: any ) => { if (!assistantEnabled) { const error = new Error('AssistantDisabled'); error.name = 'AssistantDisabledError'; throw error; } // Build payload; in WordPress mode include site_context and source: 'plugin' let payload: any = { message: messageContent, conversation_id: conversationId, new_conversation: newConversation, source: 'web' as 'web' | 'plugin', }; if (assistantEnabled && TARGET_ENV === 'wordpress') { payload.source = 'plugin'; payload.plugin_settings = { allow_log_reads: Boolean(PLUGIN_ASSISTANT_SETTINGS.allowLogReads), allow_write_operations: Boolean(PLUGIN_ASSISTANT_SETTINGS.allowWriteOperations), require_read_approval: requireReadApproval, }; // Fetch site context only once per conversation. const shouldAttemptSiteContext = !requireReadApproval && newConversation && !contextAttemptedForConversationRef.current.has(conversationId); if (shouldAttemptSiteContext) { // Mark before fetch so retries/reruns don't repeatedly hit the context endpoint. contextAttemptedForConversationRef.current.add(conversationId); try { const ctxResp = await fetch(PLUGIN_REST_ENDPOINTS.context, { method: 'GET', credentials: 'include', headers: { 'Content-Type': 'application/json', 'X-WP-Nonce': getRestNonce() }, signal, }); if (ctxResp.ok) { const context = await ctxResp.json(); payload.site_context = context; } } catch (e) { // If diagnostics fail, proceed without blocking the chat } } } // Merge any extra body (e.g., tool_resume) if (extraBody && typeof extraBody === 'object') { payload = { ...payload, ...extraBody }; } const response = await fetch(endpoint, { method: 'POST', credentials: 'include', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(payload), signal: signal, }); if (!response.ok) { if (response.status === 429) { const error = new Error('Rate limit exceeded'); error.name = 'RateLimitError'; throw error; } throw new Error(response.statusText); } return response; }; // Helper to set or clear partial text and recent chunks const setPartialContent = useCallback( (targetConversationId: string, content: string | null, chunk: string | null) => { // Update partial contents setPartialContents(old => { const oldContent = old[targetConversationId] ?? null; // If the new content is identical to old content, skip the update: if (oldContent === content) { return old; } const copy = { ...old }; if (content === null) { delete copy[targetConversationId]; } else { copy[targetConversationId] = content; } return copy; }); // Update recent chunks setLatestChunks(old => { const copy = { ...old }; if (content === null) { delete copy[targetConversationId]; } else if (chunk) { copy[targetConversationId] = chunk; } return copy; }); }, [] ); const patchConversationState = useCallback(( targetConversationId: string, updater: (current: EphemeralConversationState) => EphemeralConversationState ) => { setEphemeralMap(prev => { const current = prev[targetConversationId] ?? DEFAULT_EPHEMERAL_STATE; const next = updater(current); const nextMap = { ...prev, [targetConversationId]: next }; ephemeralMapRef.current = nextMap; return nextMap; }); }, []); const setFinalizedThinkingSummary = useCallback((messageId: string, summary: FinalizedThinkingSummary | null) => { setFinalizedThinkingMap(prev => { if (!summary) { if (!(messageId in prev)) return prev; const copy = { ...prev }; delete copy[messageId]; return copy; } return { ...prev, [messageId]: summary }; }); }, []); const applyStreamEvent = useCallback((targetConversationId: string, event: ChatStreamEvent) => { if (event.type === 'status') { patchConversationState(targetConversationId, (current) => { if (event.phase === 'drafting') { return { ...current, phase: 'STREAMING', hasStartedStreaming: true, thoughtDurationMs: current.thoughtDurationMs ?? ( current.startedAt ? Math.max(1, Date.now() - current.startedAt) : current.thoughtDurationMs ), }; } return { ...current, phase: 'THINKING', statusLabel: event.label, approvalState: event.phase === 'approval' ? current.approvalState : undefined, }; }); return; } if (event.type === 'tool_start') { patchConversationState(targetConversationId, (current) => { let activities = current.activities; let timeline = current.timeline; if (isVisibleToolActivity(event.name)) { const label = buildActivityLabel(event.name, event.target); const existingIdx = current.activities.findIndex(activity => activity.id === event.id); const nextActivity: ChatActivity = { id: event.id, name: event.name, surface: event.surface, state: 'running', label, target: event.target, batchId: event.batchId, }; activities = existingIdx < 0 ? [...current.activities, nextActivity] : current.activities.map((activity, idx) => idx === existingIdx ? { ...activity, ...nextActivity } : activity); const existingTimelineIdx = current.timeline.findIndex(item => item.id === `tool-${event.id}`); const nextTimelineEntry: ThinkingTimelineEntry = { id: `tool-${event.id}`, kind: 'tool', label, state: 'running', }; timeline = existingTimelineIdx < 0 ? [...current.timeline, nextTimelineEntry] : current.timeline.map((item, idx) => idx === existingTimelineIdx ? nextTimelineEntry : item); } return { ...current, phase: 'THINKING', statusLabel: getLiveStatusLabel(activities), approvalState: undefined, activities, timeline, }; }); return; } if (event.type === 'tool_end') { patchConversationState(targetConversationId, (current) => { const activities = current.activities.map(activity => ( activity.id === event.id ? { ...activity, state: event.ok ? 'done' as const : 'error' as const, summary: event.summary, } : activity )); const timeline = current.timeline.map(item => ( item.id === `tool-${event.id}` && item.kind === 'tool' ? { ...item, state: event.ok ? 'done' as const : 'error' as const } : item )); return { ...current, activities, timeline, statusLabel: current.hasStartedStreaming ? current.statusLabel : getLiveStatusLabel(activities, current.approvalState, { preferReasoning: Boolean(current.reasoningText?.trim()), }), }; }); return; } if (event.type === 'approval') { patchConversationState(targetConversationId, (current) => ({ ...current, phase: 'THINKING', statusLabel: event.state === 'needed' ? 'Waiting for approval' : getStatusLabelForTool(event.tool), approvalState: { state: event.state, tool: event.tool, target: event.target, }, timeline: [ ...current.timeline, { id: `approval-${event.state}-${current.timeline.length}`, kind: 'approval', state: event.state, label: getApprovalTimelineLabel(event.state, event.target), }, ], })); return; } if (event.type === 'reasoning_delta') { patchConversationState(targetConversationId, (current) => ({ ...current, statusLabel: current.hasStartedStreaming ? current.statusLabel : getLiveStatusLabel(current.activities, current.approvalState, { preferReasoning: true }), reasoningText: appendReasoningText(current.reasoningText, event.text), timeline: (() => { const sanitizedIncoming = sanitizeReasoningChunk(event.text); if (!sanitizedIncoming.trim()) return current.timeline; const lastEntry = current.timeline[current.timeline.length - 1]; if (lastEntry?.kind === 'reasoning') { const mergedText = appendReasoningText(lastEntry.text, event.text); if (!mergedText) return current.timeline; return [ ...current.timeline.slice(0, -1), { ...lastEntry, text: mergedText, }, ]; } return [ ...current.timeline, { id: `reasoning-${current.timeline.length}`, kind: 'reasoning', text: sanitizedIncoming.trim(), }, ]; })(), })); return; } if (event.type === 'message_delta') { patchConversationState(targetConversationId, (current) => ({ ...current, phase: 'STREAMING', hasStartedStreaming: true, thoughtDurationMs: current.thoughtDurationMs ?? ( current.startedAt ? Math.max(1, Date.now() - current.startedAt) : current.thoughtDurationMs ), })); return; } if (event.type === 'message_done') { patchConversationState(targetConversationId, (current) => ({ ...current, phase: 'STREAMING', hasStartedStreaming: true, thoughtDurationMs: current.thoughtDurationMs ?? ( current.startedAt ? Math.max(1, Date.now() - current.startedAt) : current.thoughtDurationMs ), })); } }, [patchConversationState]); const clearConversationLocally = useCallback((conversationId: string) => { queryClient.setQueryData(['conversations'], (old = []) => old.filter(conv => conv.conversation_id !== conversationId) ); queryClient.removeQueries({ queryKey: ['messages', conversationId], exact: true }); setPartialContent(conversationId, null, null); setEphemeralMap(prev => { if (!(conversationId in prev)) return prev; const copy = { ...prev }; delete copy[conversationId]; ephemeralMapRef.current = copy; return copy; }); setFinalizedThinkingMap(prev => { const prefix = `${conversationId}:turn:`; let changed = false; const next = Object.fromEntries( Object.entries(prev).filter(([key]) => { if (key.startsWith(prefix)) { changed = true; return false; } return true; }) ); return changed ? next : prev; }); }, [queryClient, setPartialContent]); const deleteConversationSilently = useCallback(async (conversationId: string) => { try { await fetch(`${API_BASE_URL}/conversations/${conversationId}`, { method: 'DELETE', credentials: 'include', }); } catch (error) { console.error('Error deleting conversation:', error); } finally { clearConversationLocally(conversationId); if (currentConversationId === conversationId) { setCurrentConversationId(null); } } }, [clearConversationLocally, currentConversationId, setCurrentConversationId]); const buildFinalizedThinkingSummary = useCallback((conversationId: string): FinalizedThinkingSummary | null => { const state = ephemeralMapRef.current[conversationId] ?? DEFAULT_EPHEMERAL_STATE; const durationMs = state.thoughtDurationMs ?? ( state.startedAt ? Math.max(1, Date.now() - state.startedAt) : 0 ); const visibleTimeline = state.timeline.filter(item => ( item.kind !== 'reasoning' || item.text.trim() )); if (visibleTimeline.length === 0) { return null; } return { durationMs: Math.max(1, durationMs), activities: state.activities.filter(activity => isVisibleToolActivity(activity.name)), approvalState: state.approvalState, reasoningText: state.reasoningText?.trim() || undefined, timeline: visibleTimeline, }; }, []); const getCurrentUserTurnKey = useCallback((conversationId: string): string | null => { const conversationMessages = queryClient.getQueryData(['messages', conversationId]) || EMPTY_MESSAGES; let userTurnIndex = 0; for (const message of conversationMessages) { if (message.role === 'user') { userTurnIndex += 1; } } if (userTurnIndex <= 0) return null; return buildConversationTurnKey(conversationId, userTurnIndex); }, [queryClient]); const isFailedFirstTurnConversation = useCallback((conversationMessages: Message[], userMessageId: string) => { const nonErrorMessages = conversationMessages.filter(msg => msg.role !== 'error'); return ( nonErrorMessages.length === 1 && nonErrorMessages[0].role === 'user' && nonErrorMessages[0].message_id === userMessageId ); }, []); const deletePersistedUserMessage = useCallback(async ( conversationId: string, userMessageId: string, sentAt: number ) => { try { await fetch( `${API_BASE_URL}/conversations/${conversationId}/messages/${userMessageId}?sent_at=${sentAt}`, { method: 'DELETE', credentials: 'include', } ); } catch (error) { console.error('Error deleting persisted user message:', error); } }, []); // Handles the streaming response from the server const handleStreamingResponse = useCallback(async ( initialReader: ReadableStreamDefaultReader, targetConversationId: string ) => { let assistantMessageContent = ''; const decoder = new TextDecoder(); const TOOL_PREFIX = '[TOOL_REQUEST]:'; const getToolTarget = (toolReq: DelegatedToolRequest) => { if (toolReq.name === 'read_site') { const resource = toolReq.arguments?.['resource']; return typeof resource === 'string' ? resource : undefined; } if (toolReq.name === 'propose_changes') { const goal = toolReq.arguments?.['goal']; return typeof goal === 'string' ? goal : undefined; } return undefined; }; const parseToolResponse = async (resp: Response) => { const contentType = resp.headers.get('Content-Type') || ''; if (contentType.includes('application/json')) { return await resp.json(); } const text = await resp.text(); return text ? { data: text } : {}; }; const executeToolRequests = async ( toolReqs: DelegatedToolRequest[], batchId?: string ): Promise> => { const executeSingleToolRequest = async (toolReq: DelegatedToolRequest) => { const target = getToolTarget(toolReq); applyStreamEvent(targetConversationId, { type: 'tool_start', id: toolReq.id, name: toolReq.name, surface: getSurfaceForTool(toolReq.name), target, batchId, }); let toolResult: any = null; let ok = true; let summary: string | undefined; if (toolReq.name === 'read_site') { const wpNonce = getRestNonce(); let responsePayload: any = null; try { const rsResp = await fetch(PLUGIN_REST_ENDPOINTS.readSite, { method: 'POST', credentials: 'include', headers: { 'Content-Type': 'application/json', 'X-WP-Nonce': wpNonce }, body: JSON.stringify(toolReq.arguments || {}) }); responsePayload = await parseToolResponse(rsResp); if (!rsResp.ok) { responsePayload = { ...(typeof responsePayload === 'object' && responsePayload ? responsePayload : {}), error: (responsePayload && responsePayload.error) || 'http_error', status: rsResp.status, status_text: rsResp.statusText, }; } } catch (err) { responsePayload = { error: 'read_site_request_failed', message: err instanceof Error ? err.message : String(err), }; } if (requireReadApproval) { applyStreamEvent(targetConversationId, { type: 'approval', state: 'needed', tool: toolReq.name, target, }); const decision = await openReadApproval({ request: toolReq, response: responsePayload, }); applyStreamEvent(targetConversationId, { type: 'approval', state: decision.approved ? 'granted' : 'rejected', tool: toolReq.name, target, }); if (decision.approved) { toolResult = responsePayload; } else { toolResult = { status: 'declined', message: 'The user has enabled "Per-Request Approval for Reads" and declined to share the requested information.', user_reason: decision.user_reason || null, request: toolReq.arguments || {}, }; ok = false; summary = 'Read declined'; } } else { toolResult = responsePayload; } } else if (toolReq.name === 'propose_changes') { try { applyStreamEvent(targetConversationId, { type: 'approval', state: 'needed', tool: toolReq.name, target, }); const decision = await openProposeChanges(toolReq.arguments || toolReq); applyStreamEvent(targetConversationId, { type: 'approval', state: decision.approved ? 'granted' : 'rejected', tool: toolReq.name, target, }); if (decision.approved) { const wpNonce = getRestNonce(); const execResp = await fetch(PLUGIN_REST_ENDPOINTS.executeChanges, { method: 'POST', credentials: 'include', headers: { 'Content-Type': 'application/json', 'X-WP-Nonce': wpNonce }, body: JSON.stringify(toolReq.arguments || {}) }); const responsePayload = await parseToolResponse(execResp); toolResult = execResp.ok ? responsePayload : { ...(typeof responsePayload === 'object' && responsePayload ? responsePayload : {}), error: (responsePayload && responsePayload.error) || 'http_error', status: execResp.status, status_text: execResp.statusText, }; } else { toolResult = { status: 'rejected', user_reason: decision.user_reason || null }; ok = false; summary = 'Change request rejected'; } } catch (err) { toolResult = { error: 'propose_changes_request_failed', message: err instanceof Error ? err.message : String(err), }; ok = false; } } else { toolResult = { error: 'unsupported_tool' }; ok = false; summary = 'Unsupported tool'; } if (toolResult?.error) { ok = false; summary = summary || toolResult.message || toolResult.error; } applyStreamEvent(targetConversationId, { type: 'tool_end', id: toolReq.id, ok, summary, batchId, }); return { tool_call: toolReq, result: toolResult }; }; const allReadSite = toolReqs.every(req => req?.name === 'read_site'); const resumeItems = allReadSite ? [ ...(await Promise.all(toolReqs.slice(0, MAX_PARALLEL_DELEGATED_READS).map(req => executeSingleToolRequest(req)))), ...toolReqs.slice(MAX_PARALLEL_DELEGATED_READS).map((req: DelegatedToolRequest) => { applyStreamEvent(targetConversationId, { type: 'tool_start', id: req.id, name: req.name, surface: getSurfaceForTool(req.name), target: getToolTarget(req), batchId, }); applyStreamEvent(targetConversationId, { type: 'tool_end', id: req.id, ok: false, summary: `Parallel tool limit exceeded (${MAX_PARALLEL_DELEGATED_READS})`, batchId, }); return { tool_call: req, result: { error: 'parallel_tool_limit_exceeded', message: `This delegated batch exceeded the parallel limit of ${MAX_PARALLEL_DELEGATED_READS}. Retry with fewer parallel calls.`, }, }; }), ] : await (async () => { const acc: any[] = []; for (const req of toolReqs) { acc.push(await executeSingleToolRequest(req)); } return acc; })(); const cont = await sendChatRequest( `${API_BASE_URL}/chat`, '', targetConversationId, false, abortRefs.current[targetConversationId!]?.signal ?? new AbortController().signal, resumeItems.length === 1 ? { tool_resume: resumeItems[0] } : { tool_resumes: resumeItems } ); const contReader = cont.body?.getReader(); if (!contReader) throw new Error('ReadableStream not supported (resume)'); readerRefs.current[targetConversationId!] = contReader; return contReader; }; const handleEvent = async ( event: ChatStreamEvent ): Promise | null> => { if (event.type === 'tool_request') { if (!event.requests.length) { throw new Error('Empty tool request'); } return await executeToolRequests(event.requests, event.batchId); } if (event.type === 'error') { throw new Error(event.message); } applyStreamEvent(targetConversationId, event); if (event.type === 'message_delta') { assistantMessageContent += event.text; setPartialContent(targetConversationId, assistantMessageContent, event.text); if (streamCallbackRef.current) { streamCallbackRef.current(targetConversationId, assistantMessageContent); } } return null; }; const processLegacyStream = async ( reader: ReadableStreamDefaultReader, initialChunk: string ): Promise | null> => { let first = true; let pendingChunk: string | null = initialChunk; while (true) { if (pendingChunk === null) { return null; } const chunk = pendingChunk; pendingChunk = null; if (chunk.startsWith('[ERROR]:')) { throw new Error(chunk); } if (first) { if (chunk.startsWith(TOOL_PREFIX)) { const jsonStr = chunk.slice(TOOL_PREFIX.length).trim(); let toolReqPayload: any = null; try { toolReqPayload = JSON.parse(jsonStr); } catch { throw new Error('Malformed tool request'); } const toolReqs: DelegatedToolRequest[] = Array.isArray(toolReqPayload?.requests) ? toolReqPayload.requests : [toolReqPayload]; return executeToolRequests(toolReqs); } applyStreamEvent(targetConversationId, { type: 'status', phase: 'drafting', label: '' }); first = false; } assistantMessageContent += chunk; setPartialContent(targetConversationId, assistantMessageContent, chunk); if (streamCallbackRef.current) { streamCallbackRef.current(targetConversationId, assistantMessageContent); } const next = await reader.read(); if (next.done) { return null; } pendingChunk = decoder.decode(next.value); } }; const processStream = async ( reader: ReadableStreamDefaultReader ): Promise | null> => { let mode: 'unknown' | 'ndjson' | 'legacy' = 'unknown'; let buffer = ''; let sawNdjsonEvent = false; while (true) { const { done, value } = await reader.read(); if (done) { if (mode === 'ndjson' && buffer.trim()) { const event = JSON.parse(buffer.trim()) as ChatStreamEvent; const nextReader = await handleEvent(event); if (nextReader) return nextReader; } break; } const chunk = decoder.decode(value, { stream: true }); if (mode === 'unknown') { const sample = (buffer + chunk).trimStart(); if (!sample) { continue; } mode = sample.startsWith('{') ? 'ndjson' : 'legacy'; if (mode === 'legacy') { return processLegacyStream(reader, chunk); } } buffer += chunk; while (true) { const lineBreak = buffer.indexOf('\n'); if (lineBreak < 0) break; const rawLine = buffer.slice(0, lineBreak); buffer = buffer.slice(lineBreak + 1); const line = rawLine.trim(); if (!line) continue; let event: ChatStreamEvent; try { event = JSON.parse(line) as ChatStreamEvent; } catch { if (!sawNdjsonEvent) { return processLegacyStream(reader, chunk); } throw new Error('Malformed NDJSON event stream'); } sawNdjsonEvent = true; const nextReader = await handleEvent(event); if (nextReader) return nextReader; } } return null; }; // Loop to handle possible multiple sequential tool requests let nextReader: ReadableStreamDefaultReader | null = initialReader; while (nextReader) { nextReader = await processStream(nextReader); } // Finalize assistant message if (assistantMessageContent) { const finalizedThinkingSummary = buildFinalizedThinkingSummary(targetConversationId); addNewMessage('assistant', assistantMessageContent, targetConversationId); if (finalizedThinkingSummary) { const turnKey = getCurrentUserTurnKey(targetConversationId); if (turnKey) { setFinalizedThinkingSummary(turnKey, finalizedThinkingSummary); } } } updateConversationState(targetConversationId, { ...DEFAULT_EPHEMERAL_STATE, isPersisted: true, }); if (DEBUG_UI) console.log('[UI] phase=IDLE', { conv: targetConversationId }); setPartialContent(targetConversationId, null, null); }, [ addNewMessage, applyStreamEvent, openProposeChanges, openReadApproval, requireReadApproval, sendChatRequest, buildFinalizedThinkingSummary, getCurrentUserTurnKey, setPartialContent, setFinalizedThinkingSummary, updateConversationState, ]); // Handles rate limit errors and other errors const handleChatError = useCallback(( error: Error, targetConversationId: string, userMessageId: string ) => { setPartialContent(targetConversationId!, null, null); // If it's a rate limit error... if (error.name === 'RateLimitError') { updateConversationState(targetConversationId!, { ...DEFAULT_EPHEMERAL_STATE, isPersisted: true }); // Set conversation as idle setIsRateLimitModalOpen(true); // Show rate limit modal // Refresh conversations and messages (deletes any optimistic updates not saved to server) queryClient.invalidateQueries({ queryKey: ['conversations'] }); queryClient.invalidateQueries({ queryKey: ['messages'] }); setCurrentConversationId(null); // Reset current conversation ID // Otherwise, it's a different error... } else if (userMessageId) { // Set error state updateConversationState(targetConversationId!, { ...DEFAULT_EPHEMERAL_STATE, phase: 'ERROR', }); // Add error message to conversation addNewMessage('error', 'Something went wrong. Please try again.', targetConversationId!, userMessageId); } }, [addNewMessage, queryClient, setCurrentConversationId, setIsRateLimitModalOpen, setPartialContent, updateConversationState]); // Cleanup reader and abort controller refs for a conversation const cleanup = useCallback((conversationId: string) => { if (readerRefs.current[conversationId]) { readerRefs.current[conversationId]?.cancel(); readerRefs.current[conversationId]?.releaseLock(); readerRefs.current[conversationId] = null; } if (abortRefs.current[conversationId]) { abortRefs.current[conversationId]?.abort(); abortRefs.current[conversationId] = null; } }, []); // ================== // // EXPORTED FUNCTIONS // // ================== // // Function to send a message and handle streaming const sendMessageInternal = useCallback(async ( messageContent: string, options?: { forceNewConversation?: boolean } ) => { if (!assistantEnabled) { return Promise.reject(new Error('assistant_disabled')); } // No current conversation ID means this is the first message in a new conversation const isNewConversation = Boolean(options?.forceNewConversation) || !currentConversationId; // Initialize variables let userMessageId: string | null = null; let targetConversationId: string | null = null; let currentConversationState: EphemeralConversationState | null = null // If this is the first message in a new conversation... if (isNewConversation) { // Setup new conversation targetConversationId = addNewConversation(messageContent); userMessageId = addNewMessage('user', messageContent, targetConversationId); setCurrentConversationId(targetConversationId); // If this is NOT the first message in a new conversation... } else { // Add message to existing conversation userMessageId = addNewMessage('user', messageContent, currentConversationId!); targetConversationId = currentConversationId!; } // Set conversation as thinking // Here we're getting the isPersisted state directly from the setter function return const existingConversationState = ephemeralMap[targetConversationId!] ?? DEFAULT_EPHEMERAL_STATE; currentConversationState = updateConversationState(targetConversationId!, { phase: 'THINKING', isPersisted: existingConversationState.isPersisted, statusLabel: 'Thinking', activities: [], approvalState: undefined, reasoningText: undefined, timeline: [], hasStartedStreaming: false, startedAt: Date.now(), thoughtDurationMs: undefined, }); if (DEBUG_UI) console.log('[UI] phase=THINKING', { conv: targetConversationId }); // Create an abort controller const abortController = new AbortController(); abortRefs.current[targetConversationId!] = abortController; // Try to send the message try { // Send the message const response = await sendChatRequest( `${API_BASE_URL}/chat`, messageContent, targetConversationId!, isNewConversation || !currentConversationState.isPersisted, abortController.signal ); // Get the reader for the response body const reader = response.body?.getReader(); if (!reader) { throw new Error('ReadableStream not supported'); } readerRefs.current[targetConversationId!] = reader; // Stream the response await handleStreamingResponse(reader, targetConversationId!); // Handle the error (rate limit modal or error message) } catch (error: any) { handleChatError(error, targetConversationId!, userMessageId!); // Cleanup } finally { cleanup(targetConversationId!); } }, [ addNewConversation, addNewMessage, assistantEnabled, cleanup, currentConversationId, handleChatError, handleStreamingResponse, sendChatRequest, setCurrentConversationId, updateConversationState, ephemeralMap, ]); const sendMessage = useCallback(async (messageContent: string) => { await sendMessageInternal(messageContent); }, [sendMessageInternal]); // Function to retry a message const retryError = useCallback(async (userMessageId: string) => { if (!assistantEnabled) return; if (!currentConversationId) return; const messages = queryClient.getQueryData(['messages', currentConversationId]) || EMPTY_ARRAY; const userMessage = messages.find((msg) => msg.message_id === userMessageId && msg.role === 'user'); // Remove error and user messages if (userMessage) { if (isFailedFirstTurnConversation(messages, userMessageId)) { await deleteConversationSilently(currentConversationId); await sendMessageInternal(userMessage.content, { forceNewConversation: true }); return; } await deletePersistedUserMessage(currentConversationId, userMessageId, userMessage.sent_at); queryClient.setQueryData(['messages', currentConversationId], (old = []) => old?.filter(msg => msg.message_id !== userMessageId && !(msg.role === 'error' && msg.userMessageId === userMessageId) ) ); // Resend message await sendMessageInternal(userMessage.content); } }, [ assistantEnabled, currentConversationId, deleteConversationSilently, deletePersistedUserMessage, isFailedFirstTurnConversation, queryClient, sendMessageInternal, ]); // Function to dismiss an error message (will also delete the user message that caused it) const dismissError = useCallback(async (userMessageId: string) => { if (!currentConversationId) return; const messages = queryClient.getQueryData(['messages', currentConversationId]) || EMPTY_ARRAY; const userMessage = messages.find((msg) => msg.message_id === userMessageId && msg.role === 'user'); if (isFailedFirstTurnConversation(messages, userMessageId)) { await deleteConversationSilently(currentConversationId); return; } if (userMessage) { await deletePersistedUserMessage(currentConversationId, userMessageId, userMessage.sent_at); } // Remove user message and error message queryClient.setQueryData(['messages', currentConversationId], (old = []) => old?.filter(msg => msg.message_id !== userMessageId && !(msg.role === 'error' && msg.userMessageId === userMessageId) ) ); // Set conversation as idle updateConversationState(currentConversationId, { ...DEFAULT_EPHEMERAL_STATE, isPersisted: true }); // If no messages left, reset current conversation const remainingMessages = queryClient.getQueryData(['messages', currentConversationId]) || []; if (remainingMessages.length === 0) { setCurrentConversationId(null); } }, [ currentConversationId, deleteConversationSilently, deletePersistedUserMessage, isFailedFirstTurnConversation, queryClient, setCurrentConversationId, updateConversationState, ]); // Function to delete a conversation const deleteConversation = useCallback(async (conversationId: string) => { if (!assistantEnabled) return; try { const response = await fetch(`${API_BASE_URL}/conversations/${conversationId}`, { method: 'DELETE', credentials: 'include', }); if (!response.ok) { const errorData = await response.json(); throw new Error(errorData.detail || 'Failed to delete conversation'); } queryClient.invalidateQueries({ queryKey: ['conversations'] }); queryClient.invalidateQueries({ queryKey: ['messages', conversationId] }); // If deleting current conversation if (currentConversationId === conversationId) { const remainingConversations = conversations.filter((conv) => conv.conversation_id !== conversationId); // Sort and set new current conversation remainingConversations.sort((a, b) => b.updated_at - a.updated_at); const newCurrentId = remainingConversations[0]?.conversation_id || null; setCurrentConversationId(newCurrentId); } } catch (error) { console.error('Error deleting conversation:', error); } }, [assistantEnabled, currentConversationId, conversations, queryClient, setCurrentConversationId]); // Function to clone an SEO page const cloneSeoPage = useCallback(async (seoPageId: string) => { if (!assistantEnabled) return; try { // Call the clone endpoint const response = await fetch(`${API_BASE_URL}/clone-seo-page/${seoPageId}`, { method: 'GET', credentials: 'include', }); if (!response.ok) { console.error('Failed to clone SEO page:', await response.text()); return; } const data = await response.json(); // Invalidate conversations query to fetch the new conversation await queryClient.invalidateQueries({ queryKey: ['conversations'] }); // Set the current conversation to the newly cloned one setCurrentConversationId(data.conversation_id); // If there's a URL parameter, remove it to prevent re-cloning on refresh const urlParams = new URLSearchParams(window.location.search); if (urlParams.has('continue_seo_page')) { urlParams.delete('continue_seo_page'); const newUrl = window.location.pathname + (urlParams.toString() ? '?' + urlParams.toString() : ''); window.history.replaceState({}, document.title, newUrl); } } catch (error) { console.error('Error cloning SEO page:', error); } }, [assistantEnabled, queryClient, setCurrentConversationId]); // Memoize conversations and messages const memoizedConversations = useMemo(() => rawConversations.length > 0 ? rawConversations : EMPTY_CONVERSATIONS , [rawConversations]); // Existing server conversations are not first-turn contexts. useEffect(() => { for (const conversation of memoizedConversations) { contextAttemptedForConversationRef.current.add(conversation.conversation_id); } }, [memoizedConversations]); // Memoize messages const memoizedMessages = useMemo(() => rawMessages.length > 0 ? rawMessages : EMPTY_MESSAGES , [rawMessages]); // Memoize the context value const contextValue = useMemo(() => ({ currentConversationId, setCurrentConversationId, conversations: memoizedConversations, messages: memoizedMessages, ephemeralMap, finalizedThinkingMap, sendMessage, retryError, dismissError, deleteConversation, cloneSeoPage, partialContents, setPartialContent, latestChunks, }), [ currentConversationId, setCurrentConversationId, memoizedConversations, memoizedMessages, ephemeralMap, finalizedThinkingMap, sendMessage, retryError, dismissError, deleteConversation, cloneSeoPage, partialContents, setPartialContent, latestChunks, ]); return ( {children} ); }; // ================= // // CONTEXT SELECTORS // // ================= // // Base selector hook export const useChatSelector = (selector: (state: ChatContextProps) => T): T => { return useContextSelector(ChatContext, selector); }; // Memoize individual selectors const currentConversationIdSelector = (state: ChatContextProps) => state.currentConversationId; const setCurrentConversationIdSelector = (state: ChatContextProps) => state.setCurrentConversationId; const conversationsSelector = (state: ChatContextProps) => state.conversations; const messagesSelector = (state: ChatContextProps) => state.messages; const sendMessageSelector = (state: ChatContextProps) => state.sendMessage; const retryErrorSelector = (state: ChatContextProps) => state.retryError; const dismissErrorSelector = (state: ChatContextProps) => state.dismissError; const deleteConversationSelector = (state: ChatContextProps) => state.deleteConversation; const cloneSeoPageSelector = (state: ChatContextProps) => state.cloneSeoPage; // Individual selector hooks using memoized selectors export const useCurrentConversationId = () => useChatSelector(currentConversationIdSelector); export const useSetCurrentConversationId = () => useChatSelector(setCurrentConversationIdSelector); export const useConversations = () => useChatSelector(conversationsSelector); export const useMessages = () => useChatSelector(messagesSelector); export const useSendMessage = () => useChatSelector(sendMessageSelector); export const useRetryError = () => useChatSelector(retryErrorSelector); export const useDismissError = () => useChatSelector(dismissErrorSelector); export const useDeleteConversation = () => useChatSelector(deleteConversationSelector); export const useCloneSeoPage = () => useChatSelector(cloneSeoPageSelector); // Ephemeral conversation state for only current conversation export function useCurrentConversationState() { const currentId = useContextSelector(ChatContext, ctx => ctx.currentConversationId); return useContextSelector(ChatContext, ctx => { if (!currentId) return null; return ctx.ephemeralMap[currentId] ?? DEFAULT_EPHEMERAL_STATE; }); } // Streaming message content only for the current conversation export function useCurrentPartialContent() { const currentId = useContextSelector(ChatContext, ctx => ctx.currentConversationId); return useContextSelector(ChatContext, ctx => { if (!currentId) return null; return ctx.partialContents[currentId] ?? ''; }); } export function useAssistantThinkingSummary(messageId?: string) { return useContextSelector(ChatContext, ctx => { if (!messageId) return null; return ctx.finalizedThinkingMap[messageId] ?? null; }); } // Latest streaming chunk for current conversation export function useCurrentLatestChunk() { const currentId = useContextSelector(ChatContext, ctx => ctx.currentConversationId); return useContextSelector(ChatContext, ctx => { if (!currentId) return null; return ctx.latestChunks[currentId] ?? ''; }); }