@copilotkit/react-core 1.4.8-no-pino-redact.1 → 1.5.0-coagents-v0-3.0

package/src/hooks/use-chat.ts

@@ -18,14 +18,16 @@ import {
   Role,
   CopilotRequestType,
   ActionInputAvailability,
+  loadMessagesFromJsonRepresentation,
 } from "@copilotkit/runtime-client-gql";
 
 import { CopilotApiConfig } from "../context";
-import { FrontendAction } from "../types/frontend-action";
+import { FrontendAction, processActionsForRuntimeRequest } from "../types/frontend-action";
 import { CoagentState } from "../types/coagent-state";
 import { AgentSession } from "../context/copilot-context";
 import { useToast } from "../components/toast/toast-provider";
 import { useCopilotRuntimeClient } from "./use-copilot-runtime-client";
+import { useAsyncCallback } from "../components/error-boundary/error-utils";
 
 export type UseChatOptions = {
   /**
@@ -97,6 +99,27 @@ export type UseChatOptions = {
    * setState-powered method to update the agent session
    */
   setAgentSession: React.Dispatch<React.SetStateAction<AgentSession | null>>;
+
+  /**
+   * The current thread ID.
+   */
+  threadId: string | null;
+  /**
+   * set the current thread ID
+   */
+  setThreadId: (threadId: string | null) => void;
+  /**
+   * The current run ID.
+   */
+  runId: string | null;
+  /**
+   * set the current run ID
+   */
+  setRunId: (runId: string | null) => void;
+  /**
+   * The global chat abort controller.
+   */
+  chatAbortControllerRef: React.MutableRefObject<AbortController | null>;
 };
 
 export type UseChatHelpers = {
@@ -139,19 +162,23 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
     coagentStatesRef,
     agentSession,
     setAgentSession,
+    threadId,
+    setThreadId,
+    runId,
+    setRunId,
+    chatAbortControllerRef,
   } = options;
-
-  const abortControllerRef = useRef<AbortController>();
-  const threadIdRef = useRef<string | null>(null);
-  const runIdRef = useRef<string | null>(null);
   const { addGraphQLErrorsToast } = useToast();
-
   const runChatCompletionRef = useRef<(previousMessages: Message[]) => Promise<Message[]>>();
   // We need to keep a ref of coagent states and session because of renderAndWait - making sure
   // the latest state is sent to the API
   // This is a workaround and needs to be addressed in the future
   const agentSessionRef = useRef<AgentSession | null>(agentSession);
   agentSessionRef.current = agentSession;
+  const threadIdRef = useRef<string | null>(threadId);
+  threadIdRef.current = threadId;
+  const runIdRef = useRef<string | null>(runId);
+  runIdRef.current = runId;
 
   const publicApiKey = copilotConfig.publicApiKey;
 
@@ -167,292 +194,371 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
     credentials: copilotConfig.credentials,
   });
 
-  const runChatCompletion = async (previousMessages: Message[]): Promise<Message[]> => {
-    setIsLoading(true);
-
-    // this message is just a placeholder. It will disappear once the first real message
-    // is received
-    let newMessages: Message[] = [
-      new TextMessage({
-        content: "",
-        role: Role.Assistant,
-      }),
-    ];
-    const abortController = new AbortController();
-    abortControllerRef.current = abortController;
-
-    setMessages([...previousMessages, ...newMessages]);
-
-    const systemMessage = makeSystemMessageCallback();
-
-    const messagesWithContext = [systemMessage, ...(initialMessages || []), ...previousMessages];
-
-    const stream = runtimeClient.asStream(
-      runtimeClient.generateCopilotResponse({
-        data: {
-          frontend: {
-            actions: actions
-              .filter(
-                (action) =>
-                  action.available !== ActionInputAvailability.Disabled || !action.disabled,
-              )
-              .map((action) => {
-                let available: ActionInputAvailability | undefined =
-                  ActionInputAvailability.Enabled;
-                if (action.disabled) {
-                  available = ActionInputAvailability.Disabled;
-                } else if (action.available === "disabled") {
-                  available = ActionInputAvailability.Disabled;
-                } else if (action.available === "remote") {
-                  available = ActionInputAvailability.Remote;
-                }
-                return {
-                  name: action.name,
-                  description: action.description || "",
-                  jsonSchema: JSON.stringify(actionParametersToJsonSchema(action.parameters || [])),
-                  available,
-                };
-              }),
-            url: window.location.href,
-          },
-          threadId: threadIdRef.current,
-          runId: runIdRef.current,
-          messages: convertMessagesToGqlInput(filterAgentStateMessages(messagesWithContext)),
-          ...(copilotConfig.cloud
-            ? {
-                cloud: {
-                  ...(copilotConfig.cloud.guardrails?.input?.restrictToTopic?.enabled
-                    ? {
-                        guardrails: {
-                          inputValidationRules: {
-                            allowList:
-                              copilotConfig.cloud.guardrails.input.restrictToTopic.validTopics,
-                            denyList:
-                              copilotConfig.cloud.guardrails.input.restrictToTopic.invalidTopics,
+  const runChatCompletion = useAsyncCallback(
+    async (previousMessages: Message[]): Promise<Message[]> => {
+      setIsLoading(true);
+
+      // this message is just a placeholder. It will disappear once the first real message
+      // is received
+      let newMessages: Message[] = [
+        new TextMessage({
+          content: "",
+          role: Role.Assistant,
+        }),
+      ];
+
+      chatAbortControllerRef.current = new AbortController();
+
+      setMessages([...previousMessages, ...newMessages]);
+
+      const systemMessage = makeSystemMessageCallback();
+
+      const messagesWithContext = [systemMessage, ...(initialMessages || []), ...previousMessages];
+
+      const isAgentRun = agentSessionRef.current !== null;
+
+      const stream = runtimeClient.asStream(
+        runtimeClient.generateCopilotResponse({
+          data: {
+            frontend: {
+              actions: processActionsForRuntimeRequest(actions),
+              url: window.location.href,
+            },
+            threadId: threadIdRef.current,
+            runId: runIdRef.current,
+            messages: convertMessagesToGqlInput(filterAgentStateMessages(messagesWithContext)),
+            ...(copilotConfig.cloud
+              ? {
+                  cloud: {
+                    ...(copilotConfig.cloud.guardrails?.input?.restrictToTopic?.enabled
+                      ? {
+                          guardrails: {
+                            inputValidationRules: {
+                              allowList:
+                                copilotConfig.cloud.guardrails.input.restrictToTopic.validTopics,
+                              denyList:
+                                copilotConfig.cloud.guardrails.input.restrictToTopic.invalidTopics,
+                            },
                           },
-                        },
-                      }
-                    : {}),
-                },
-              }
-            : {}),
-          metadata: {
-            requestType: CopilotRequestType.Chat,
+                        }
+                      : {}),
+                  },
+                }
+              : {}),
+            metadata: {
+              requestType: CopilotRequestType.Chat,
+            },
+            ...(agentSessionRef.current
+              ? {
+                  agentSession: agentSessionRef.current,
+                }
+              : {}),
+            agentStates: Object.values(coagentStatesRef.current!).map((state) => ({
+              agentName: state.name,
+              state: JSON.stringify(state.state),
+            })),
           },
-          ...(agentSessionRef.current
-            ? {
-                agentSession: agentSessionRef.current,
-              }
-            : {}),
-          agentStates: Object.values(coagentStatesRef.current!).map((state) => ({
-            agentName: state.name,
-            state: JSON.stringify(state.state),
-          })),
-        },
-        properties: copilotConfig.properties,
-        signal: abortControllerRef.current?.signal,
-      }),
-    );
-
-    const guardrailsEnabled =
-      copilotConfig.cloud?.guardrails?.input?.restrictToTopic.enabled || false;
-
-    const reader = stream.getReader();
-
-    let actionResults: { [id: string]: string } = {};
-    let executedCoAgentStateRenders: string[] = [];
-    let followUp: FrontendAction["followUp"] = undefined;
-
-    try {
-      while (true) {
-        let done, value;
-
-        try {
-          const readResult = await reader.read();
-          done = readResult.done;
-          value = readResult.value;
-        } catch (readError) {
-          break;
-        }
+          properties: copilotConfig.properties,
+          signal: chatAbortControllerRef.current?.signal,
+        }),
+      );
 
-        if (done) {
-          break;
-        }
+      const guardrailsEnabled =
+        copilotConfig.cloud?.guardrails?.input?.restrictToTopic.enabled || false;
 
-        if (!value?.generateCopilotResponse) {
-          continue;
-        }
+      const reader = stream.getReader();
 
-        threadIdRef.current = value.generateCopilotResponse.threadId || null;
-        runIdRef.current = value.generateCopilotResponse.runId || null;
+      let executedCoAgentStateRenders: string[] = [];
+      let followUp: FrontendAction["followUp"] = undefined;
 
-        const messages = convertGqlOutputToMessages(
-          filterAdjacentAgentStateMessages(value.generateCopilotResponse.messages),
-        );
+      let messages: Message[] = [];
+      let syncedMessages: Message[] = [];
 
-        if (messages.length === 0) {
-          continue;
-        }
+      try {
+        while (true) {
+          let done, value;
 
-        newMessages = [];
+          try {
+            const readResult = await reader.read();
+            done = readResult.done;
+            value = readResult.value;
+          } catch (readError) {
+            break;
+          }
 
-        // request failed, display error message
-        if (
-          value.generateCopilotResponse.status?.__typename === "FailedResponseStatus" &&
-          value.generateCopilotResponse.status.reason === "GUARDRAILS_VALIDATION_FAILED"
-        ) {
-          newMessages = [
-            new TextMessage({
-              role: MessageRole.Assistant,
-              content: value.generateCopilotResponse.status.details?.guardrailsReason || "",
-            }),
-          ];
+          if (done) {
+            if (chatAbortControllerRef.current.signal.aborted) {
+              return [];
+            }
+            break;
+          }
+
+          if (!value?.generateCopilotResponse) {
+            continue;
+          }
+
+          threadIdRef.current = value.generateCopilotResponse.threadId || null;
+          runIdRef.current = value.generateCopilotResponse.runId || null;
+
+          setThreadId(threadIdRef.current);
+          setRunId(runIdRef.current);
+
+          messages = convertGqlOutputToMessages(
+            filterAdjacentAgentStateMessages(value.generateCopilotResponse.messages),
+          );
+
+          if (messages.length === 0) {
+            continue;
+          }
+
+          newMessages = [];
+
+          // request failed, display error message and quit
+          if (
+            value.generateCopilotResponse.status?.__typename === "FailedResponseStatus" &&
+            value.generateCopilotResponse.status.reason === "GUARDRAILS_VALIDATION_FAILED"
+          ) {
+            newMessages = [
+              new TextMessage({
+                role: MessageRole.Assistant,
+                content: value.generateCopilotResponse.status.details?.guardrailsReason || "",
+              }),
+            ];
+            setMessages([...previousMessages, ...newMessages]);
+            break;
+          }
+
+          // add messages to the chat
+          else {
+            newMessages = [...messages];
+
+            for (const message of messages) {
+              // execute onCoAgentStateRender handler
+              if (
+                message.isAgentStateMessage() &&
+                !message.active &&
+                !executedCoAgentStateRenders.includes(message.id) &&
+                onCoAgentStateRender
+              ) {
+                // Do not execute a coagent action if guardrails are enabled but the status is not known
+                if (guardrailsEnabled && value.generateCopilotResponse.status === undefined) {
+                  break;
+                }
+                // execute coagent action
+                await onCoAgentStateRender({
+                  name: message.agentName,
+                  nodeName: message.nodeName,
+                  state: message.state,
+                });
+                executedCoAgentStateRenders.push(message.id);
+              }
+            }
+
+            const lastAgentStateMessage = [...messages]
+              .reverse()
+              .find((message) => message.isAgentStateMessage());
+
+            if (lastAgentStateMessage) {
+              if (
+                lastAgentStateMessage.state.messages &&
+                lastAgentStateMessage.state.messages.length > 0
+              ) {
+                syncedMessages = loadMessagesFromJsonRepresentation(
+                  lastAgentStateMessage.state.messages,
+                );
+              }
+              setCoagentStatesWithRef((prevAgentStates) => ({
+                ...prevAgentStates,
+                [lastAgentStateMessage.agentName]: {
+                  name: lastAgentStateMessage.agentName,
+                  state: lastAgentStateMessage.state,
+                  running: lastAgentStateMessage.running,
+                  active: lastAgentStateMessage.active,
+                  threadId: lastAgentStateMessage.threadId,
+                  nodeName: lastAgentStateMessage.nodeName,
+                  runId: lastAgentStateMessage.runId,
+                },
+              }));
+              if (lastAgentStateMessage.running) {
+                setAgentSession({
+                  threadId: lastAgentStateMessage.threadId,
+                  agentName: lastAgentStateMessage.agentName,
+                  nodeName: lastAgentStateMessage.nodeName,
+                });
+              } else {
+                setAgentSession(null);
+              }
+            }
+          }
+
+          if (newMessages.length > 0) {
+            // Update message state
+            setMessages([...previousMessages, ...newMessages]);
+          }
         }
+        const finalMessages = constructFinalMessages(syncedMessages, previousMessages, newMessages);
+
+        let didExecuteAction = false;
 
-        // add messages to the chat
-        else {
-          for (const message of messages) {
-            newMessages.push(message);
-            // execute regular action executions
+        // execute regular action executions that are specific to the frontend (last actions)
+        if (onFunctionCall) {
+          // Find consecutive action execution messages at the end
+          const lastMessages = [];
+          for (let i = finalMessages.length - 1; i >= 0; i--) {
+            const message = finalMessages[i];
             if (
               message.isActionExecutionMessage() &&
-              message.status.code !== MessageStatusCode.Pending &&
-              message.scope === "client" &&
-              onFunctionCall
+              message.status.code !== MessageStatusCode.Pending
             ) {
-              if (!(message.id in actionResults)) {
-                // Do not execute a function call if guardrails are enabled but the status is not known
-                if (guardrailsEnabled && value.generateCopilotResponse.status === undefined) {
-                  break;
-                }
-                // execute action
-                try {
-                  // We update the message state before calling the handler so that the render
-                  // function can be called with `executing` state
-                  setMessages([...previousMessages, ...newMessages]);
+              lastMessages.unshift(message);
+            } else {
+              break;
+            }
+          }
 
-                  const action = actions.find((action) => action.name === message.name);
+          for (const message of lastMessages) {
+            // We update the message state before calling the handler so that the render
+            // function can be called with `executing` state
+            setMessages(finalMessages);
 
-                  if (action) {
-                    followUp = action.followUp;
-                  }
+            const action = actions.find((action) => action.name === message.name);
 
-                  const result = await onFunctionCall({
+            if (action) {
+              followUp = action.followUp;
+              let result: any;
+              try {
+                result = await Promise.race([
+                  onFunctionCall({
                     messages: previousMessages,
                     name: message.name,
                     args: message.arguments,
-                  });
-                  actionResults[message.id] = result;
-                } catch (e) {
-                  actionResults[message.id] = `Failed to execute action ${message.name}`;
-                  console.error(`Failed to execute action ${message.name}: ${e}`);
-                }
+                  }),
+                  new Promise((resolve) =>
+                    chatAbortControllerRef.current?.signal.addEventListener("abort", () =>
+                      resolve("Operation was aborted by the user"),
+                    ),
+                  ),
+                  // if the user stopped generation, we also abort consecutive actions
+                  new Promise((resolve) => {
+                    if (chatAbortControllerRef.current?.signal.aborted) {
+                      resolve("Operation was aborted by the user");
+                    }
+                  }),
+                ]);
+              } catch (e) {
+                result = `Failed to execute action ${message.name}`;
+                console.error(`Failed to execute action ${message.name}: ${e}`);
               }
-              // add the result message
-              newMessages.push(
+              didExecuteAction = true;
+              const messageIndex = finalMessages.findIndex((msg) => msg.id === message.id);
+              finalMessages.splice(
+                messageIndex + 1,
+                0,
                 new ResultMessage({
-                  result: ResultMessage.encodeResult(actionResults[message.id]),
+                  id: "result-" + message.id,
+                  result: ResultMessage.encodeResult(result),
                   actionExecutionId: message.id,
                   actionName: message.name,
                 }),
               );
             }
-            // execute coagent actions
-            if (
-              message.isAgentStateMessage() &&
-              !message.active &&
-              !executedCoAgentStateRenders.includes(message.id) &&
-              onCoAgentStateRender
-            ) {
-              // Do not execute a coagent action if guardrails are enabled but the status is not known
-              if (guardrailsEnabled && value.generateCopilotResponse.status === undefined) {
-                break;
-              }
-              // execute coagent action
-              await onCoAgentStateRender({
-                name: message.agentName,
-                nodeName: message.nodeName,
-                state: message.state,
-              });
-              executedCoAgentStateRenders.push(message.id);
-            }
           }
 
-          const lastAgentStateMessage = [...messages]
-            .reverse()
-            .find((message) => message.isAgentStateMessage());
-
-          if (lastAgentStateMessage) {
-            setCoagentStatesWithRef((prevAgentStates) => ({
-              ...prevAgentStates,
-              [lastAgentStateMessage.agentName]: {
-                name: lastAgentStateMessage.agentName,
-                state: lastAgentStateMessage.state,
-                running: lastAgentStateMessage.running,
-                active: lastAgentStateMessage.active,
-                threadId: lastAgentStateMessage.threadId,
-                nodeName: lastAgentStateMessage.nodeName,
-                runId: lastAgentStateMessage.runId,
-              },
-            }));
-            if (lastAgentStateMessage.running) {
-              setAgentSession({
-                threadId: lastAgentStateMessage.threadId,
-                agentName: lastAgentStateMessage.agentName,
-                nodeName: lastAgentStateMessage.nodeName,
-              });
-            } else {
-              setAgentSession(null);
-            }
-          }
+          setMessages(finalMessages);
         }
 
-        if (newMessages.length > 0) {
-          // Update message state
-          setMessages([...previousMessages, ...newMessages]);
+        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
+          !chatAbortControllerRef.current?.signal.aborted
+        ) {
+          // run the completion again and return the result
+
+          // wait for next tick to make sure all the react state updates
+          // - tried using react-dom's flushSync, but it did not work
+          await new Promise((resolve) => setTimeout(resolve, 10));
+
+          return await runChatCompletionRef.current!(finalMessages);
+        } else if (chatAbortControllerRef.current?.signal.aborted) {
+          // filter out all the action execution messages that do not have a consecutive matching result message
+          const repairedMessages = finalMessages.filter((message, actionExecutionIndex) => {
+            if (message.isActionExecutionMessage()) {
+              return finalMessages.find(
+                (msg, resultIndex) =>
+                  msg.isResultMessage() &&
+                  msg.actionExecutionId === message.id &&
+                  resultIndex === actionExecutionIndex + 1,
+              );
+            }
+            return true;
+          });
+          const repairedMessageIds = repairedMessages.map((message) => message.id);
+          setMessages(repairedMessages);
+
+          if (agentSessionRef.current?.nodeName) {
+            setAgentSession({
+              threadId: agentSessionRef.current.threadId,
+              agentName: agentSessionRef.current.agentName,
+              nodeName: "__end__",
+            });
+          }
+          // only return new messages that were not filtered out
+          return newMessages.filter((message) => repairedMessageIds.includes(message.id));
+        } else {
+          return newMessages.slice();
         }
+      } finally {
+        setIsLoading(false);
       }
-
-      if (
-        // if followUp is not explicitly false
-        followUp !== false &&
-        // if we have client side results
-        (Object.values(actionResults).length ||
-          // or the last message we received is a result
-          (newMessages.length && newMessages[newMessages.length - 1].isResultMessage()))
-      ) {
-        // run the completion again and return the result
-
-        // wait for next tick to make sure all the react state updates
-        // - tried using react-dom's flushSync, but it did not work
-        await new Promise((resolve) => setTimeout(resolve, 10));
-
-        return await runChatCompletionRef.current!([...previousMessages, ...newMessages]);
-      } else {
-        return newMessages.slice();
-      }
-    } finally {
-      setIsLoading(false);
-    }
-  };
+    },
+    [
+      messages,
+      setMessages,
+      makeSystemMessageCallback,
+      copilotConfig,
+      setIsLoading,
+      initialMessages,
+      isLoading,
+      actions,
+      onFunctionCall,
+      onCoAgentStateRender,
+      setCoagentStatesWithRef,
+      coagentStatesRef,
+      agentSession,
+      setAgentSession,
+    ],
+  );
 
   runChatCompletionRef.current = runChatCompletion;
 
-  const runChatCompletionAndHandleFunctionCall = async (messages: Message[]): Promise<void> => {
-    await runChatCompletionRef.current!(messages);
-  };
+  const runChatCompletionAndHandleFunctionCall = useAsyncCallback(
+    async (messages: Message[]): Promise<void> => {
+      await runChatCompletionRef.current!(messages);
+    },
+    [messages],
+  );
+
+  const append = useAsyncCallback(
+    async (message: Message): Promise<void> => {
+      if (isLoading) {
+        return;
+      }
 
-  const append = async (message: Message): Promise<void> => {
-    if (isLoading) {
-      return;
-    }
+      const newMessages = [...messages, message];
+      setMessages(newMessages);
+      return runChatCompletionAndHandleFunctionCall(newMessages);
+    },
+    [isLoading, messages, setMessages, runChatCompletionAndHandleFunctionCall],
+  );
 
-    const newMessages = [...messages, message];
-    setMessages(newMessages);
-    return runChatCompletionAndHandleFunctionCall(newMessages);
-  };
-
-  const reload = async (): Promise<void> => {
+  const reload = useAsyncCallback(async (): Promise<void> => {
     if (isLoading || messages.length === 0) {
       return;
     }
@@ -466,10 +572,10 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
     setMessages(newMessages);
 
     return runChatCompletionAndHandleFunctionCall(newMessages);
-  };
+  }, [isLoading, messages, setMessages, runChatCompletionAndHandleFunctionCall]);
 
   const stop = (): void => {
-    abortControllerRef.current?.abort();
+    chatAbortControllerRef.current?.abort("Stop was called");
   };
 
   return {
@@ -479,3 +585,32 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
     runChatCompletion: () => runChatCompletionRef.current!(messages),
   };
 }
+
+function constructFinalMessages(
+  syncedMessages: Message[],
+  previousMessages: Message[],
+  newMessages: Message[],
+): Message[] {
+  const finalMessages =
+    syncedMessages.length > 0 ? [...syncedMessages] : [...previousMessages, ...newMessages];
+
+  if (syncedMessages.length > 0) {
+    const messagesWithAgentState = [...previousMessages, ...newMessages];
+
+    let previousMessageId: string | undefined = undefined;
+
+    for (const message of messagesWithAgentState) {
+      if (message.isAgentStateMessage()) {
+        // insert this message into finalMessages after the position of previousMessageId
+        const index = finalMessages.findIndex((msg) => msg.id === previousMessageId);
+        if (index !== -1) {
+          finalMessages.splice(index + 1, 0, message);
+        }
+      }
+
+      previousMessageId = message.id;
+    }
+  }
+
+  return finalMessages;
+}