@copilotkit/react-core 1.4.8 → 1.5.0-coagents-v0-3.0

package/src/hooks/use-chat.ts

@@ -18,10 +18,11 @@ 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";
@@ -98,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 = {
@@ -140,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;
 
@@ -180,8 +206,8 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
           role: Role.Assistant,
         }),
       ];
-      const abortController = new AbortController();
-      abortControllerRef.current = abortController;
+
+      chatAbortControllerRef.current = new AbortController();
 
       setMessages([...previousMessages, ...newMessages]);
 
@@ -189,34 +215,13 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
 
       const messagesWithContext = [systemMessage, ...(initialMessages || []), ...previousMessages];
 
+      const isAgentRun = agentSessionRef.current !== null;
+
       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,
-                  };
-                }),
+              actions: processActionsForRuntimeRequest(actions),
               url: window.location.href,
             },
             threadId: threadIdRef.current,
@@ -254,7 +259,7 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
             })),
           },
           properties: copilotConfig.properties,
-          signal: abortControllerRef.current?.signal,
+          signal: chatAbortControllerRef.current?.signal,
         }),
       );
 
@@ -263,10 +268,12 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
 
       const reader = stream.getReader();
 
-      let actionResults: { [id: string]: string } = {};
       let executedCoAgentStateRenders: string[] = [];
       let followUp: FrontendAction["followUp"] = undefined;
 
+      let messages: Message[] = [];
+      let syncedMessages: Message[] = [];
+
       try {
         while (true) {
           let done, value;
@@ -280,6 +287,9 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
           }
 
           if (done) {
+            if (chatAbortControllerRef.current.signal.aborted) {
+              return [];
+            }
             break;
           }
 
@@ -290,7 +300,10 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
           threadIdRef.current = value.generateCopilotResponse.threadId || null;
           runIdRef.current = value.generateCopilotResponse.runId || null;
 
-          const messages = convertGqlOutputToMessages(
+          setThreadId(threadIdRef.current);
+          setRunId(runIdRef.current);
+
+          messages = convertGqlOutputToMessages(
             filterAdjacentAgentStateMessages(value.generateCopilotResponse.messages),
           );
 
@@ -300,7 +313,7 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
 
           newMessages = [];
 
-          // request failed, display error message
+          // request failed, display error message and quit
           if (
             value.generateCopilotResponse.status?.__typename === "FailedResponseStatus" &&
             value.generateCopilotResponse.status.reason === "GUARDRAILS_VALIDATION_FAILED"
@@ -311,57 +324,16 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
                 content: value.generateCopilotResponse.status.details?.guardrailsReason || "",
               }),
             ];
+            setMessages([...previousMessages, ...newMessages]);
+            break;
           }
 
           // add messages to the chat
           else {
-            for (const message of messages) {
-              newMessages.push(message);
-              // execute regular action executions
-              if (
-                message.isActionExecutionMessage() &&
-                message.status.code !== MessageStatusCode.Pending &&
-                message.scope === "client" &&
-                onFunctionCall
-              ) {
-                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]);
-
-                    const action = actions.find((action) => action.name === message.name);
-
-                    if (action) {
-                      followUp = action.followUp;
-                    }
+            newMessages = [...messages];
 
-                    const result = await 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}`);
-                  }
-                }
-                // add the result message
-                newMessages.push(
-                  new ResultMessage({
-                    result: ResultMessage.encodeResult(actionResults[message.id]),
-                    actionExecutionId: message.id,
-                    actionName: message.name,
-                  }),
-                );
-              }
-              // execute coagent actions
+            for (const message of messages) {
+              // execute onCoAgentStateRender handler
               if (
                 message.isAgentStateMessage() &&
                 !message.active &&
@@ -387,6 +359,14 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
               .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]: {
@@ -416,14 +396,88 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
             setMessages([...previousMessages, ...newMessages]);
           }
         }
+        const finalMessages = constructFinalMessages(syncedMessages, previousMessages, newMessages);
+
+        let didExecuteAction = false;
+
+        // 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
+            ) {
+              lastMessages.unshift(message);
+            } else {
+              break;
+            }
+          }
+
+          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);
+
+            const action = actions.find((action) => action.name === message.name);
+
+            if (action) {
+              followUp = action.followUp;
+              let result: any;
+              try {
+                result = await Promise.race([
+                  onFunctionCall({
+                    messages: previousMessages,
+                    name: message.name,
+                    args: message.arguments,
+                  }),
+                  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}`);
+              }
+              didExecuteAction = true;
+              const messageIndex = finalMessages.findIndex((msg) => msg.id === message.id);
+              finalMessages.splice(
+                messageIndex + 1,
+                0,
+                new ResultMessage({
+                  id: "result-" + message.id,
+                  result: ResultMessage.encodeResult(result),
+                  actionExecutionId: message.id,
+                  actionName: message.name,
+                }),
+              );
+            }
+          }
+
+          setMessages(finalMessages);
+        }
 
         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()))
+          // 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
 
@@ -431,7 +485,32 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
           // - tried using react-dom's flushSync, but it did not work
           await new Promise((resolve) => setTimeout(resolve, 10));
 
-          return await runChatCompletionRef.current!([...previousMessages, ...newMessages]);
+          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();
         }
@@ -496,7 +575,7 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
   }, [isLoading, messages, setMessages, runChatCompletionAndHandleFunctionCall]);
 
   const stop = (): void => {
-    abortControllerRef.current?.abort();
+    chatAbortControllerRef.current?.abort("Stop was called");
   };
 
   return {
@@ -506,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;
+}

package/src/hooks/use-coagent.ts

@@ -265,7 +265,11 @@ export function useCoAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentRe
     } else if (coagentStates[name] === undefined) {
       setState(options.initialState === undefined ? {} : options.initialState);
     }
-  }, [isExternalStateManagement(options) ? JSON.stringify(options.state) : undefined]);
+  }, [
+    isExternalStateManagement(options) ? JSON.stringify(options.state) : undefined,
+    // reset initialstate on reset
+    coagentStates[name] === undefined,
+  ]);
 
   const runAgentCallback = useAsyncCallback(
     async (hint?: HintFunction) => {
@@ -288,23 +292,36 @@ export function useCoAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentRe
   };
 }
 
-function startAgent(name: string, context: CopilotContextParams) {
+export function startAgent(name: string, context: CopilotContextParams) {
   const { setAgentSession } = context;
   setAgentSession({
     agentName: name,
   });
 }
 
-function stopAgent(name: string, context: CopilotContextParams) {
+export function stopAgent(name: string, context: CopilotContextParams) {
   const { agentSession, setAgentSession } = context;
   if (agentSession && agentSession.agentName === name) {
     setAgentSession(null);
+    context.setCoagentStates((prevAgentStates) => {
+      return {
+        ...prevAgentStates,
+        [name]: {
+          ...prevAgentStates[name],
+          running: false,
+          active: false,
+          threadId: undefined,
+          nodeName: undefined,
+          runId: undefined,
+        },
+      };
+    });
   } else {
     console.warn(`No agent session found for ${name}`);
   }
 }
 
-async function runAgent(
+export async function runAgent(
   name: string,
   context: CopilotContextParams & CopilotMessagesContextParams,
   appendMessage: (message: Message) => Promise<void>,

package/src/hooks/use-copilot-action.ts

@@ -71,6 +71,15 @@
  *     );
  *   },
  * });
+ *
+ * @example
+ * // Catch all action allows you to render actions that are not defined in the frontend
+ * useCopilotAction({
+ *   name: "*",
+ *   render: ({ name, args, status, result, handler, respond }) => {
+ *     return <div>Rendering action: {name}</div>;
+ *   },
+ * });
  */
 
 /**
@@ -129,6 +138,7 @@ import {
   ActionRenderProps,
   ActionRenderPropsNoArgsWait,
   ActionRenderPropsWait,
+  CatchAllFrontendAction,
   FrontendAction,
 } from "../types/frontend-action";
 
@@ -142,7 +152,7 @@ import {
 // useCallback, useMemo or other memoization techniques are not suitable here,
 // because they will cause a infinite rerender loop.
 export function useCopilotAction<const T extends Parameter[] | [] = []>(
-  action: FrontendAction<T>,
+  action: FrontendAction<T> | CatchAllFrontendAction,
   dependencies?: any[],
 ): void {
   const { setAction, removeAction, actions, chatComponentsCache } = useCopilotContext();
@@ -152,9 +162,14 @@ export function useCopilotAction<const T extends Parameter[] | [] = []>(
   // clone the action to avoid mutating the original object
   action = { ...action };
 
-  // If the developer provides a renderAndWait function, we transform the action
+  // If the developer provides a renderAndWaitForResponse function, we transform the action
   // to use a promise internally, so that we can treat it like a normal action.
-  if (action.renderAndWait || action.renderAndWaitForResponse) {
+  if (
+    // renderAndWaitForResponse is not available for catch all actions
+    isFrontendAction(action) &&
+    // check if renderAndWaitForResponse is set
+    (action.renderAndWait || action.renderAndWaitForResponse)
+  ) {
     const renderAndWait = action.renderAndWait || action.renderAndWaitForResponse;
     // remove the renderAndWait function from the action
     action.renderAndWait = undefined;
@@ -212,10 +227,15 @@ export function useCopilotAction<const T extends Parameter[] | [] = []>(
   // This ensures that any captured variables in the handler are up to date.
   if (dependencies === undefined) {
     if (actions[idRef.current]) {
-      actions[idRef.current].handler = action.handler as any;
+      // catch all actions don't have a handler
+      if (isFrontendAction(action)) {
+        actions[idRef.current].handler = action.handler as any;
+      }
       if (typeof action.render === "function") {
         if (chatComponentsCache.current !== null) {
-          chatComponentsCache.current.actions[action.name] = action.render;
+          // TODO: using as any here because the type definitions are getting to tricky
+          // not wasting time on this now - we know the types are compatible
+          chatComponentsCache.current.actions[action.name] = action.render as any;
         }
       }
     }
@@ -224,23 +244,25 @@ export function useCopilotAction<const T extends Parameter[] | [] = []>(
   useEffect(() => {
     setAction(idRef.current, action as any);
     if (chatComponentsCache.current !== null && action.render !== undefined) {
-      chatComponentsCache.current.actions[action.name] = action.render;
+      // see comment about type safety above
+      chatComponentsCache.current.actions[action.name] = action.render as any;
     }
     return () => {
       // NOTE: For now, we don't remove the chatComponentsCache entry when the action is removed.
       // This is because we currently don't have access to the messages array in CopilotContext.
+      // UPDATE: We now have access, we should remove the entry if not referenced by any message.
       removeAction(idRef.current);
     };
   }, [
     setAction,
     removeAction,
-    action.description,
+    isFrontendAction(action) ? action.description : undefined,
     action.name,
-    action.disabled,
-    action.available,
+    isFrontendAction(action) ? action.disabled : undefined,
+    isFrontendAction(action) ? action.available : undefined,
     // This should be faster than deep equality checking
     // In addition, all major JS engines guarantee the order of object keys
-    JSON.stringify(action.parameters),
+    JSON.stringify(isFrontendAction(action) ? action.parameters : []),
     // include render only if it's a string
     typeof action.render === "string" ? action.render : undefined,
     // dependencies set by the developer
@@ -248,6 +270,12 @@ export function useCopilotAction<const T extends Parameter[] | [] = []>(
   ]);
 }
 
+function isFrontendAction<T extends Parameter[]>(
+  action: FrontendAction<T> | CatchAllFrontendAction,
+): action is FrontendAction<T> {
+  return action.name !== "*";
+}
+
 interface RenderAndWaitForResponse {
   promise: Promise<any>;
   resolve: (result: any) => void;

package/src/hooks/use-copilot-chat.ts

@@ -39,7 +39,7 @@
  * ```
  */
 import { useRef, useEffect, useCallback } from "react";
-import { useCopilotContext } from "../context/copilot-context";
+import { AgentSession, useCopilotContext } from "../context/copilot-context";
 import { Message, Role, TextMessage } from "@copilotkit/runtime-client-gql";
 import { SystemMessageFunction } from "../types";
 import { useChat } from "./use-chat";
@@ -79,6 +79,7 @@ export interface UseCopilotChatReturn {
   deleteMessage: (messageId: string) => void;
   reloadMessages: () => Promise<void>;
   stopGeneration: () => void;
+  reset: () => void;
   isLoading: boolean;
   runChatCompletion: () => Promise<Message[]>;
 }
@@ -100,6 +101,12 @@ export function useCopilotChat({
     coAgentStateRenders,
     agentSession,
     setAgentSession,
+    agentLock,
+    threadId,
+    setThreadId,
+    runId,
+    setRunId,
+    chatAbortControllerRef,
   } = useCopilotContext();
   const { messages, setMessages } = useCopilotMessagesContext();
 
@@ -158,11 +165,16 @@ export function useCopilotChat({
     setCoagentStatesWithRef,
     agentSession,
     setAgentSession,
+    threadId,
+    setThreadId,
+    runId,
+    setRunId,
+    chatAbortControllerRef,
   });
 
-  // this is a workaround born out of a bug that Athena insessently ran into.
+  // this is a workaround born out of a bug that Athena incessantly ran into.
   // We could not find the origin of the bug, however, it was clear that an outdated version of the append function was being used somehow --
-  // it referecned the old state of the messages array, and not the latest one.
+  // it referenced the old state of the messages array, and not the latest one.
   //
   // We want to make copilotkit as abuse-proof as possible, so we are adding this workaround to ensure that the latest version of the append function is always used.
   //
@@ -207,12 +219,40 @@ export function useCopilotChat({
     return await latestRunChatCompletion.current!();
   }, [latestRunChatCompletion]);
 
+  const reset = useCallback(() => {
+    latestStopFunc();
+    setMessages([]);
+    setThreadId(null);
+    setRunId(null);
+    setCoagentStatesWithRef({});
+    let initialAgentSession: AgentSession | null = null;
+    if (agentLock) {
+      initialAgentSession = {
+        agentName: agentLock,
+      };
+    }
+    setAgentSession(initialAgentSession);
+  }, [
+    latestStopFunc,
+    setMessages,
+    setThreadId,
+    setCoagentStatesWithRef,
+    setAgentSession,
+    agentLock,
+  ]);
+
+  const latestReset = useUpdatedRef(reset);
+  const latestResetFunc = useCallback(() => {
+    return latestReset.current();
+  }, [latestReset]);
+
   return {
     visibleMessages: messages,
     appendMessage: latestAppendFunc,
     setMessages: latestSetMessagesFunc,
     reloadMessages: latestReloadFunc,
     stopGeneration: latestStopFunc,
+    reset: latestResetFunc,
     deleteMessage: latestDeleteFunc,
     runChatCompletion: latestRunChatCompletionFunc,
     isLoading,