@copilotkit/react-core 1.1.1-feat-runtime-remote-actions.4 → 1.1.1-feat-runtime-remote-actions.5

package/src/components/copilot-provider/copilotkit.tsx

@@ -36,6 +36,7 @@ import {
   CopilotApiConfig,
   InChatRenderFunction,
   ChatComponentsCache,
+  AgentSession,
 } from "../../context/copilot-context";
 import useTree from "../../hooks/use-tree";
 import { CopilotChatSuggestionConfiguration, DocumentPointer } from "../../types";
@@ -50,7 +51,8 @@ import { AgentStateMessage, Message } from "@copilotkit/runtime-client-gql";
 import { FrontendAction } from "../../types/frontend-action";
 import useFlatCategoryStore from "../../hooks/use-flat-category-store";
 import { CopilotKitProps } from "./copilotkit-props";
-import { CopilotChatUI } from "../../hooks/use-copilot-chat-ui";
+import { CoagentAction } from "../../types/coagent-action";
+import { CoagentState } from "../../types/coagent-state";
 
 export function CopilotKit({ children, ...props }: CopilotKitProps) {
   // Compute all the functions and properties that we need to pass
@@ -65,7 +67,11 @@ export function CopilotKit({ children, ...props }: CopilotKitProps) {
   const chatApiEndpoint = props.runtimeUrl || COPILOT_CLOUD_CHAT_URL;
 
   const [actions, setActions] = useState<Record<string, FrontendAction<any>>>({});
-  const chatComponentsCache = useRef<ChatComponentsCache>({ actions: {}, chatUI: {} });
+  const [coagentActions, setCoagentActions] = useState<Record<string, CoagentAction<any>>>({});
+  const chatComponentsCache = useRef<ChatComponentsCache>({
+    actions: {},
+    coagentActions: {},
+  });
   const { addElement, removeElement, printTree } = useTree();
   const [messages, setMessages] = useState<Message[]>([]);
   const [isLoading, setIsLoading] = useState(false);
@@ -94,6 +100,23 @@ export function CopilotKit({ children, ...props }: CopilotKitProps) {
     });
   }, []);
 
+  const setCoagentAction = useCallback((id: string, action: CoagentAction<any>) => {
+    setCoagentActions((prevPoints) => {
+      return {
+        ...prevPoints,
+        [id]: action,
+      };
+    });
+  }, []);
+
+  const removeCoagentAction = useCallback((id: string) => {
+    setCoagentActions((prevPoints) => {
+      const newPoints = { ...prevPoints };
+      delete newPoints[id];
+      return newPoints;
+    });
+  }, []);
+
   const getContextString = useCallback(
     (documents: DocumentPointer[], categories: string[]) => {
       const documentsString = documents
@@ -208,9 +231,15 @@ export function CopilotKit({ children, ...props }: CopilotKitProps) {
     });
   };
 
-  const [agentStates, setAgentStates] = useState<Record<string, AgentStateMessage | null>>({});
+  const [coagentStates, setCoagentStates] = useState<Record<string, CoagentState>>({});
+  let initialAgentSession: AgentSession | null = null;
+  if (props.agent) {
+    initialAgentSession = {
+      agentName: props.agent,
+    };
+  }
 
-  const [chatUI, setChatUI] = useState<CopilotChatUI[]>([]);
+  const [agentSession, setAgentSession] = useState<AgentSession | null>(initialAgentSession);
 
   return (
     <CopilotContext.Provider
@@ -220,6 +249,9 @@ export function CopilotKit({ children, ...props }: CopilotKitProps) {
         getFunctionCallHandler,
         setAction,
         removeAction,
+        coagentActions,
+        setCoagentAction,
+        removeCoagentAction,
         getContextString,
         addContext,
         removeContext,
@@ -237,10 +269,10 @@ export function CopilotKit({ children, ...props }: CopilotKitProps) {
         chatInstructions,
         setChatInstructions,
         showDevConsole: props.showDevConsole === undefined ? "auto" : props.showDevConsole,
-        agentStates,
-        setAgentStates,
-        chatUI,
-        setChatUI,
+        coagentStates,
+        setCoagentStates,
+        agentSession,
+        setAgentSession,
       }}
     >
       {children}
@@ -263,7 +295,7 @@ function entryPointsToFunctionCallHandler(actions: FrontendAction<any>[]): Funct
       await new Promise<void>((resolve, reject) => {
         flushSync(async () => {
           try {
-            result = await action.handler(args);
+            result = await action.handler?.(args);
             resolve();
           } catch (error) {
             reject(error);

package/src/context/copilot-context.tsx

@@ -1,11 +1,12 @@
 import { CopilotCloudConfig, FunctionCallHandler } from "@copilotkit/shared";
-import { AgentStateMessage, Message } from "@copilotkit/runtime-client-gql";
+import { Message } from "@copilotkit/runtime-client-gql";
 import { ActionRenderProps, FrontendAction } from "../types/frontend-action";
 import React from "react";
 import { TreeNodeId } from "../hooks/use-tree";
 import { DocumentPointer } from "../types";
 import { CopilotChatSuggestionConfiguration } from "../types/chat-suggestion-configuration";
-import { CopilotChatUI } from "../hooks/use-copilot-chat-ui";
+import { CoagentAction, CoagentActionRenderProps } from "../types/coagent-action";
+import { CoagentState } from "../types/coagent-state";
 
 /**
  * Interface for the configuration of the Copilot API.
@@ -68,10 +69,19 @@ export interface CopilotApiConfig {
 }
 
 export type InChatRenderFunction = (props: ActionRenderProps<any>) => string | JSX.Element;
+export type CoagentInChatRenderFunction = (
+  props: CoagentActionRenderProps<any>,
+) => string | JSX.Element;
 
 export interface ChatComponentsCache {
   actions: Record<string, InChatRenderFunction | string>;
-  chatUI: Record<string, Function | string>;
+  coagentActions: Record<string, CoagentInChatRenderFunction | string>;
+}
+
+export interface AgentSession {
+  agentName: string;
+  threadId?: string;
+  nodeName?: string;
 }
 
 export interface CopilotContextParams {
@@ -79,6 +89,12 @@ export interface CopilotContextParams {
   actions: Record<string, FrontendAction<any>>;
   setAction: (id: string, action: FrontendAction<any>) => void;
   removeAction: (id: string) => void;
+
+  // coagent actions
+  coagentActions: Record<string, CoagentAction<any>>;
+  setCoagentAction: (id: string, action: CoagentAction<any>) => void;
+  removeCoagentAction: (id: string) => void;
+
   chatComponentsCache: React.RefObject<ChatComponentsCache>;
 
   getFunctionCallHandler: (
@@ -117,13 +133,11 @@ export interface CopilotContextParams {
 
   showDevConsole: boolean | "auto";
 
-  // agent states
-  agentStates: Record<string, AgentStateMessage | null>;
-  setAgentStates: React.Dispatch<React.SetStateAction<Record<string, AgentStateMessage | null>>>;
-
-  // chat ui
-  chatUI: CopilotChatUI[];
-  setChatUI: React.Dispatch<React.SetStateAction<CopilotChatUI[]>>;
+  // agents
+  coagentStates: Record<string, CoagentState>;
+  setCoagentStates: React.Dispatch<React.SetStateAction<Record<string, CoagentState>>>;
+  agentSession: AgentSession | null;
+  setAgentSession: React.Dispatch<React.SetStateAction<AgentSession | null>>;
 }
 
 const emptyCopilotContext: CopilotContextParams = {
@@ -131,7 +145,11 @@ const emptyCopilotContext: CopilotContextParams = {
   setAction: () => {},
   removeAction: () => {},
 
-  chatComponentsCache: { current: { actions: {}, chatUI: {} } },
+  coagentActions: {},
+  setCoagentAction: () => {},
+  removeCoagentAction: () => {},
+
+  chatComponentsCache: { current: { actions: {}, coagentActions: {} } },
   getContextString: (documents: DocumentPointer[], categories: string[]) =>
     returnAndThrowInDebug(""),
   addContext: () => "",
@@ -169,10 +187,11 @@ const emptyCopilotContext: CopilotContextParams = {
   addChatSuggestionConfiguration: () => {},
   removeChatSuggestionConfiguration: () => {},
   showDevConsole: "auto",
-  agentStates: {},
-  setAgentStates: () => {},
-  chatUI: [],
-  setChatUI: () => {},
+  coagentStates: {},
+  setCoagentStates: () => {},
+
+  agentSession: null,
+  setAgentSession: () => {},
 };
 
 export const CopilotContext = React.createContext<CopilotContextParams>(emptyCopilotContext);

package/src/hooks/index.ts

@@ -3,8 +3,8 @@ export type { UseCopilotChatOptions } from "./use-copilot-chat";
 export type { UseCopilotChatReturn } from "./use-copilot-chat";
 
 export { useCopilotAction } from "./use-copilot-action";
+export { useCoagentAction } from "./use-coagent-action";
 export { useMakeCopilotDocumentReadable } from "./use-make-copilot-document-readable";
 export { type UseChatHelpers } from "./use-chat";
 export { useCopilotReadable } from "./use-copilot-readable";
-export { useAgentState } from "./use-agent-state";
-export { useCopilotChatUI } from "./use-copilot-chat-ui";
+export { useCoagent } from "./use-coagent";

package/src/hooks/use-chat.ts

@@ -2,8 +2,8 @@ import { useRef } from "react";
 import {
   FunctionCallHandler,
   COPILOT_CLOUD_PUBLIC_API_KEY_HEADER,
-  Action,
   actionParametersToJsonSchema,
+  CoagentActionHandler,
 } from "@copilotkit/shared";
 import {
   Message,
@@ -21,6 +21,9 @@ import {
 } from "@copilotkit/runtime-client-gql";
 
 import { CopilotApiConfig } from "../context";
+import { FrontendAction } from "../types/frontend-action";
+import { CoagentState } from "../types/coagent-state";
+import { AgentSession } from "../context/copilot-context";
 
 export type UseChatOptions = {
   /**
@@ -33,10 +36,16 @@ export type UseChatOptions = {
    * automatically to the API and will be used to update the chat.
    */
   onFunctionCall?: FunctionCallHandler;
+
+  /**
+   * Callback function to be called when a coagent action is received.
+   */
+  onCoagentAction?: CoagentActionHandler;
+
   /**
    * Function definitions to be sent to the API.
    */
-  actions: Action[];
+  actions: FrontendAction<any>[];
 
   /**
    * The CopilotKit API configuration.
@@ -67,10 +76,25 @@ export type UseChatOptions = {
    */
   setIsLoading: React.Dispatch<React.SetStateAction<boolean>>;
 
+  /**
+   * The current list of coagent states.
+   */
+  coagentStates: Record<string, CoagentState>;
+
   /**
    * setState-powered method to update the agent states
    */
-  setAgentStates: React.Dispatch<React.SetStateAction<Record<string, AgentStateMessage | null>>>;
+  setCoagentStates: React.Dispatch<React.SetStateAction<Record<string, CoagentState>>>;
+
+  /**
+   * The current agent session.
+   */
+  agentSession: AgentSession | null;
+
+  /**
+   * setState-powered method to update the agent session
+   */
+  setAgentSession: React.Dispatch<React.SetStateAction<AgentSession | null>>;
 };
 
 export type UseChatHelpers = {
@@ -92,12 +116,6 @@ export type UseChatHelpers = {
   stop: () => void;
 };
 
-interface AgentSession {
-  threadId: string;
-  agentName: string;
-  nodeName: string;
-}
-
 export function useChat(options: UseChatOptions): UseChatHelpers {
   const {
     messages,
@@ -109,13 +127,16 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
     isLoading,
     actions,
     onFunctionCall,
-    setAgentStates,
+    onCoagentAction,
+    setCoagentStates,
+    coagentStates,
+    agentSession,
+    setAgentSession,
   } = options;
 
   const abortControllerRef = useRef<AbortController>();
   const threadIdRef = useRef<string | null>(null);
   const runIdRef = useRef<string | null>(null);
-  const agentSessionRef = useRef<AgentSession | null>(null);
 
   const publicApiKey = copilotConfig.publicApiKey;
 
@@ -185,15 +206,15 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
           metadata: {
             requestType: CopilotRequestType.Chat,
           },
-          ...(agentSessionRef.current
+          ...(agentSession
             ? {
-                agentSession: {
-                  threadId: agentSessionRef.current.threadId,
-                  agentName: agentSessionRef.current.agentName,
-                  nodeName: agentSessionRef.current.nodeName,
-                },
+                agentSession,
               }
             : {}),
+          agentStates: Object.values(coagentStates).map((state) => ({
+            agentName: state.name,
+            state: JSON.stringify(state.state),
+          })),
         },
         properties: copilotConfig.properties,
         signal: abortControllerRef.current?.signal,
@@ -205,7 +226,8 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
 
     const reader = stream.getReader();
 
-    let results: { [id: string]: string } = {};
+    let actionResults: { [id: string]: string } = {};
+    let executedCoagentActions: string[] = [];
 
     try {
       while (true) {
@@ -250,22 +272,25 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
 
             if (message instanceof AgentStateMessage) {
               if (message.running) {
-                setAgentStates((prevAgentStates) => ({
+                setCoagentStates((prevAgentStates) => ({
                   ...prevAgentStates,
-                  [message.agentName]: message,
+                  [message.agentName]: {
+                    name: message.agentName,
+                    state: message.state,
+                    running: message.running,
+                    active: message.active,
+                    threadId: message.threadId,
+                    nodeName: message.nodeName,
+                    runId: message.runId,
+                  },
                 }));
-                agentSessionRef.current = {
+                setAgentSession({
                   threadId: message.threadId,
                   agentName: message.agentName,
                   nodeName: message.nodeName,
-                };
-              } else {
-                setAgentStates((prevAgentStates) => {
-                  const newAgentStates = { ...prevAgentStates };
-                  delete newAgentStates[message.agentName];
-                  return newAgentStates;
                 });
-                agentSessionRef.current = null;
+              } else {
+                setAgentSession(null);
               }
             }
 
@@ -276,7 +301,7 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
               message.scope === "client" &&
               onFunctionCall
             ) {
-              if (!(message.id in results)) {
+              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;
@@ -287,18 +312,38 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
                   name: message.name,
                   args: message.arguments,
                 });
-                results[message.id] = result;
+                actionResults[message.id] = result;
               }
 
               // add the result message
               newMessages.push(
                 new ResultMessage({
-                  result: ResultMessage.encodeResult(results[message.id]),
+                  result: ResultMessage.encodeResult(actionResults[message.id]),
                   actionExecutionId: message.id,
                   actionName: message.name,
                 }),
               );
             }
+
+            // execute coagent actions
+            if (
+              message instanceof AgentStateMessage &&
+              !message.active &&
+              !executedCoagentActions.includes(message.id) &&
+              onCoagentAction
+            ) {
+              // 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 onCoagentAction({
+                name: message.agentName,
+                nodeName: message.nodeName,
+                state: message.state,
+              });
+              executedCoagentActions.push(message.id);
+            }
           }
         }
 
@@ -308,11 +353,16 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
           const filteredMessages = [...previousMessages, ...newMessages].reduce(
             (acc: Message[], message: Message) => {
               if (
-                message instanceof AgentStateMessage && // Check if the current message is an AgentStateMessage
-                acc.length > 0 && // Ensure there is at least one message in the accumulator
-                acc[acc.length - 1] instanceof AgentStateMessage && // Check if the last message in the accumulator is also an AgentStateMessage
-                (acc[acc.length - 1] as AgentStateMessage).agentName === message.agentName && // Check if the agentName is the same
-                (acc[acc.length - 1] as AgentStateMessage).nodeName === message.nodeName // Check if the nodeName is the same
+                // If the current message is an AgentStateMessage
+                message instanceof AgentStateMessage &&
+                // And there is at least one message in the accumulator
+                acc.length > 0 &&
+                // And the last message in the accumulator is also an AgentStateMessage
+                acc[acc.length - 1] instanceof AgentStateMessage &&
+                // And the agentName, nodeName, and runId are the same
+                (acc[acc.length - 1] as AgentStateMessage).agentName === message.agentName &&
+                (acc[acc.length - 1] as AgentStateMessage).nodeName === message.nodeName &&
+                (acc[acc.length - 1] as AgentStateMessage).runId === message.runId
               ) {
                 // If the conditions are met, replace the last message in the accumulator with the current message
                 acc[acc.length - 1] = message;
@@ -320,7 +370,7 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
                 // Otherwise, add the current message to the accumulator
                 acc.push(message);
               }
-              return acc; // Return the accumulator for the next iteration
+              return acc;
             },
             [],
           );
@@ -332,7 +382,7 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
 
       if (
         // if we have client side results
-        Object.values(results).length ||
+        Object.values(actionResults).length ||
         // or the last message we received is a result
         (newMessages.length && newMessages[newMessages.length - 1] instanceof ResultMessage)
       ) {

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

@@ -0,0 +1,44 @@
+import { useRef, useContext, useEffect } from "react";
+import { CopilotContext } from "../context/copilot-context";
+import { randomId } from "@copilotkit/shared";
+import { CoagentAction } from "../types/coagent-action";
+
+// We implement useCoagentAction dependency handling so that
+// the developer has the option to not provide any dependencies.
+// see useCopilotAction for more details about this approach.
+export function useCoagentAction<T = any>(action: CoagentAction<T>, dependencies?: any[]): void {
+  const { setCoagentAction, removeCoagentAction, coagentActions, chatComponentsCache } =
+    useContext(CopilotContext);
+  const idRef = useRef<string>(randomId());
+
+  const key = `${action.name}-${action.nodeName || "global"}`;
+
+  if (dependencies === undefined) {
+    if (coagentActions[idRef.current]) {
+      coagentActions[idRef.current].handler = action.handler as any;
+      if (typeof action.render === "function") {
+        if (chatComponentsCache.current !== null) {
+          chatComponentsCache.current.coagentActions[key] = action.render;
+        }
+      }
+    }
+  }
+
+  useEffect(() => {
+    setCoagentAction(idRef.current, action as any);
+    if (chatComponentsCache.current !== null && action.render !== undefined) {
+      chatComponentsCache.current.coagentActions[key] = action.render;
+    }
+    return () => {
+      removeCoagentAction(idRef.current);
+    };
+  }, [
+    setCoagentAction,
+    removeCoagentAction,
+    action.name,
+    // include render only if it's a string
+    typeof action.render === "string" ? action.render : undefined,
+    // dependencies set by the developer
+    ...(dependencies || []),
+  ]);
+}

package/src/hooks/use-coagent.ts

@@ -1,40 +1,28 @@
-import { useState } from "react";
+import { useEffect } from "react";
+import { useCopilotContext } from "../context";
+import { CoagentState } from "../types/coagent-state";
 
-export function useCoagent<T = any>(options: UseCoagentOptions<T>): UseCoagentReturnType<T> {
-  throw new Error("Not implemented");
+interface WithInternalStateManagementAndInitial<T> {
+  name: string;
+  initialState: T;
 }
 
-const [appState, setAppState] = useState({
-  story: "dasfsdf",
-});
-
-const [story, setStory] = useState({
-  story: "dasfsdf",
-});
-
-const { nodeName } = useCoagent({
-  name: "myAgent",
-});
-
-// what do we need from useCoagent?
-// - nodeName
-// - threadId
-// - running
-// - state
-
-export function useCopilotAction<T = any>(coagentAction: CoagentAction<T>) {
-  throw new Error("Not implemented");
+interface WithInternalStateManagement {
+  name: string;
+  initialState?: any; // Optional initialState with default type any
 }
 
-// <CopilotKit agent="lockedInAgentName" />
-
-export interface UseCoagentOptions<T> {
+interface WithExternalStateManagement<T> {
   name: string;
-  initialState?: T;
-  state?: T;
-  setState?: (newState: T | ((prevState: T | undefined) => T)) => void;
+  state: T;
+  setState: (newState: T | ((prevState: T | undefined) => T)) => void;
 }
 
+type UseCoagentOptions<T> =
+  | WithInternalStateManagementAndInitial<T>
+  | WithInternalStateManagement
+  | WithExternalStateManagement<T>;
+
 export interface UseCoagentReturnType<T> {
   name: string;
   nodeName?: string;
@@ -46,15 +34,100 @@ export interface UseCoagentReturnType<T> {
   stop: () => void;
 }
 
-export interface CoagentActionRenderProps<T> {
-  status: "executing" | "complete" | "in_progress";
-  args: T;
-  result: any;
+export function useCoagent<T = any>(options: UseCoagentOptions<T>): UseCoagentReturnType<T> {
+  const isExternalStateManagement = (
+    options: UseCoagentOptions<T>,
+  ): options is WithExternalStateManagement<T> => {
+    return "state" in options && "setState" in options;
+  };
+
+  const { name } = options;
+
+  const isInternalStateManagementWithInitial = (
+    options: UseCoagentOptions<T>,
+  ): options is WithInternalStateManagementAndInitial<T> => {
+    return "initialState" in options;
+  };
+
+  const { coagentStates, setCoagentStates } = useCopilotContext();
+
+  const getCoagentState = (coagentStates: Record<string, CoagentState>, name: string) => {
+    if (coagentStates[name]) {
+      return coagentStates[name];
+    } else {
+      return {
+        name,
+        state: isInternalStateManagementWithInitial(options) ? options.initialState : {},
+        running: false,
+        active: false,
+        threadId: undefined,
+        nodeName: undefined,
+        runId: undefined,
+      };
+    }
+  };
+
+  // if we manage state internally, we need to provide a function to set the state
+  const setState = (newState: T | ((prevState: T | undefined) => T)) => {
+    setCoagentStates((prevAgentStates) => {
+      let coagentState: CoagentState = getCoagentState(prevAgentStates, name);
+
+      const updatedState =
+        typeof newState === "function" ? (newState as Function)(coagentState.state) : newState;
+
+      return {
+        ...prevAgentStates,
+        [name]: {
+          ...coagentState,
+          state: updatedState,
+        },
+      };
+    });
+  };
+
+  const coagentState = getCoagentState(coagentStates, name);
+
+  const state = isExternalStateManagement(options) ? options.state : coagentState.state;
+
+  // Sync internal state with external state if state management is external
+  useEffect(() => {
+    if (isExternalStateManagement(options)) {
+      setState(options.state);
+    } else if (coagentStates[name] === undefined) {
+      setState(options.initialState === undefined ? {} : options.initialState);
+    }
+  }, [isExternalStateManagement(options) ? JSON.stringify(options.state) : undefined]);
+
+  // Return the state and setState function
+  return {
+    name,
+    nodeName: coagentState.nodeName,
+    state,
+    setState,
+    running: coagentState.running,
+    start: () => {
+      startAgent(name);
+    },
+    stop: () => {
+      stopAgent(name);
+    },
+  };
+}
+
+function startAgent(name: string) {
+  const { setAgentSession } = useCopilotContext();
+  setAgentSession({
+    agentName: name,
+  });
 }
 
-export interface CoagentAction<T> {
-  name: string;
-  nodeName?: string;
-  handler?: (args: T) => any | Promise<any>;
-  render?: (props: CoagentActionRenderProps<T>) => string | React.ReactElement;
+function stopAgent(name: string) {
+  const { agentSession, setAgentSession } = useCopilotContext();
+  if (agentSession && agentSession.agentName === name) {
+    setAgentSession(null);
+  } else {
+    console.warn(`No agent session found for ${name}`);
+  }
 }
+
+// <CopilotKit agent="lockedInAgentName" />