Skip to content

refactor: use TanStack Query cache as single source of truth for chat messages #218

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chubes4 merged 1 commit into from
Feb 17, 2026
Merged

refactor: use TanStack Query cache as single source of truth for chat messages #218

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 19 additions & 98 deletions inc/Api/Chat/Chat.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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' ),
Expand Down Expand Up @@ -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' ),
Expand All @@ -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,
Expand All @@ -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,
)
);

Expand All @@ -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,
),
)
Expand Down
114 changes: 84 additions & 30 deletions inc/Core/Admin/Pages/Pipelines/assets/react/components/chat/ChatSidebar.jsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
*/
Expand Down Expand Up @@ -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();
Expand All @@ -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 ) => {
Expand All @@ -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';
Expand All @@ -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(
Expand All @@ -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
);
Expand All @@ -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();
Expand All @@ -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 ]
Expand All @@ -216,7 +270,7 @@ export default function ChatSidebar() {

const handleSessionDeleted = useCallback( () => {
clearChatSession();
setMessages( [] );
setPendingUserMessage( null );
}, [ clearChatSession ] );

const handleCopyChat = useCallback( () => {
Expand Down
Loading