From 80414a8b61fa7cd1ad621fd7324d3fa1429ff7c7 Mon Sep 17 00:00:00 2001 From: Sarai Chinwag Date: Tue, 17 Feb 2026 01:39:53 +0000 Subject: [PATCH] refactor: use TanStack Query cache as single source of truth for chat messages (#210) --- inc/Api/Chat/Chat.php | 117 +++--------------- .../react/components/chat/ChatSidebar.jsx | 114 ++++++++++++----- .../assets/react/hooks/useChatTurn.js | 44 +++++-- 3 files changed, 134 insertions(+), 141 deletions(-) diff --git a/inc/Api/Chat/Chat.php b/inc/Api/Chat/Chat.php index f9ac9816e..d07fb6366 100644 --- a/inc/Api/Chat/Chat.php +++ b/inc/Api/Chat/Chat.php @@ -288,14 +288,10 @@ public static function handle_ping( WP_REST_Request $request ) { $prompt = sanitize_textarea_field( wp_unslash( $request->get_param( 'prompt' ) ?? '' ) ); $context = $request->get_param( 'context' ) ?? array(); - AgentContext::set( AgentType::CHAT ); - - $provider = PluginSettings::get( 'default_provider', '' ); - $model = PluginSettings::get( 'default_model', '' ); - $max_turns = PluginSettings::get( 'max_turns', 12 ); + $provider = PluginSettings::get( 'default_provider', '' ); + $model = PluginSettings::get( 'default_model', '' ); if ( empty( $provider ) || empty( $model ) ) { - AgentContext::clear(); return new WP_Error( 'provider_required', __( 'Default AI provider and model must be configured.', 'data-machine' ), @@ -329,7 +325,6 @@ public static function handle_ping( WP_REST_Request $request ) { ); if ( empty( $session_id ) ) { - AgentContext::clear(); return new WP_Error( 'session_creation_failed', __( 'Failed to create chat session.', 'data-machine' ), @@ -353,100 +348,26 @@ public static function handle_ping( WP_REST_Request $request ) { $model ); - $tool_manager = new ToolManager(); - $all_tools = $tool_manager->getAvailableToolsForChat(); - - try { - // Run FULL multi-turn loop (not single_turn) so the response is complete. - $loop = new AIConversationLoop(); - $loop_result = $loop->execute( - $messages, - $all_tools, - $provider, - $model, - AgentType::CHAT, - array( 'session_id' => $session_id ), - $max_turns, - false // multi-turn: run to completion - ); - - if ( isset( $loop_result['error'] ) ) { - $chat_db->update_session( - $session_id, - $messages, - array( - 'status' => 'error', - 'error_message' => $loop_result['error'], - 'last_activity' => current_time( 'mysql', true ), - 'message_count' => count( $messages ), - ), - $provider, - $model - ); - - do_action( - 'datamachine_log', - 'error', - 'Chat ping AI loop returned error', - array( - 'session_id' => $session_id, - 'error' => $loop_result['error'], - 'agent_type' => AgentType::CHAT, - ) - ); - - AgentContext::clear(); - return new WP_Error( - 'ping_ai_error', - $loop_result['error'], - array( 'status' => 500 ) - ); - } - } catch ( \Throwable $e ) { - do_action( - 'datamachine_log', - 'error', - 'Chat ping AI loop exception', - array( - 'session_id' => $session_id, - 'error' => $e->getMessage(), - 'agent_type' => AgentType::CHAT, - ) - ); - - $chat_db->update_session( - $session_id, - $messages, - array( - 'status' => 'error', - 'error_message' => $e->getMessage(), - 'last_activity' => current_time( 'mysql', true ), - 'message_count' => count( $messages ), - ), - $provider, - $model - ); + $result = self::executeConversationTurn( + $session_id, + $messages, + $provider, + $model, + array( 'agent_type' => AgentType::CHAT ) + ); - AgentContext::clear(); - return new WP_Error( - 'ping_error', - $e->getMessage(), - array( 'status' => 500 ) - ); - } finally { - AgentContext::clear(); + if ( is_wp_error( $result ) ) { + return $result; } - $messages = $loop_result['messages']; - $final_content = $loop_result['final_content']; - + // Update session to completed with ping source. $chat_db->update_session( $session_id, - $messages, + $result['messages'], array( 'status' => 'completed', 'last_activity' => current_time( 'mysql', true ), - 'message_count' => count( $messages ), + 'message_count' => count( $result['messages'] ), 'source' => 'ping', ), $provider, @@ -464,9 +385,9 @@ public static function handle_ping( WP_REST_Request $request ) { 'info', 'Chat ping completed', array( - 'session_id' => $session_id, - 'turns' => $loop_result['turn_count'] ?? 1, - 'agent_type' => AgentType::CHAT, + 'session_id' => $session_id, + 'turns' => $result['turn_count'], + 'agent_type' => AgentType::CHAT, ) ); @@ -475,8 +396,8 @@ public static function handle_ping( WP_REST_Request $request ) { 'success' => true, 'data' => array( 'session_id' => $session_id, - 'response' => $final_content, - 'turns' => $loop_result['turn_count'] ?? 1, + 'response' => $result['final_content'], + 'turns' => $result['turn_count'], 'completed' => true, ), ) diff --git a/inc/Core/Admin/Pages/Pipelines/assets/react/components/chat/ChatSidebar.jsx b/inc/Core/Admin/Pages/Pipelines/assets/react/components/chat/ChatSidebar.jsx index 2f6ff0b22..a900bce94 100644 --- a/inc/Core/Admin/Pages/Pipelines/assets/react/components/chat/ChatSidebar.jsx +++ b/inc/Core/Admin/Pages/Pipelines/assets/react/components/chat/ChatSidebar.jsx @@ -3,16 +3,20 @@ * * Collapsible right sidebar for chat interface. * Manages conversation state, session switching, and API interactions. - * Persists conversation across page refreshes via session storage. + * Uses TanStack Query cache as single source of truth for messages. */ /** * WordPress dependencies */ -import { useState, useCallback, useEffect, useRef } from '@wordpress/element'; +import { useState, useCallback, useRef } from '@wordpress/element'; import { Button } from '@wordpress/components'; import { close, copy } from '@wordpress/icons'; import { __ } from '@wordpress/i18n'; +/** + * External dependencies + */ +import { useQueryClient } from '@tanstack/react-query'; /** * Internal dependencies */ @@ -73,7 +77,8 @@ export default function ChatSidebar() { clearChatSession, selectedPipelineId, } = useUIStore(); - const [ messages, setMessages ] = useState( [] ); + const queryClient = useQueryClient(); + const [ pendingUserMessage, setPendingUserMessage ] = useState( null ); const [ isCopied, setIsCopied ] = useState( false ); const [ view, setView ] = useState( 'chat' ); // 'chat' | 'sessions' const chatMutation = useChatMutation(); @@ -88,18 +93,11 @@ export default function ChatSidebar() { const isCreatingSessionRef = useRef( false ); const loadingSessionRef = useRef( null ); - useEffect( () => { - if ( sessionQuery.data?.conversation ) { - setMessages( sessionQuery.data.conversation ); - } - }, [ sessionQuery.data ] ); - - useEffect( () => { - if ( sessionQuery.error?.message?.includes( 'not found' ) ) { - clearChatSession(); - setMessages( [] ); - } - }, [ sessionQuery.error, clearChatSession ] ); + // Messages from query cache, with pending message for new sessions + const cachedMessages = sessionQuery.data?.conversation ?? []; + const messages = pendingUserMessage + ? [ ...cachedMessages, pendingUserMessage ] + : cachedMessages; const handleSend = useCallback( async ( message ) => { @@ -115,7 +113,28 @@ export default function ChatSidebar() { const requestId = generateRequestId(); const userMessage = { role: 'user', content: message }; - setMessages( ( prev ) => [ ...prev, userMessage ] ); + + if ( isNewSession ) { + // No session yet — use temp pending state + setPendingUserMessage( userMessage ); + } else { + // Optimistic update to query cache + queryClient.setQueryData( + [ 'chat-session', chatSessionId ], + ( old ) => { + if ( ! old ) { + return old; + } + return { + ...old, + conversation: [ + ...( old.conversation || [] ), + userMessage, + ], + }; + } + ); + } // Track which session is loading for session-aware UI loadingSessionRef.current = chatSessionId || 'new'; @@ -129,15 +148,29 @@ export default function ChatSidebar() { requestId, } ); + const responseSessionId = response.session_id; + if ( - response.session_id && - response.session_id !== chatSessionId + responseSessionId && + responseSessionId !== chatSessionId ) { - setChatSessionId( response.session_id ); + setChatSessionId( responseSessionId ); } if ( response.conversation ) { - setMessages( response.conversation ); + // Server returned full conversation — seed the cache + queryClient.setQueryData( + [ 'chat-session', responseSessionId ], + ( old ) => ( { + ...( old || {} ), + conversation: response.conversation, + } ) + ); + } + + // Clear pending message now that session exists + if ( isNewSession ) { + setPendingUserMessage( null ); } invalidateFromToolCalls( @@ -146,14 +179,10 @@ export default function ChatSidebar() { ); // Continue processing if not complete (turn-by-turn polling) - if ( ! response.completed && response.session_id ) { + if ( ! response.completed && responseSessionId ) { await processToCompletion( - response.session_id, - ( newMessages ) => - setMessages( ( prev ) => [ - ...prev, - ...newMessages, - ] ), + responseSessionId, + queryClient, response.max_turns, selectedPipelineId ); @@ -169,7 +198,30 @@ export default function ChatSidebar() { role: 'assistant', content: errorContent, }; - setMessages( ( prev ) => [ ...prev, errorMessage ] ); + + // Clear pending message on error + if ( isNewSession ) { + setPendingUserMessage( null ); + } + + const targetSessionId = chatSessionId; + if ( targetSessionId ) { + queryClient.setQueryData( + [ 'chat-session', targetSessionId ], + ( old ) => { + if ( ! old ) { + return old; + } + return { + ...old, + conversation: [ + ...( old.conversation || [] ), + errorMessage, + ], + }; + } + ); + } if ( error.message?.includes( 'not found' ) ) { clearChatSession(); @@ -189,18 +241,20 @@ export default function ChatSidebar() { selectedPipelineId, invalidateFromToolCalls, processToCompletion, + queryClient, ] ); const handleNewConversation = useCallback( () => { clearChatSession(); - setMessages( [] ); + setPendingUserMessage( null ); setView( 'chat' ); }, [ clearChatSession ] ); const handleSelectSession = useCallback( ( sessionId ) => { setChatSessionId( sessionId ); + setPendingUserMessage( null ); setView( 'chat' ); }, [ setChatSessionId ] @@ -216,7 +270,7 @@ export default function ChatSidebar() { const handleSessionDeleted = useCallback( () => { clearChatSession(); - setMessages( [] ); + setPendingUserMessage( null ); }, [ clearChatSession ] ); const handleCopyChat = useCallback( () => { diff --git a/inc/Core/Admin/Pages/Pipelines/assets/react/hooks/useChatTurn.js b/inc/Core/Admin/Pages/Pipelines/assets/react/hooks/useChatTurn.js index 403eeb2d9..5b28dbc84 100644 --- a/inc/Core/Admin/Pages/Pipelines/assets/react/hooks/useChatTurn.js +++ b/inc/Core/Admin/Pages/Pipelines/assets/react/hooks/useChatTurn.js @@ -3,6 +3,7 @@ * * Manages turn-by-turn chat execution for async responses. * Polls /chat/continue endpoint until conversation completes. + * Updates TanStack Query cache directly instead of using callbacks. * * @since 0.12.0 */ @@ -18,6 +19,25 @@ import apiFetch from '@wordpress/api-fetch'; */ import { useChatQueryInvalidation } from './useChatQueryInvalidation'; +/** + * Update chat session cache with new messages + * + * @param {Object} queryClient TanStack Query client + * @param {string} sessionId Session ID + * @param {Array} newMessages New messages to append + */ +function appendMessagesToCache( queryClient, sessionId, newMessages ) { + queryClient.setQueryData( [ 'chat-session', sessionId ], ( old ) => { + if ( ! old ) { + return old; + } + return { + ...old, + conversation: [ ...( old.conversation || [] ), ...newMessages ], + }; + } ); +} + /** * Hook for managing turn-by-turn chat execution * @@ -32,13 +52,13 @@ export function useChatTurn() { /** * Execute a single continuation turn * - * @param {string} sessionId Session ID to continue - * @param {Function} onNewMessages Callback for new messages - * @param {number} selectedPipelineId Pipeline ID for context + * @param {string} sessionId Session ID to continue + * @param {Object} queryClient TanStack Query client + * @param {number} selectedPipelineId Pipeline ID for context * @return {Object} API response */ const continueTurn = useCallback( - async ( sessionId, onNewMessages, selectedPipelineId ) => { + async ( sessionId, queryClient, selectedPipelineId ) => { const response = await apiFetch( { path: '/datamachine/v1/chat/continue', method: 'POST', @@ -54,9 +74,7 @@ export function useChatTurn() { const data = response.data; if ( data.new_messages?.length ) { - onNewMessages( data.new_messages ); - - // Invalidate queries for any mutations that occurred + appendMessagesToCache( queryClient, sessionId, data.new_messages ); invalidateFromToolCalls( data.tool_calls, selectedPipelineId ); } @@ -68,14 +86,14 @@ export function useChatTurn() { /** * Process turns until completion or max turns reached * - * @param {string} sessionId Session ID to continue - * @param {Function} onNewMessages Callback for new messages (receives array) - * @param {number} maxTurns Maximum turns from server response - * @param {number} selectedPipelineId Pipeline ID for context + * @param {string} sessionId Session ID to continue + * @param {Object} queryClient TanStack Query client + * @param {number} maxTurns Maximum turns from server response + * @param {number} selectedPipelineId Pipeline ID for context * @return {Object} Result with completed status and turn count */ const processToCompletion = useCallback( - async ( sessionId, onNewMessages, maxTurns, selectedPipelineId ) => { + async ( sessionId, queryClient, maxTurns, selectedPipelineId ) => { setProcessingSessionId( sessionId ); setIsProcessing( true ); setTurnCount( 0 ); @@ -88,7 +106,7 @@ export function useChatTurn() { while ( ! completed && turns < effectiveMaxTurns ) { const response = await continueTurn( sessionId, - onNewMessages, + queryClient, selectedPipelineId );