@copilotkit/react-core 1.9.2-next.2 → 1.9.2-next.20

package/dist/hooks/use-chat.js

@@ -83,6 +83,7 @@ __export(use_chat_exports, {
 });
 module.exports = __toCommonJS(use_chat_exports);
 var import_react5 = require("react");
+var import_react_dom = require("react-dom");
 var import_shared4 = require("@copilotkit/shared");
 var import_runtime_client_gql3 = require("@copilotkit/runtime-client-gql");
 
@@ -111,16 +112,104 @@ function processActionsForRuntimeRequest(actions) {
   return filteredActions;
 }
 
+// src/context/copilot-context.tsx
+var import_react = __toESM(require("react"));
+var emptyCopilotContext = {
+  actions: {},
+  setAction: () => {
+  },
+  removeAction: () => {
+  },
+  coAgentStateRenders: {},
+  setCoAgentStateRender: () => {
+  },
+  removeCoAgentStateRender: () => {
+  },
+  chatComponentsCache: { current: { actions: {}, coAgentStateRenders: {} } },
+  getContextString: (documents, categories) => returnAndThrowInDebug(""),
+  addContext: () => "",
+  removeContext: () => {
+  },
+  getAllContext: () => [],
+  getFunctionCallHandler: () => returnAndThrowInDebug(() => __async(void 0, null, function* () {
+  })),
+  isLoading: false,
+  setIsLoading: () => returnAndThrowInDebug(false),
+  chatInstructions: "",
+  setChatInstructions: () => returnAndThrowInDebug(""),
+  additionalInstructions: [],
+  setAdditionalInstructions: () => returnAndThrowInDebug([]),
+  getDocumentsContext: (categories) => returnAndThrowInDebug([]),
+  addDocumentContext: () => returnAndThrowInDebug(""),
+  removeDocumentContext: () => {
+  },
+  runtimeClient: {},
+  copilotApiConfig: new class {
+    get chatApiEndpoint() {
+      throw new Error("Remember to wrap your app in a `<CopilotKit> {...} </CopilotKit>` !!!");
+    }
+    get headers() {
+      return {};
+    }
+    get body() {
+      return {};
+    }
+  }(),
+  chatSuggestionConfiguration: {},
+  addChatSuggestionConfiguration: () => {
+  },
+  removeChatSuggestionConfiguration: () => {
+  },
+  showDevConsole: false,
+  coagentStates: {},
+  setCoagentStates: () => {
+  },
+  coagentStatesRef: { current: {} },
+  setCoagentStatesWithRef: () => {
+  },
+  agentSession: null,
+  setAgentSession: () => {
+  },
+  forwardedParameters: {},
+  agentLock: null,
+  threadId: "",
+  setThreadId: () => {
+  },
+  runId: null,
+  setRunId: () => {
+  },
+  chatAbortControllerRef: { current: null },
+  availableAgents: [],
+  extensions: {},
+  setExtensions: () => {
+  },
+  langGraphInterruptAction: null,
+  setLangGraphInterruptAction: () => null,
+  removeLangGraphInterruptAction: () => null,
+  onError: void 0
+};
+var CopilotContext = import_react.default.createContext(emptyCopilotContext);
+function useCopilotContext() {
+  const context = import_react.default.useContext(CopilotContext);
+  if (context === emptyCopilotContext) {
+    throw new Error("Remember to wrap your app in a `<CopilotKit> {...} </CopilotKit>` !!!");
+  }
+  return context;
+}
+function returnAndThrowInDebug(_value) {
+  throw new Error("Remember to wrap your app in a `<CopilotKit> {...} </CopilotKit>` !!!");
+}
+
 // src/hooks/use-copilot-runtime-client.ts
 var import_runtime_client_gql2 = require("@copilotkit/runtime-client-gql");
 
 // src/components/toast/toast-provider.tsx
-var import_react = require("react");
+var import_react2 = require("react");
 var import_shared2 = require("@copilotkit/shared");
 var import_jsx_runtime = require("react/jsx-runtime");
-var ToastContext = (0, import_react.createContext)(void 0);
+var ToastContext = (0, import_react2.createContext)(void 0);
 function useToast() {
-  const context = (0, import_react.useContext)(ToastContext);
+  const context = (0, import_react2.useContext)(ToastContext);
   if (!context) {
     throw new Error("useToast must be used within a ToastProvider");
   }
@@ -128,7 +217,7 @@ function useToast() {
 }
 
 // src/hooks/use-copilot-runtime-client.ts
-var import_react2 = require("react");
+var import_react3 = require("react");
 var import_shared3 = require("@copilotkit/shared");
 
 // src/utils/dev-console.ts
@@ -139,9 +228,36 @@ function shouldShowDevConsole(showDevConsole) {
 // src/hooks/use-copilot-runtime-client.ts
 var useCopilotRuntimeClient = (options) => {
   const { setBannerError } = useToast();
-  const _a = options, { showDevConsole } = _a, runtimeOptions = __objRest(_a, ["showDevConsole"]);
-  const lastStructuredErrorRef = (0, import_react2.useRef)(null);
-  const runtimeClient = (0, import_react2.useMemo)(() => {
+  const _a = options, { showDevConsole, onError } = _a, runtimeOptions = __objRest(_a, ["showDevConsole", "onError"]);
+  const lastStructuredErrorRef = (0, import_react3.useRef)(null);
+  const traceUIError = (error, originalError) => __async(void 0, null, function* () {
+    if (!onError || !runtimeOptions.publicApiKey)
+      return;
+    try {
+      const errorEvent = {
+        type: "error",
+        timestamp: Date.now(),
+        context: {
+          source: "ui",
+          request: {
+            operation: "runtimeClient",
+            url: runtimeOptions.url,
+            startTime: Date.now()
+          },
+          technical: {
+            environment: "browser",
+            userAgent: typeof navigator !== "undefined" ? navigator.userAgent : void 0,
+            stackTrace: originalError instanceof Error ? originalError.stack : void 0
+          }
+        },
+        error
+      };
+      yield onError(errorEvent);
+    } catch (error2) {
+      console.error("Error in onError handler:", error2);
+    }
+  });
+  const runtimeClient = (0, import_react3.useMemo)(() => {
     return new import_runtime_client_gql2.CopilotRuntimeClient(__spreadProps(__spreadValues({}, runtimeOptions), {
       handleGQLErrors: (error) => {
         var _a2;
@@ -168,12 +284,14 @@ var useCopilotRuntimeClient = (options) => {
             const ckError = createStructuredError(gqlError);
             if (ckError) {
               setBannerError(ckError);
+              traceUIError(ckError, gqlError);
             } else {
               const fallbackError = new import_shared3.CopilotKitError({
                 message: gqlError.message,
                 code: import_shared3.CopilotKitErrorCode.UNKNOWN
               });
               setBannerError(fallbackError);
+              traceUIError(fallbackError, gqlError);
             }
           };
           graphQLErrors.forEach(routeError);
@@ -187,6 +305,7 @@ var useCopilotRuntimeClient = (options) => {
               code: import_shared3.CopilotKitErrorCode.UNKNOWN
             });
             setBannerError(fallbackError);
+            traceUIError(fallbackError, error);
           }
         }
       },
@@ -199,7 +318,7 @@ var useCopilotRuntimeClient = (options) => {
         setBannerError(warningError);
       }
     }));
-  }, [runtimeOptions, setBannerError, showDevConsole]);
+  }, [runtimeOptions, setBannerError, showDevConsole, onError]);
   return runtimeClient;
 };
 function createStructuredError(gqlError) {
@@ -226,93 +345,6 @@ function createStructuredError(gqlError) {
   return null;
 }
 
-// src/context/copilot-context.tsx
-var import_react3 = __toESM(require("react"));
-var emptyCopilotContext = {
-  actions: {},
-  setAction: () => {
-  },
-  removeAction: () => {
-  },
-  coAgentStateRenders: {},
-  setCoAgentStateRender: () => {
-  },
-  removeCoAgentStateRender: () => {
-  },
-  chatComponentsCache: { current: { actions: {}, coAgentStateRenders: {} } },
-  getContextString: (documents, categories) => returnAndThrowInDebug(""),
-  addContext: () => "",
-  removeContext: () => {
-  },
-  getAllContext: () => [],
-  getFunctionCallHandler: () => returnAndThrowInDebug(() => __async(void 0, null, function* () {
-  })),
-  isLoading: false,
-  setIsLoading: () => returnAndThrowInDebug(false),
-  chatInstructions: "",
-  setChatInstructions: () => returnAndThrowInDebug(""),
-  additionalInstructions: [],
-  setAdditionalInstructions: () => returnAndThrowInDebug([]),
-  getDocumentsContext: (categories) => returnAndThrowInDebug([]),
-  addDocumentContext: () => returnAndThrowInDebug(""),
-  removeDocumentContext: () => {
-  },
-  runtimeClient: {},
-  copilotApiConfig: new class {
-    get chatApiEndpoint() {
-      throw new Error("Remember to wrap your app in a `<CopilotKit> {...} </CopilotKit>` !!!");
-    }
-    get headers() {
-      return {};
-    }
-    get body() {
-      return {};
-    }
-  }(),
-  chatSuggestionConfiguration: {},
-  addChatSuggestionConfiguration: () => {
-  },
-  removeChatSuggestionConfiguration: () => {
-  },
-  showDevConsole: false,
-  coagentStates: {},
-  setCoagentStates: () => {
-  },
-  coagentStatesRef: { current: {} },
-  setCoagentStatesWithRef: () => {
-  },
-  agentSession: null,
-  setAgentSession: () => {
-  },
-  forwardedParameters: {},
-  agentLock: null,
-  threadId: "",
-  setThreadId: () => {
-  },
-  runId: null,
-  setRunId: () => {
-  },
-  chatAbortControllerRef: { current: null },
-  availableAgents: [],
-  extensions: {},
-  setExtensions: () => {
-  },
-  langGraphInterruptAction: null,
-  setLangGraphInterruptAction: () => null,
-  removeLangGraphInterruptAction: () => null
-};
-var CopilotContext = import_react3.default.createContext(emptyCopilotContext);
-function useCopilotContext() {
-  const context = import_react3.default.useContext(CopilotContext);
-  if (context === emptyCopilotContext) {
-    throw new Error("Remember to wrap your app in a `<CopilotKit> {...} </CopilotKit>` !!!");
-  }
-  return context;
-}
-function returnAndThrowInDebug(_value) {
-  throw new Error("Remember to wrap your app in a `<CopilotKit> {...} </CopilotKit>` !!!");
-}
-
 // src/components/error-boundary/error-utils.tsx
 var import_react4 = require("react");
 
@@ -458,6 +490,35 @@ function useChat(options) {
   } = options;
   const runChatCompletionRef = (0, import_react5.useRef)();
   const addErrorToast = useErrorToast();
+  const { setBannerError } = useToast();
+  const { onError } = useCopilotContext();
+  const traceUIError = (error, originalError) => __async(this, null, function* () {
+    if (!onError || !(copilotConfig == null ? void 0 : copilotConfig.publicApiKey))
+      return;
+    try {
+      const traceEvent = {
+        type: "error",
+        timestamp: Date.now(),
+        context: {
+          source: "ui",
+          request: {
+            operation: "useChatCompletion",
+            url: copilotConfig.chatApiEndpoint,
+            startTime: Date.now()
+          },
+          technical: {
+            environment: "browser",
+            userAgent: typeof navigator !== "undefined" ? navigator.userAgent : void 0,
+            stackTrace: originalError instanceof Error ? originalError.stack : void 0
+          }
+        },
+        error
+      };
+      yield onError(traceEvent);
+    } catch (traceError) {
+      console.error("Error in use-chat onError handler:", traceError);
+    }
+  });
   const agentSessionRef = (0, import_react5.useRef)(agentSession);
   agentSessionRef.current = agentSession;
   const runIdRef = (0, import_react5.useRef)(runId);
@@ -477,7 +538,7 @@ function useChat(options) {
   const pendingAppendsRef = (0, import_react5.useRef)([]);
   const runChatCompletion = useAsyncCallback(
     (previousMessages) => __async(this, null, function* () {
-      var _a, _b, _c, _d, _e, _f, _g, _h, _i, _j, _k, _l, _m, _n, _o;
+      var _a, _b, _c, _d, _e, _f, _g, _h, _i, _j, _k, _l, _m, _n, _o, _p, _q, _r, _s, _t;
       setIsLoading(true);
       const interruptEvent = langGraphInterruptAction == null ? void 0 : langGraphInterruptAction.event;
       if ((interruptEvent == null ? void 0 : interruptEvent.name) === import_runtime_client_gql3.MetaEventName.LangGraphInterruptEvent && (interruptEvent == null ? void 0 : interruptEvent.value) && !(interruptEvent == null ? void 0 : interruptEvent.response) && agentSessionRef.current) {
@@ -609,20 +670,53 @@ function useChat(options) {
           messages2 = (0, import_runtime_client_gql3.convertGqlOutputToMessages)(
             (0, import_runtime_client_gql3.filterAdjacentAgentStateMessages)(rawMessagesResponse)
           );
-          if (messages2.length === 0) {
-            continue;
-          }
           newMessages = [];
           if (((_k = value.generateCopilotResponse.status) == null ? void 0 : _k.__typename) === "FailedResponseStatus" && value.generateCopilotResponse.status.reason === "GUARDRAILS_VALIDATION_FAILED") {
+            const guardrailsReason = ((_l = value.generateCopilotResponse.status.details) == null ? void 0 : _l.guardrailsReason) || "";
             newMessages = [
               new import_runtime_client_gql3.TextMessage({
                 role: import_runtime_client_gql3.MessageRole.Assistant,
-                content: ((_l = value.generateCopilotResponse.status.details) == null ? void 0 : _l.guardrailsReason) || ""
+                content: guardrailsReason
               })
             ];
+            const guardrailsError = new import_shared4.CopilotKitError({
+              message: `Guardrails validation failed: ${guardrailsReason}`,
+              code: import_shared4.CopilotKitErrorCode.MISUSE
+            });
+            yield traceUIError(guardrailsError, {
+              statusReason: value.generateCopilotResponse.status.reason,
+              statusDetails: value.generateCopilotResponse.status.details
+            });
             setMessages([...previousMessages, ...newMessages]);
             break;
-          } else {
+          }
+          if (((_m = value.generateCopilotResponse.status) == null ? void 0 : _m.__typename) === "FailedResponseStatus" && value.generateCopilotResponse.status.reason === "UNKNOWN_ERROR") {
+            const errorMessage = ((_n = value.generateCopilotResponse.status.details) == null ? void 0 : _n.description) || "An unknown error occurred";
+            const statusDetails = value.generateCopilotResponse.status.details;
+            const originalError = (statusDetails == null ? void 0 : statusDetails.originalError) || (statusDetails == null ? void 0 : statusDetails.error);
+            const originalCode = (originalError == null ? void 0 : originalError.code) || ((_o = originalError == null ? void 0 : originalError.extensions) == null ? void 0 : _o.code);
+            const originalSeverity = (originalError == null ? void 0 : originalError.severity) || ((_p = originalError == null ? void 0 : originalError.extensions) == null ? void 0 : _p.severity);
+            const originalVisibility = (originalError == null ? void 0 : originalError.visibility) || ((_q = originalError == null ? void 0 : originalError.extensions) == null ? void 0 : _q.visibility);
+            let errorCode = import_shared4.CopilotKitErrorCode.NETWORK_ERROR;
+            if (originalCode && Object.values(import_shared4.CopilotKitErrorCode).includes(originalCode)) {
+              errorCode = originalCode;
+            }
+            const structuredError = new import_shared4.CopilotKitError({
+              message: errorMessage,
+              code: errorCode,
+              severity: originalSeverity,
+              visibility: originalVisibility
+            });
+            setBannerError(structuredError);
+            yield traceUIError(structuredError, {
+              statusReason: value.generateCopilotResponse.status.reason,
+              statusDetails: value.generateCopilotResponse.status.details,
+              originalErrorCode: originalCode,
+              preservedStructure: !!originalCode
+            });
+            setIsLoading(false);
+            break;
+          } else if (messages2.length > 0) {
             newMessages = [...messages2];
             for (const message of messages2) {
               if (message.isAgentStateMessage() && !message.active && !executedCoAgentStateRenders.includes(message.id) && onCoAgentStateRender) {
@@ -684,6 +778,39 @@ function useChat(options) {
           newMessages
         );
         let didExecuteAction = false;
+        const executeActionFromMessage = (currentAction, actionMessage) => __async(this, null, function* () {
+          var _a2;
+          const isInterruptAction = interruptMessages.find((m) => m.id === actionMessage.id);
+          followUp = (_a2 = currentAction == null ? void 0 : currentAction.followUp) != null ? _a2 : !isInterruptAction;
+          if (currentAction == null ? void 0 : currentAction._setActivatingMessageId) {
+            currentAction._setActivatingMessageId(actionMessage.id);
+          }
+          const resultMessage = yield executeAction({
+            onFunctionCall,
+            message: actionMessage,
+            chatAbortControllerRef,
+            onError: (error) => {
+              addErrorToast([error]);
+              console.error(`Failed to execute action ${actionMessage.name}: ${error}`);
+            },
+            setMessages,
+            getFinalMessages: () => finalMessages,
+            isRenderAndWait: (currentAction == null ? void 0 : currentAction._isRenderAndWait) || false
+          });
+          didExecuteAction = true;
+          const messageIndex = finalMessages.findIndex((msg) => msg.id === actionMessage.id);
+          finalMessages.splice(messageIndex + 1, 0, resultMessage);
+          if (currentAction == null ? void 0 : currentAction._isRenderAndWait) {
+            const messagesForImmediateUpdate = [...finalMessages];
+            (0, import_react_dom.flushSync)(() => {
+              setMessages(messagesForImmediateUpdate);
+            });
+          }
+          if (currentAction == null ? void 0 : currentAction._setActivatingMessageId) {
+            currentAction._setActivatingMessageId(null);
+          }
+          return resultMessage;
+        });
         if (onFunctionCall) {
           const lastMessages = [];
           for (let i = finalMessages.length - 1; i >= 0; i--) {
@@ -700,37 +827,28 @@ function useChat(options) {
               (action2) => action2.name === message.name
             );
             const currentResultMessagePairedFeAction = message.isResultMessage() ? getPairedFeAction(actions, message) : null;
-            const executeActionFromMessage = (action2, message2) => __async(this, null, function* () {
-              var _a2;
-              const isInterruptAction = interruptMessages.find((m) => m.id === message2.id);
-              followUp = (_a2 = action2 == null ? void 0 : action2.followUp) != null ? _a2 : !isInterruptAction;
-              const resultMessage = yield executeAction({
-                onFunctionCall,
-                previousMessages,
-                message: message2,
-                chatAbortControllerRef,
-                onError: (error) => {
-                  addErrorToast([error]);
-                  console.error(`Failed to execute action ${message2.name}: ${error}`);
-                }
-              });
-              didExecuteAction = true;
-              const messageIndex = finalMessages.findIndex((msg) => msg.id === message2.id);
-              finalMessages.splice(messageIndex + 1, 0, resultMessage);
-              return resultMessage;
-            });
             if (action && message.isActionExecutionMessage()) {
-              const resultMessage = yield executeActionFromMessage(action, message);
-              const pairedFeAction = getPairedFeAction(actions, resultMessage);
-              if (pairedFeAction) {
-                const newExecutionMessage = new import_runtime_client_gql3.ActionExecutionMessage({
-                  name: pairedFeAction.name,
-                  arguments: (0, import_shared4.parseJson)(resultMessage.result, resultMessage.result),
-                  status: message.status,
-                  createdAt: message.createdAt,
-                  parentMessageId: message.parentMessageId
-                });
-                yield executeActionFromMessage(pairedFeAction, newExecutionMessage);
+              const isRenderAndWaitAction = (action == null ? void 0 : action._isRenderAndWait) || false;
+              const alreadyProcessed = isRenderAndWaitAction && finalMessages.some(
+                (fm) => fm.isResultMessage() && fm.actionExecutionId === message.id
+              );
+              if (alreadyProcessed) {
+              } else {
+                const resultMessage = yield executeActionFromMessage(
+                  action,
+                  message
+                );
+                const pairedFeAction = getPairedFeAction(actions, resultMessage);
+                if (pairedFeAction) {
+                  const newExecutionMessage = new import_runtime_client_gql3.ActionExecutionMessage({
+                    name: pairedFeAction.name,
+                    arguments: (0, import_shared4.parseJson)(resultMessage.result, resultMessage.result),
+                    status: message.status,
+                    createdAt: message.createdAt,
+                    parentMessageId: message.parentMessageId
+                  });
+                  yield executeActionFromMessage(pairedFeAction, newExecutionMessage);
+                }
               }
             } else if (message.isResultMessage() && currentResultMessagePairedFeAction) {
               const newExecutionMessage = new import_runtime_client_gql3.ActionExecutionMessage({
@@ -748,16 +866,12 @@ function useChat(options) {
           }
           setMessages(finalMessages);
         }
-        if (
-          // if followUp is not explicitly false
-          followUp !== false && // and we executed an action
-          (didExecuteAction || // the last message is a server side result
-          !isAgentRun && finalMessages.length && finalMessages[finalMessages.length - 1].isResultMessage()) && // the user did not stop generation
-          !((_m = chatAbortControllerRef.current) == null ? void 0 : _m.signal.aborted)
-        ) {
+        if (followUp !== false && (didExecuteAction || // the last message is a server side result
+        !isAgentRun && finalMessages.length && finalMessages[finalMessages.length - 1].isResultMessage()) && // the user did not stop generation
+        !((_r = chatAbortControllerRef.current) == null ? void 0 : _r.signal.aborted)) {
           yield new Promise((resolve) => setTimeout(resolve, 10));
           return yield runChatCompletionRef.current(finalMessages);
-        } else if ((_n = chatAbortControllerRef.current) == null ? void 0 : _n.signal.aborted) {
+        } else if ((_s = chatAbortControllerRef.current) == null ? void 0 : _s.signal.aborted) {
           const repairedMessages = finalMessages.filter((message, actionExecutionIndex) => {
             if (message.isActionExecutionMessage()) {
               return finalMessages.find(
@@ -768,7 +882,7 @@ function useChat(options) {
           });
           const repairedMessageIds = repairedMessages.map((message) => message.id);
           setMessages(repairedMessages);
-          if ((_o = agentSessionRef.current) == null ? void 0 : _o.nodeName) {
+          if ((_t = agentSessionRef.current) == null ? void 0 : _t.nodeName) {
             setAgentSession({
               threadId: agentSessionRef.current.threadId,
               agentName: agentSessionRef.current.agentName,
@@ -862,21 +976,35 @@ function useChat(options) {
     [isLoading, messages, setMessages, runChatCompletionAndHandleFunctionCall]
   );
   const reload = useAsyncCallback(
-    (messageId) => __async(this, null, function* () {
+    (reloadMessageId) => __async(this, null, function* () {
       if (isLoading || messages.length === 0) {
         return;
       }
-      const index = messages.findIndex((msg) => msg.id === messageId);
-      if (index === -1) {
-        console.warn(`Message with id ${messageId} not found`);
+      const reloadMessageIndex = messages.findIndex((msg) => msg.id === reloadMessageId);
+      if (reloadMessageIndex === -1) {
+        console.warn(`Message with id ${reloadMessageId} not found`);
         return;
       }
-      let newMessages = messages.slice(0, index);
-      if (newMessages.length > 0 && newMessages[newMessages.length - 1].isAgentStateMessage()) {
-        newMessages = newMessages.slice(0, newMessages.length - 1);
+      const reloadMessageRole = messages[reloadMessageIndex].role;
+      if (reloadMessageRole !== import_runtime_client_gql3.MessageRole.Assistant) {
+        console.warn(`Regenerate cannot be performed on ${reloadMessageRole} role`);
+        return;
       }
-      setMessages(newMessages);
-      return runChatCompletionAndHandleFunctionCall(newMessages);
+      let historyCutoff = [];
+      if (messages.length > 2) {
+        const lastUserMessageBeforeRegenerate = messages.slice(0, reloadMessageIndex).reverse().find(
+          (msg) => (
+            // @ts-expect-error -- message has role
+            msg.role === import_runtime_client_gql3.MessageRole.User
+          )
+        );
+        const indexOfLastUserMessageBeforeRegenerate = messages.findIndex(
+          (msg) => msg.id === lastUserMessageBeforeRegenerate.id
+        );
+        historyCutoff = messages.slice(0, indexOfLastUserMessageBeforeRegenerate + 1);
+      }
+      setMessages(historyCutoff);
+      return runChatCompletionAndHandleFunctionCall(historyCutoff);
     }),
     [isLoading, messages, setMessages, runChatCompletionAndHandleFunctionCall]
   );
@@ -911,20 +1039,31 @@ function constructFinalMessages(syncedMessages, previousMessages, newMessages) {
 function executeAction(_0) {
   return __async(this, arguments, function* ({
     onFunctionCall,
-    previousMessages,
     message,
     chatAbortControllerRef,
-    onError
+    onError,
+    setMessages,
+    getFinalMessages,
+    isRenderAndWait
   }) {
     let result;
     let error = null;
+    const currentMessagesForHandler = getFinalMessages();
+    const handlerReturnedPromise = onFunctionCall({
+      messages: currentMessagesForHandler,
+      name: message.name,
+      args: message.arguments
+    });
+    if (isRenderAndWait) {
+      const currentMessagesForRender = getFinalMessages();
+      (0, import_react_dom.flushSync)(() => {
+        setMessages([...currentMessagesForRender]);
+      });
+    }
     try {
       result = yield Promise.race([
-        onFunctionCall({
-          messages: previousMessages,
-          name: message.name,
-          args: message.arguments
-        }),
+        handlerReturnedPromise,
+        // Await the promise returned by the handler
         new Promise(
           (resolve) => {
             var _a;