Merge pull request #6726 from remix-project-org/fix_tool_cb

more display time

Author: STetsing

libs/remix-ai-core/src/helpers/streamHandler.ts

@@ -147,8 +147,11 @@ export const HandleOpenAIResponse = async (aiResponse: IAIStreamResponse | any,
             const toolCallsArray = Array.from(toolCalls.values());
             const response = await tool_callback(toolCallsArray, uiToolCallback)
 
+            // Preserve the uiToolCallback from the response if it exists (from subsequent calls)
             if (response && typeof response === 'object') {
-              response.uiToolCallback = uiToolCallback;
+              if (!response.uiToolCallback && uiToolCallback) {
+                response.uiToolCallback = uiToolCallback;
+              }
             }
             cb("\n\n");
             HandleOpenAIResponse(response, cb, done_cb)
@@ -244,8 +247,11 @@ export const HandleMistralAIResponse = async (aiResponse: IAIStreamResponse | an
             const toolCalls = json.choices[0].delta.tool_calls;
             const response = await tool_callback(toolCalls, uiToolCallback)
 
+            // Preserve the uiToolCallback from the response if it exists (from subsequent calls)
             if (response && typeof response === 'object') {
-              response.uiToolCallback = uiToolCallback;
+              if (!response.uiToolCallback && uiToolCallback) {
+                response.uiToolCallback = uiToolCallback;
+              }
             }
             HandleMistralAIResponse(response, cb, done_cb)
           } else if (json.choices[0].delta.content){
@@ -346,12 +352,11 @@ export const HandleAnthropicResponse = async (aiResponse: IAIStreamResponse | an
             }));
 
             if (toolCalls.length > 0) {
-              uiToolCallback?.(true);
-              const response = await tool_callback(toolCalls)
-              uiToolCallback?.(false);
-              // Keep the callback attached for recursive calls
+              const response = await tool_callback(toolCalls, uiToolCallback)
               if (response && typeof response === 'object') {
-                response.uiToolCallback = uiToolCallback;
+                if (!response.uiToolCallback && uiToolCallback) {
+                  response.uiToolCallback = uiToolCallback;
+                }
               }
               cb("\n\n");
               HandleAnthropicResponse(response, cb, done_cb)
@@ -417,8 +422,11 @@ export const HandleOllamaResponse = async (aiResponse: IAIStreamResponse | any,
             const toolCalls = parsed.message.tool_calls;
             const response = await tool_callback(toolCalls, uiToolCallback)
             // Keep the callback attached for recursive calls
+            // Preserve the uiToolCallback from the response if it exists (from subsequent calls)
             if (response && typeof response === 'object') {
-              response.uiToolCallback = uiToolCallback;
+              if (!response.uiToolCallback && uiToolCallback) {
+                response.uiToolCallback = uiToolCallback;
+              }
             }
             cb("\n\n");
             HandleOllamaResponse(response, cb, done_cb, reasoning_cb)

libs/remix-ai-core/src/inferencers/mcp/mcpInferencer.ts

@@ -506,20 +506,23 @@ export class MCPInferencer extends RemoteInferencer implements ICompletions, IGe
 
             const followUpOptions = {
               ...enhancedOptions,
-              toolsMessages: toolsMessagesArray
+              toolsMessages: toolsMessagesArray,
+              chatHistory: options.provider === 'anthropic'
+                ? [...(enhancedOptions.chatHistory || []), { role: 'user', content: prompt }]
+                : enhancedOptions.chatHistory
             };
 
-            // Send empty prompt - the tool results are in toolsMessages
-            // Don't add extra prompts as they cause Anthropic to summarize instead of using full tool results
             if (options.provider === 'openai' || options.provider === 'mistralai') {
               return {
                 streamResponse: await this.baseInferencer.answer(prompt, followUpOptions),
-                callback: toolExecutionStatusCallback
+                callback: toolExecutionStatusCallback,
+                uiToolCallback: uiCallback
               } as IAIStreamResponse;
             } else {
               return {
                 streamResponse: await this.baseInferencer.answer("", followUpOptions),
-                callback: toolExecutionStatusCallback
+                callback: toolExecutionStatusCallback,
+                uiToolCallback: uiCallback
               } as IAIStreamResponse;
             }
           }
@@ -832,7 +835,6 @@ ${toolsList}`,
           }
 
           if (uiCallback){
-            console.log('on UI tool callback', innerToolCall.name, innerToolCall.arguments)
             uiCallback(true, innerToolCall.name, innerToolCall.arguments);
           }
           const result = await this.executeTool(targetServer, innerToolCall);

libs/remix-ai-core/src/inferencers/remote/remoteInference.ts

@@ -188,7 +188,9 @@ export class RemoteInferencer implements ICompletions, IGeneration {
   }
 
   async answer(prompt, options:IParams=GenerationParams): Promise<any> {
-    options.chatHistory = buildChatPrompt()
+    if (!options.toolsMessages) {
+      options.chatHistory = buildChatPrompt()
+    }
     const payload = { 'prompt': prompt, "endpoint":"answer", ...options }
     if (options.stream_result) return this._streamInferenceRequest(payload, AIRequestType.GENERAL)
     else return this._makeRequest(payload, AIRequestType.GENERAL)

libs/remix-ui/remix-ai-assistant/src/components/remix-ui-remix-ai-assistant.tsx

@@ -466,7 +466,7 @@ export const RemixUiRemixAiAssistant = React.forwardRef<
         let toolExecutionStartTime: number | null = null
 
         const uiToolCallback = (isExecuting: boolean, toolName?: string, toolArgs?: Record<string, any>) => {
-          const MIN_DISPLAY_TIME = 2000 // 2 seconds
+          const MIN_DISPLAY_TIME = 30000 // 30 seconds
 
           // Clear any pending timeout
           if (clearToolTimeout) {
@@ -475,16 +475,17 @@ export const RemixUiRemixAiAssistant = React.forwardRef<
           }
 
           if (isExecuting) {
-            // Tool execution starting or updating - show immediately
             if (!toolExecutionStartTime) {
               toolExecutionStartTime = Date.now()
             }
+
             setMessages(prev =>
               prev.map(m => (m.id === assistantId ? {
                 ...m,
-                isExecutingTools: isExecuting,
-                executingToolName: toolName,
-                executingToolArgs: toolArgs
+                // Only show tool execution indicator if no content has arrived yet
+                isExecutingTools: m.content.length === 0 ? isExecuting : m.isExecutingTools,
+                executingToolName: m.content.length === 0 ? toolName : m.executingToolName,
+                executingToolArgs: m.content.length === 0 ? toolArgs : m.executingToolArgs
               } : m))
             )
           } else {
@@ -601,7 +602,8 @@ export const RemixUiRemixAiAssistant = React.forwardRef<
             }
           )
         }
-        setIsStreaming(false)
+        // Note: setIsStreaming(false) is called in each handler's completion callback
+        // DO NOT call it here as it would stop the spinner before the response completes
       }
       catch (error) {
         console.error('Error sending prompt:', error)