@red-hat-developer-hub/backstage-plugin-lightspeed 2.5.0 → 2.6.0

package/dist/hooks/useConversationMessages.esm.js

@@ -6,6 +6,7 @@ import { TEMP_CONVERSATION_ID } from '../const.esm.js';
 import botAvatar from '../images/bot-avatar.svg';
 import userAvatar from '../images/user-avatar.svg';
 import { getConversationsData, createUserMessage, createBotMessage, transformDocumentsToSources, getTimestamp } from '../utils/lightspeed-chatbox-utils.esm.js';
+import { getSharedToolCallsCache, setSharedToolCallsCache, migrateSharedToolCallsCacheSessionPrefixToConversation, clearSharedToolCallsCacheSessionPrefix } from './toolCallsCacheStore.esm.js';
 import { useCreateConversationMessage } from './useCreateCoversationMessage.esm.js';
 
 const toolCallIdKey = (id) => {
@@ -18,6 +19,11 @@ const isMcpStyleToolCallPayload = (data) => {
 const isLegacyToolResultToken = (token) => {
   return !!token && typeof token === "object" && !Array.isArray(token) && typeof token.tool_name === "string" && token.tool_name.length > 0;
 };
+let tempToolCallsCachePrefixFallbackSeq = 0;
+function createTempToolCallsCacheSessionPrefix() {
+  const suffix = globalThis.crypto?.randomUUID?.() ?? `${Date.now()}-${++tempToolCallsCachePrefixFallbackSeq}`;
+  return `lightspeed-temp:${suffix}`;
+}
 const legacyToolResultToString = (response) => {
   if (!response) return "";
   if (typeof response === "string") return response;
@@ -49,13 +55,17 @@ const useConversationMessages = (conversationId, userName, selectedModel, select
   const streamingConversations = useRef({
     [currentConversation]: []
   });
+  const isTempStreamInProgressRef = useRef(false);
+  const [streamingConversationId, setStreamingConversationId] = useState(null);
   const pendingToolCalls = useRef({});
-  const toolCallsCache = useRef({});
   useEffect(() => {
     if (currentConversation !== conversationId) {
       setCurrentConversation(conversationId);
       setConversations((prev) => {
         if (conversationId === TEMP_CONVERSATION_ID) {
+          if (isTempStreamInProgressRef.current && (prev[TEMP_CONVERSATION_ID]?.length ?? 0) > 0) {
+            return prev;
+          }
           return { ...prev, [TEMP_CONVERSATION_ID]: [] };
         }
         if (prev[conversationId]) return prev;
@@ -95,7 +105,7 @@ const useConversationMessages = (conversationId, userName, selectedModel, select
           )
         });
         const cacheKey = `${currentConversation}-${i}`;
-        const cachedToolCalls = toolCallsCache.current[cacheKey];
+        const cachedToolCalls = getSharedToolCallsCache(cacheKey);
         if (cachedToolCalls && cachedToolCalls.length > 0) {
           botMsg.toolCalls = cachedToolCalls;
         }
@@ -106,7 +116,10 @@ const useConversationMessages = (conversationId, userName, selectedModel, select
           ...streamingConversations.current[currentConversation]
         );
       }
-      setConversations(_conversations);
+      setConversations((prev) => ({
+        ...prev,
+        ..._conversations
+      }));
     }
   }, [
     conversationsData,
@@ -118,116 +131,254 @@ const useConversationMessages = (conversationId, userName, selectedModel, select
   ]);
   const handleInputPrompt = useCallback(
     async (prompt, attachments = []) => {
-      let newConversationId = "";
-      let requestId = "";
-      const conversationTuple = [
-        createUserMessage({
-          avatar,
-          name: userName,
-          content: prompt,
-          timestamp: getTimestamp(Date.now()) ?? ""
-        }),
-        createBotMessage({
-          avatar: botAvatar,
-          isLoading: true,
-          name: selectedModel,
-          content: "",
-          timestamp: ""
-        })
-      ];
-      streamingConversations.current = {
-        ...streamingConversations.current,
-        [currentConversation]: conversationTuple
-      };
-      setConversations((prevConv) => {
-        return {
-          ...prevConv,
-          [currentConversation]: [
-            ...prevConv?.[currentConversation] ?? [],
-            ...conversationTuple
-          ]
-        };
-      });
-      setTimeout(() => {
-        scrollToBottomRef.current?.scrollToBottom();
-      }, 0);
-      const finalMessages = [];
-      let buffer = "";
+      const streamStartedOnTemp = currentConversation === TEMP_CONVERSATION_ID;
+      if (streamStartedOnTemp) {
+        isTempStreamInProgressRef.current = true;
+      }
       try {
-        const reader = await createMessage({
-          prompt,
-          selectedModel,
-          selectedProvider,
-          currentConversation,
-          attachments
+        let newConversationId = "";
+        let requestId = "";
+        setStreamingConversationId(currentConversation);
+        const toolCallsCacheKeyPrefix = currentConversation === TEMP_CONVERSATION_ID ? createTempToolCallsCacheSessionPrefix() : currentConversation;
+        const conversationTuple = [
+          createUserMessage({
+            avatar,
+            name: userName,
+            content: prompt,
+            timestamp: getTimestamp(Date.now()) ?? ""
+          }),
+          createBotMessage({
+            avatar: botAvatar,
+            isLoading: true,
+            name: selectedModel,
+            content: "",
+            timestamp: ""
+          })
+        ];
+        streamingConversations.current = {
+          ...streamingConversations.current,
+          [currentConversation]: conversationTuple
+        };
+        setConversations((prevConv) => {
+          return {
+            ...prevConv,
+            [currentConversation]: [
+              ...prevConv?.[currentConversation] ?? [],
+              ...conversationTuple
+            ]
+          };
         });
-        const decoder = new TextDecoder("utf-8");
-        let streamEnded = false;
-        while (!streamEnded) {
-          const { value, done } = await reader.read();
-          if (done) {
-            streamEnded = true;
-            break;
-          }
-          buffer += decoder.decode(value, { stream: true });
-          const parts = buffer.split("\n\n");
-          buffer = parts.pop();
-          for (const part of parts) {
-            const lines = part.split("\n").filter((line) => line.startsWith("data:"));
-            const jsonString = lines.map((line) => line.trim().slice(5).trim()).join("");
-            try {
-              const { event, data } = JSON.parse(jsonString);
-              if (event === "start") {
-                requestId = data?.request_id;
-                onRequestIdReady?.(requestId);
-                if (currentConversation === TEMP_CONVERSATION_ID) {
-                  newConversationId = data?.conversation_id;
+        setTimeout(() => {
+          scrollToBottomRef.current?.scrollToBottom();
+        }, 0);
+        const finalMessages = [];
+        let buffer = "";
+        try {
+          const reader = await createMessage({
+            prompt,
+            selectedModel,
+            selectedProvider,
+            currentConversation,
+            attachments
+          });
+          const decoder = new TextDecoder("utf-8");
+          let streamEnded = false;
+          while (!streamEnded) {
+            const { value, done } = await reader.read();
+            if (done) {
+              streamEnded = true;
+              break;
+            }
+            buffer += decoder.decode(value, { stream: true });
+            const parts = buffer.split("\n\n");
+            buffer = parts.pop();
+            for (const part of parts) {
+              const lines = part.split("\n").filter((line) => line.startsWith("data:"));
+              const jsonString = lines.map((line) => line.trim().slice(5).trim()).join("");
+              try {
+                const { event, data } = JSON.parse(jsonString);
+                if (event === "start") {
+                  requestId = data?.request_id;
+                  onRequestIdReady?.(requestId);
+                  if (currentConversation === TEMP_CONVERSATION_ID) {
+                    newConversationId = data?.conversation_id;
+                  }
                 }
-              }
-              if (event === "tool_call") {
-                const toolCallData = data?.token;
-                const legacyObjectCall = typeof toolCallData === "object" && toolCallData !== null && !Array.isArray(toolCallData) && toolCallData.tool_name;
-                const mcpStyle = isMcpStyleToolCallPayload(data);
-                const rawArgs = data?.args ?? data?.arguments;
-                const mcpArgs = rawArgs && typeof rawArgs === "object" && !Array.isArray(rawArgs) ? rawArgs : {};
-                let toolCall;
-                if (legacyObjectCall && data.id !== null) {
-                  toolCall = {
-                    id: data.id,
-                    toolName: toolCallData.tool_name,
-                    arguments: toolCallData.arguments || {},
-                    startTime: Date.now(),
-                    isLoading: true
-                  };
-                } else if (mcpStyle) {
-                  toolCall = {
-                    id: data.id,
-                    toolName: data.name.trim(),
-                    description: typeof data.type === "string" && data.type !== data.name ? data.type : void 0,
-                    arguments: mcpArgs,
-                    startTime: Date.now(),
-                    isLoading: true
-                  };
+                if (event === "tool_call") {
+                  const toolCallData = data?.token;
+                  const legacyObjectCall = typeof toolCallData === "object" && toolCallData !== null && !Array.isArray(toolCallData) && toolCallData.tool_name;
+                  const mcpStyle = isMcpStyleToolCallPayload(data);
+                  const rawArgs = data?.args ?? data?.arguments;
+                  const mcpArgs = rawArgs && typeof rawArgs === "object" && !Array.isArray(rawArgs) ? rawArgs : {};
+                  let toolCall;
+                  if (legacyObjectCall && data.id !== null) {
+                    toolCall = {
+                      id: data.id,
+                      toolName: toolCallData.tool_name,
+                      arguments: toolCallData.arguments || {},
+                      startTime: Date.now(),
+                      isLoading: true
+                    };
+                  } else if (mcpStyle) {
+                    toolCall = {
+                      id: data.id,
+                      toolName: data.name.trim(),
+                      description: typeof data.type === "string" && data.type !== data.name ? data.type : void 0,
+                      arguments: mcpArgs,
+                      startTime: Date.now(),
+                      isLoading: true
+                    };
+                  }
+                  if (toolCall && data.id !== null) {
+                    const newToolCall = toolCall;
+                    pendingToolCalls.current[toolCallIdKey(data.id)] = newToolCall;
+                    setConversations((prevConversations) => {
+                      const conversation = prevConversations[currentConversation] ?? [];
+                      const lastMessageIndex = conversation.length - 1;
+                      if (lastMessageIndex < 0) return prevConversations;
+                      const lastMessage = { ...conversation[lastMessageIndex] };
+                      const existingToolCalls = normalizeToolCalls(
+                        lastMessage.toolCalls
+                      );
+                      const nextToolCalls = [
+                        ...existingToolCalls,
+                        newToolCall
+                      ];
+                      lastMessage.toolCalls = nextToolCalls;
+                      const messageIndex = Math.floor(lastMessageIndex / 2);
+                      const cacheKey = `${toolCallsCacheKeyPrefix}-${messageIndex}`;
+                      setSharedToolCallsCache(cacheKey, nextToolCalls);
+                      const updatedConversation = [
+                        ...conversation.slice(0, lastMessageIndex),
+                        lastMessage
+                      ];
+                      return {
+                        ...prevConversations,
+                        [currentConversation]: updatedConversation
+                      };
+                    });
+                    const [humanMessage, aiMessage] = streamingConversations.current[currentConversation] || [];
+                    if (aiMessage) {
+                      const existingStreamingToolCalls = normalizeToolCalls(
+                        aiMessage.toolCalls
+                      );
+                      streamingConversations.current[currentConversation] = [
+                        humanMessage,
+                        {
+                          ...aiMessage,
+                          toolCalls: [
+                            ...existingStreamingToolCalls,
+                            newToolCall
+                          ]
+                        }
+                      ];
+                    }
+                  }
+                }
+                if (event === "tool_result") {
+                  const tokenResult = data?.token;
+                  const legacyResult = isLegacyToolResultToken(tokenResult);
+                  const mcpHasContent = data?.id !== null && data.content !== void 0 && !legacyResult;
+                  let responsePayload;
+                  let matchToolName;
+                  let toolIdKey;
+                  if (legacyResult) {
+                    responsePayload = legacyToolResultToString(
+                      tokenResult.response
+                    );
+                    matchToolName = tokenResult.tool_name;
+                    toolIdKey = data?.id !== null ? toolCallIdKey(data.id) : void 0;
+                  } else if (mcpHasContent) {
+                    toolIdKey = toolCallIdKey(data.id);
+                    responsePayload = typeof data.content === "string" ? data.content : JSON.stringify(data.content);
+                    if (typeof data.status === "string" && data.status !== "success") {
+                      responsePayload = `[${data.status}] ${responsePayload}`;
+                    }
+                  }
+                  if (responsePayload !== void 0 && toolIdKey !== void 0) {
+                    const pendingCall = pendingToolCalls.current[toolIdKey];
+                    const endTime = Date.now();
+                    const executionTime = pendingCall ? (endTime - pendingCall.startTime) / 1e3 : 0;
+                    setConversations((prevConversations) => {
+                      const conversation = prevConversations[currentConversation] ?? [];
+                      const lastMessageIndex = conversation.length - 1;
+                      if (lastMessageIndex < 0) return prevConversations;
+                      const lastMessage = { ...conversation[lastMessageIndex] };
+                      const toolCalls = lastMessage.toolCalls || [];
+                      const updatedToolCalls = toolCalls.map((tc) => {
+                        const idMatches = toolCallIdKey(tc.id) === toolIdKey || matchToolName !== void 0 && tc.toolName === matchToolName;
+                        if (idMatches) {
+                          return {
+                            ...tc,
+                            response: responsePayload,
+                            endTime,
+                            executionTime,
+                            isLoading: false
+                          };
+                        }
+                        return tc;
+                      });
+                      lastMessage.toolCalls = updatedToolCalls;
+                      const messageIndex = Math.floor(lastMessageIndex / 2);
+                      const cacheKey = `${toolCallsCacheKeyPrefix}-${messageIndex}`;
+                      setSharedToolCallsCache(cacheKey, updatedToolCalls);
+                      const updatedConversation = [
+                        ...conversation.slice(0, lastMessageIndex),
+                        lastMessage
+                      ];
+                      return {
+                        ...prevConversations,
+                        [currentConversation]: updatedConversation
+                      };
+                    });
+                    const [humanMessage, aiMessage] = streamingConversations.current[currentConversation] || [];
+                    if (aiMessage) {
+                      const toolCalls = aiMessage.toolCalls || [];
+                      const updatedToolCalls = toolCalls.map((tc) => {
+                        const idMatches = toolCallIdKey(tc.id) === toolIdKey || matchToolName !== void 0 && tc.toolName === matchToolName;
+                        if (idMatches) {
+                          return {
+                            ...tc,
+                            response: responsePayload,
+                            endTime,
+                            executionTime,
+                            isLoading: false
+                          };
+                        }
+                        return tc;
+                      });
+                      streamingConversations.current[currentConversation] = [
+                        humanMessage,
+                        { ...aiMessage, toolCalls: updatedToolCalls }
+                      ];
+                    }
+                    delete pendingToolCalls.current[toolIdKey];
+                  }
                 }
-                if (toolCall && data.id !== null) {
-                  const newToolCall = toolCall;
-                  pendingToolCalls.current[toolCallIdKey(data.id)] = newToolCall;
+                if (event === "token") {
+                  const content = data?.token || "";
+                  finalMessages.push(content);
+                  const [humanMessage, aiMessage] = streamingConversations.current[currentConversation];
+                  streamingConversations.current[currentConversation] = [
+                    humanMessage,
+                    { ...aiMessage, content: aiMessage.content + content }
+                  ];
                   setConversations((prevConversations) => {
                     const conversation = prevConversations[currentConversation] ?? [];
                     const lastMessageIndex = conversation.length - 1;
-                    if (lastMessageIndex < 0) return prevConversations;
-                    const lastMessage = { ...conversation[lastMessageIndex] };
-                    const existingToolCalls = normalizeToolCalls(
-                      lastMessage.toolCalls
+                    const lastMessage = conversation.length === 0 ? createBotMessage({
+                      content: "",
+                      timestamp: getTimestamp(Date.now())
+                    }) : { ...conversation[lastMessageIndex] };
+                    if ((lastMessage?.content ?? "").trim().length > 0) {
+                      lastMessage.isLoading = false;
+                    }
+                    lastMessage.content += content;
+                    lastMessage.name = data?.response_metadata?.model || selectedModel;
+                    lastMessage.timestamp = getTimestamp(
+                      // TODO: To be fixed in the query response
+                      data?.response_metadata?.created_at || Date.now()
                     );
-                    const nextToolCalls = [
-                      ...existingToolCalls,
-                      newToolCall
-                    ];
-                    lastMessage.toolCalls = nextToolCalls;
-                    const messageIndex = Math.floor(lastMessageIndex / 2);
-                    const cacheKey = `${currentConversation}-${messageIndex}`;
-                    toolCallsCache.current[cacheKey] = nextToolCalls;
                     const updatedConversation = [
                       ...conversation.slice(0, lastMessageIndex),
                       lastMessage
@@ -237,68 +388,22 @@ const useConversationMessages = (conversationId, userName, selectedModel, select
                       [currentConversation]: updatedConversation
                     };
                   });
-                  const [humanMessage, aiMessage] = streamingConversations.current[currentConversation] || [];
-                  if (aiMessage) {
-                    const existingStreamingToolCalls = normalizeToolCalls(
-                      aiMessage.toolCalls
-                    );
-                    streamingConversations.current[currentConversation] = [
-                      humanMessage,
-                      {
-                        ...aiMessage,
-                        toolCalls: [...existingStreamingToolCalls, newToolCall]
-                      }
-                    ];
-                  }
                 }
-              }
-              if (event === "tool_result") {
-                const tokenResult = data?.token;
-                const legacyResult = isLegacyToolResultToken(tokenResult);
-                const mcpHasContent = data?.id !== null && data.content !== void 0 && !legacyResult;
-                let responsePayload;
-                let matchToolName;
-                let toolIdKey;
-                if (legacyResult) {
-                  responsePayload = legacyToolResultToString(
-                    tokenResult.response
-                  );
-                  matchToolName = tokenResult.tool_name;
-                  toolIdKey = data?.id !== null ? toolCallIdKey(data.id) : void 0;
-                } else if (mcpHasContent) {
-                  toolIdKey = toolCallIdKey(data.id);
-                  responsePayload = typeof data.content === "string" ? data.content : JSON.stringify(data.content);
-                  if (typeof data.status === "string" && data.status !== "success") {
-                    responsePayload = `[${data.status}] ${responsePayload}`;
+                if (event === "interrupted") {
+                  if (currentConversation === TEMP_CONVERSATION_ID && data?.conversation_id) {
+                    newConversationId = data.conversation_id;
                   }
-                }
-                if (responsePayload !== void 0 && toolIdKey !== void 0) {
-                  const pendingCall = pendingToolCalls.current[toolIdKey];
-                  const endTime = Date.now();
-                  const executionTime = pendingCall ? (endTime - pendingCall.startTime) / 1e3 : 0;
                   setConversations((prevConversations) => {
                     const conversation = prevConversations[currentConversation] ?? [];
                     const lastMessageIndex = conversation.length - 1;
-                    if (lastMessageIndex < 0) return prevConversations;
-                    const lastMessage = { ...conversation[lastMessageIndex] };
-                    const toolCalls = lastMessage.toolCalls || [];
-                    const updatedToolCalls = toolCalls.map((tc) => {
-                      const idMatches = toolCallIdKey(tc.id) === toolIdKey || matchToolName !== void 0 && tc.toolName === matchToolName;
-                      if (idMatches) {
-                        return {
-                          ...tc,
-                          response: responsePayload,
-                          endTime,
-                          executionTime,
-                          isLoading: false
-                        };
-                      }
-                      return tc;
-                    });
-                    lastMessage.toolCalls = updatedToolCalls;
-                    const messageIndex = Math.floor(lastMessageIndex / 2);
-                    const cacheKey = `${currentConversation}-${messageIndex}`;
-                    toolCallsCache.current[cacheKey] = updatedToolCalls;
+                    const lastMessage = conversation.length === 0 ? createBotMessage({
+                      content: "",
+                      isLoading: false,
+                      timestamp: getTimestamp(Date.now())
+                    }) : {
+                      ...conversation[lastMessageIndex],
+                      isLoading: false
+                    };
                     const updatedConversation = [
                       ...conversation.slice(0, lastMessageIndex),
                       lastMessage
@@ -308,177 +413,105 @@ const useConversationMessages = (conversationId, userName, selectedModel, select
                       [currentConversation]: updatedConversation
                     };
                   });
-                  const [humanMessage, aiMessage] = streamingConversations.current[currentConversation] || [];
-                  if (aiMessage) {
-                    const toolCalls = aiMessage.toolCalls || [];
-                    const updatedToolCalls = toolCalls.map((tc) => {
-                      const idMatches = toolCallIdKey(tc.id) === toolIdKey || matchToolName !== void 0 && tc.toolName === matchToolName;
-                      if (idMatches) {
-                        return {
-                          ...tc,
-                          response: responsePayload,
-                          endTime,
-                          executionTime,
-                          isLoading: false
-                        };
-                      }
-                      return tc;
-                    });
-                    streamingConversations.current[currentConversation] = [
-                      humanMessage,
-                      { ...aiMessage, toolCalls: updatedToolCalls }
+                  streamEnded = true;
+                  break;
+                }
+                if (event === "end") {
+                  const documents = data?.referenced_documents || [];
+                  setConversations((prevConversations) => {
+                    const conversation = prevConversations[currentConversation] ?? [];
+                    const lastMessageIndex = conversation.length - 1;
+                    const lastMessage = conversation.length === 0 ? createBotMessage({
+                      content: "",
+                      isLoading: false,
+                      timestamp: getTimestamp(Date.now())
+                    }) : {
+                      ...conversation[lastMessageIndex],
+                      isLoading: false
+                    };
+                    if (documents.length) {
+                      lastMessage.sources = {
+                        sources: documents.map((doc) => ({
+                          title: doc.doc_title,
+                          link: doc.doc_url,
+                          body: doc.doc_description
+                        }))
+                      };
+                    }
+                    const updatedConversation = [
+                      ...conversation.slice(0, lastMessageIndex),
+                      lastMessage
                     ];
-                  }
-                  delete pendingToolCalls.current[toolIdKey];
+                    return {
+                      ...prevConversations,
+                      [currentConversation]: updatedConversation
+                    };
+                  });
                 }
-              }
-              if (event === "token") {
-                const content = data?.token || "";
-                finalMessages.push(content);
-                const [humanMessage, aiMessage] = streamingConversations.current[currentConversation];
-                streamingConversations.current[currentConversation] = [
-                  humanMessage,
-                  { ...aiMessage, content: aiMessage.content + content }
-                ];
-                setConversations((prevConversations) => {
-                  const conversation = prevConversations[currentConversation] ?? [];
-                  const lastMessageIndex = conversation.length - 1;
-                  const lastMessage = conversation.length === 0 ? createBotMessage({
-                    content: "",
-                    timestamp: getTimestamp(Date.now())
-                  }) : { ...conversation[lastMessageIndex] };
-                  if ((lastMessage?.content ?? "").trim().length > 0) {
-                    lastMessage.isLoading = false;
-                  }
-                  lastMessage.content += content;
-                  lastMessage.name = data?.response_metadata?.model || selectedModel;
-                  lastMessage.timestamp = getTimestamp(
-                    // TODO: To be fixed in the query response
-                    data?.response_metadata?.created_at || Date.now()
-                  );
-                  const updatedConversation = [
-                    ...conversation.slice(0, lastMessageIndex),
-                    lastMessage
-                  ];
-                  return {
-                    ...prevConversations,
-                    [currentConversation]: updatedConversation
-                  };
-                });
-              }
-              if (event === "interrupted") {
-                if (currentConversation === TEMP_CONVERSATION_ID && data?.conversation_id) {
-                  newConversationId = data.conversation_id;
+              } catch (error) {
+                console.warn("Error parsing JSON:", error);
+                if (typeof onComplete === "function") {
+                  onComplete("Invalid JSON received");
                 }
-                setConversations((prevConversations) => {
-                  const conversation = prevConversations[currentConversation] ?? [];
-                  const lastMessageIndex = conversation.length - 1;
-                  const lastMessage = conversation.length === 0 ? createBotMessage({
-                    content: "",
-                    isLoading: false,
-                    timestamp: getTimestamp(Date.now())
-                  }) : { ...conversation[lastMessageIndex], isLoading: false };
-                  const updatedConversation = [
-                    ...conversation.slice(0, lastMessageIndex),
-                    lastMessage
-                  ];
-                  return {
-                    ...prevConversations,
-                    [currentConversation]: updatedConversation
-                  };
-                });
-                streamEnded = true;
-                break;
-              }
-              if (event === "end") {
-                const documents = data?.referenced_documents || [];
-                setConversations((prevConversations) => {
-                  const conversation = prevConversations[currentConversation] ?? [];
-                  const lastMessageIndex = conversation.length - 1;
-                  const lastMessage = conversation.length === 0 ? createBotMessage({
-                    content: "",
-                    isLoading: false,
-                    timestamp: getTimestamp(Date.now())
-                  }) : { ...conversation[lastMessageIndex], isLoading: false };
-                  if (documents.length) {
-                    lastMessage.sources = {
-                      sources: documents.map((doc) => ({
-                        title: doc.doc_title,
-                        link: doc.doc_url,
-                        body: doc.doc_description
-                      }))
-                    };
-                  }
-                  const updatedConversation = [
-                    ...conversation.slice(0, lastMessageIndex),
-                    lastMessage
-                  ];
-                  return {
-                    ...prevConversations,
-                    [currentConversation]: updatedConversation
-                  };
-                });
-              }
-            } catch (error) {
-              console.warn("Error parsing JSON:", error);
-              if (typeof onComplete === "function") {
-                onComplete("Invalid JSON received");
               }
             }
+            if (streamEnded) break;
           }
-          if (streamEnded) break;
+        } catch (e) {
+          setConversations((prevConversations) => {
+            const conversation = prevConversations[currentConversation] ?? [];
+            const lastMessageIndex = conversation.length - 1;
+            const lastMessage = conversation.length === 0 ? createBotMessage({
+              content: "",
+              timestamp: getTimestamp(Date.now())
+            }) : { ...conversation[lastMessageIndex] };
+            lastMessage.isLoading = false;
+            lastMessage.content += e;
+            lastMessage.error = {
+              title: e.message
+            };
+            lastMessage.timestamp = getTimestamp(Date.now());
+            const updatedConversation = [
+              ...conversation.slice(0, lastMessageIndex),
+              lastMessage
+            ];
+            finalMessages.push(`${e}`);
+            return {
+              ...prevConversations,
+              [newConversationId.length > 0 ? newConversationId : currentConversation]: updatedConversation
+            };
+          });
         }
-      } catch (e) {
-        setConversations((prevConversations) => {
-          const conversation = prevConversations[currentConversation] ?? [];
-          const lastMessageIndex = conversation.length - 1;
-          const lastMessage = conversation.length === 0 ? createBotMessage({
-            content: "",
-            timestamp: getTimestamp(Date.now())
-          }) : { ...conversation[lastMessageIndex] };
-          lastMessage.isLoading = false;
-          lastMessage.content += e;
-          lastMessage.error = {
-            title: e.message
-          };
-          lastMessage.timestamp = getTimestamp(Date.now());
-          const updatedConversation = [
-            ...conversation.slice(0, lastMessageIndex),
-            lastMessage
-          ];
-          finalMessages.push(`${e}`);
-          return {
-            ...prevConversations,
-            [newConversationId.length > 0 ? newConversationId : currentConversation]: updatedConversation
-          };
-        });
-      }
-      streamingConversations.current[currentConversation] = [];
-      if (typeof onComplete === "function") {
-        onComplete(finalMessages.join(""));
-      }
-      if (currentConversation === TEMP_CONVERSATION_ID && newConversationId) {
-        Object.keys(toolCallsCache.current).forEach((key) => {
-          if (key.startsWith(`${TEMP_CONVERSATION_ID}-`)) {
-            const messageIndex = key.replace(`${TEMP_CONVERSATION_ID}-`, "");
-            const newKey = `${newConversationId}-${messageIndex}`;
-            toolCallsCache.current[newKey] = toolCallsCache.current[key];
-            delete toolCallsCache.current[key];
-          }
-        });
-        setConversations((prevConversations) => {
-          return {
-            ...prevConversations,
-            [newConversationId]: prevConversations[TEMP_CONVERSATION_ID]
-          };
-        });
-        onStart?.(newConversationId);
-        setTimeout(() => {
-          setConversations((prev) => {
-            const { [TEMP_CONVERSATION_ID]: _, ...rest } = prev;
-            return rest;
+        streamingConversations.current[currentConversation] = [];
+        if (typeof onComplete === "function") {
+          onComplete(finalMessages.join(""));
+        }
+        if (currentConversation === TEMP_CONVERSATION_ID && newConversationId) {
+          migrateSharedToolCallsCacheSessionPrefixToConversation(
+            toolCallsCacheKeyPrefix,
+            newConversationId
+          );
+          setConversations((prevConversations) => {
+            return {
+              ...prevConversations,
+              [newConversationId]: prevConversations[TEMP_CONVERSATION_ID]
+            };
           });
-        }, 0);
+          onStart?.(newConversationId);
+          setTimeout(() => {
+            setConversations((prev) => {
+              const { [TEMP_CONVERSATION_ID]: _, ...rest } = prev;
+              return rest;
+            });
+          }, 0);
+        } else if (currentConversation === TEMP_CONVERSATION_ID) {
+          clearSharedToolCallsCacheSessionPrefix(toolCallsCacheKeyPrefix);
+        }
+      } finally {
+        if (streamStartedOnTemp) {
+          isTempStreamInProgressRef.current = false;
+        }
+        setStreamingConversationId(null);
       }
     },
     [
@@ -498,6 +531,7 @@ const useConversationMessages = (conversationId, userName, selectedModel, select
     handleInputPrompt,
     conversations,
     scrollToBottomRef,
+    streamingConversationId,
     ...queryProps
   };
 };