Load agent state (CopilotKit#1251)

Author: mme

CopilotKit/.changeset/loud-pillows-watch.md

@@ -0,0 +1,7 @@
+---
+"@copilotkit/react-core": patch
+"@copilotkit/runtime-client-gql": patch
+"@copilotkit/runtime": patch
+---
+
+- Load the previous state of an agent if `threadId` is provided to CopilotKit, including all messages

CopilotKit/packages/react-core/src/components/copilot-provider/copilot-messages.tsx

@@ -2,14 +2,36 @@
  * An internal context to separate the messages state (which is constantly changing) from the rest of CopilotKit context
  */
 
-import { useState } from "react";
+import { useEffect, useState } from "react";
 import { CopilotMessagesContext } from "../../context/copilot-messages-context";
-import { Message } from "@copilotkit/runtime-client-gql";
+import { loadMessagesFromJsonRepresentation, Message } from "@copilotkit/runtime-client-gql";
 import { CopilotKitProps } from "./copilotkit-props";
+import { useCopilotContext } from "../../context/copilot-context";
 
 export function CopilotMessages({ children, ...props }: CopilotKitProps) {
   const [messages, setMessages] = useState<Message[]>([]);
 
+  const { threadId, agentSession, runtimeClient } = useCopilotContext();
+
+  useEffect(() => {
+    if (agentSession?.agentName) {
+      // reload messages
+      const fetchAgentState = async () => {
+        const result = await runtimeClient.loadAgentState({
+          threadId,
+          agentName: agentSession.agentName,
+        });
+        if (result.data?.loadAgentState?.threadExists) {
+          const messages = loadMessagesFromJsonRepresentation(
+            JSON.parse(result.data?.loadAgentState?.messages || "[]"),
+          );
+          setMessages(messages);
+        }
+      };
+      void fetchAgentState();
+    }
+  }, [threadId, agentSession?.agentName !== undefined]);
+
   return (
     <CopilotMessagesContext.Provider
       value={{

CopilotKit/packages/react-core/src/components/copilot-provider/copilotkit-props.tsx

@@ -94,4 +94,9 @@ export interface CopilotKitProps {
       onSignInComplete: (authState: AuthState) => void;
     }>;
   };
+
+  /**
+   * The thread id to use for the CopilotKit.
+   */
+  threadId?: string;
 }

CopilotKit/packages/react-core/src/components/copilot-provider/copilotkit.tsx

@@ -14,7 +14,7 @@
  * ```
  */
 
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState, SetStateAction } from "react";
 import {
   CopilotContext,
   CopilotApiConfig,
@@ -30,6 +30,7 @@ import {
   CopilotCloudConfig,
   FunctionCallHandler,
   COPILOT_CLOUD_PUBLIC_API_KEY_HEADER,
+  randomUUID,
 } from "@copilotkit/shared";
 
 import { FrontendAction } from "../../types/frontend-action";
@@ -42,7 +43,11 @@ import { ToastProvider } from "../toast/toast-provider";
 import { useCopilotRuntimeClient } from "../../hooks/use-copilot-runtime-client";
 import { shouldShowDevConsole } from "../../utils";
 import { CopilotErrorBoundary } from "../error-boundary/error-boundary";
-import { Agent } from "@copilotkit/runtime-client-gql";
+import {
+  Agent,
+  ExtensionsInput,
+  loadMessagesFromJsonRepresentation,
+} from "@copilotkit/runtime-client-gql";
 
 export function CopilotKit({ children, ...props }: CopilotKitProps) {
   const showDevConsole = props.showDevConsole === undefined ? "auto" : props.showDevConsole;
@@ -80,6 +85,7 @@ export function CopilotKitInternal({ children, ...props }: CopilotKitProps) {
   const [isLoading, setIsLoading] = useState(false);
   const [chatInstructions, setChatInstructions] = useState("");
   const [authStates, setAuthStates] = useState<Record<string, AuthState>>({});
+  const [extensions, setExtensions] = useState<ExtensionsInput>({});
 
   const {
     addElement: addDocument,
@@ -314,7 +320,25 @@ export function CopilotKitInternal({ children, ...props }: CopilotKitProps) {
   }
 
   const [agentSession, setAgentSession] = useState<AgentSession | null>(initialAgentSession);
-  const [threadId, setThreadId] = useState<string | null>(null);
+
+  const [internalThreadId, setInternalThreadId] = useState<string>(props.threadId || randomUUID());
+  const setThreadId = useCallback(
+    (value: SetStateAction<string>) => {
+      if (props.threadId) {
+        throw new Error("Cannot call setThreadId() when threadId is provided via props.");
+      }
+      setInternalThreadId(value);
+    },
+    [props.threadId],
+  );
+
+  // update the internal threadId if the props.threadId changes
+  useEffect(() => {
+    if (props.threadId !== undefined) {
+      setInternalThreadId(props.threadId);
+    }
+  }, [props.threadId]);
+
   const [runId, setRunId] = useState<string | null>(null);
 
   const chatAbortControllerRef = useRef<AbortController | null>(null);
@@ -356,7 +380,7 @@ export function CopilotKitInternal({ children, ...props }: CopilotKitProps) {
         runtimeClient,
         forwardedParameters: props.forwardedParameters || {},
         agentLock: props.agent || null,
-        threadId,
+        threadId: internalThreadId,
         setThreadId,
         runId,
         setRunId,
@@ -365,6 +389,8 @@ export function CopilotKitInternal({ children, ...props }: CopilotKitProps) {
         authConfig: props.authConfig,
         authStates,
         setAuthStates,
+        extensions,
+        setExtensions,
       }}
     >
       <CopilotMessages>{children}</CopilotMessages>

CopilotKit/packages/react-core/src/context/copilot-context.tsx

@@ -10,7 +10,11 @@ import { DocumentPointer } from "../types";
 import { CopilotChatSuggestionConfiguration } from "../types/chat-suggestion-configuration";
 import { CoAgentStateRender, CoAgentStateRenderProps } from "../types/coagent-action";
 import { CoagentState } from "../types/coagent-state";
-import { CopilotRuntimeClient, ForwardedParametersInput } from "@copilotkit/runtime-client-gql";
+import {
+  CopilotRuntimeClient,
+  ExtensionsInput,
+  ForwardedParametersInput,
+} from "@copilotkit/runtime-client-gql";
 import { Agent } from "@copilotkit/runtime-client-gql";
 
 /**
@@ -160,8 +164,8 @@ export interface CopilotContextParams {
 
   agentLock: string | null;
 
-  threadId: string | null;
-  setThreadId: React.Dispatch<React.SetStateAction<string | null>>;
+  threadId: string;
+  setThreadId: React.Dispatch<React.SetStateAction<string>>;
 
   runId: string | null;
   setRunId: React.Dispatch<React.SetStateAction<string | null>>;
@@ -193,6 +197,9 @@ export interface CopilotContextParams {
       onSignInComplete: (authState: AuthState) => void;
     }>;
   };
+
+  extensions: ExtensionsInput;
+  setExtensions: React.Dispatch<React.SetStateAction<ExtensionsInput>>;
 }
 
 const emptyCopilotContext: CopilotContextParams = {
@@ -248,12 +255,14 @@ const emptyCopilotContext: CopilotContextParams = {
   setAgentSession: () => {},
   forwardedParameters: {},
   agentLock: null,
-  threadId: null,
+  threadId: "",
   setThreadId: () => {},
   runId: null,
   setRunId: () => {},
   chatAbortControllerRef: { current: null },
   availableAgents: [],
+  extensions: {},
+  setExtensions: () => {},
 };
 
 export const CopilotContext = React.createContext<CopilotContextParams>(emptyCopilotContext);

CopilotKit/packages/react-core/src/hooks/use-chat.ts

@@ -19,6 +19,8 @@ import {
   CopilotRequestType,
   ForwardedParametersInput,
   loadMessagesFromJsonRepresentation,
+  ExtensionsInput,
+  CopilotRuntimeClient,
 } from "@copilotkit/runtime-client-gql";
 
 import { CopilotApiConfig } from "../context";
@@ -107,11 +109,11 @@ export type UseChatOptions = {
   /**
    * The current thread ID.
    */
-  threadId: string | null;
+  threadId: string;
   /**
    * set the current thread ID
    */
-  setThreadId: (threadId: string | null) => void;
+  setThreadId: (threadId: string) => void;
   /**
    * The current run ID.
    */
@@ -128,6 +130,14 @@ export type UseChatOptions = {
    * The agent lock.
    */
   agentLock: string | null;
+  /**
+   * The extensions.
+   */
+  extensions: ExtensionsInput;
+  /**
+   * The setState-powered method to update the extensions.
+   */
+  setExtensions: React.Dispatch<React.SetStateAction<ExtensionsInput>>;
 };
 
 export type UseChatHelpers = {
@@ -183,6 +193,8 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
     setRunId,
     chatAbortControllerRef,
     agentLock,
+    extensions,
+    setExtensions,
   } = options;
   const runChatCompletionRef = useRef<(previousMessages: Message[]) => Promise<Message[]>>();
   const addErrorToast = useErrorToast();
@@ -191,10 +203,11 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
   // 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 extensionsRef = useRef<ExtensionsInput>(extensions);
+  extensionsRef.current = extensions;
 
   const publicApiKey = copilotConfig.publicApiKey;
 
@@ -240,8 +253,9 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
               actions: processActionsForRuntimeRequest(actions),
               url: window.location.href,
             },
-            threadId: threadIdRef.current,
+            threadId: threadId,
             runId: runIdRef.current,
+            extensions: extensionsRef.current,
             messages: convertMessagesToGqlInput(filterAgentStateMessages(messagesWithContext)),
             ...(copilotConfig.cloud
               ? {
@@ -314,11 +328,17 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
             continue;
           }
 
-          threadIdRef.current = value.generateCopilotResponse.threadId || null;
           runIdRef.current = value.generateCopilotResponse.runId || null;
 
-          setThreadId(threadIdRef.current);
+          // in the output, graphql inserts __typename, which leads to an error when sending it along
+          // as input to the next request.
+          extensionsRef.current = CopilotRuntimeClient.removeGraphQLTypename(
+            value.generateCopilotResponse.extensions || {},
+          );
+
+          // setThreadId(threadIdRef.current);
           setRunId(runIdRef.current);
+          setExtensions(extensionsRef.current);
 
           messages = convertGqlOutputToMessages(
             filterAdjacentAgentStateMessages(value.generateCopilotResponse.messages),

CopilotKit/packages/react-core/src/hooks/use-coagent.ts

@@ -97,10 +97,11 @@ import {
 } from "../context";
 import { CoagentState } from "../types/coagent-state";
 import { useCopilotChat } from "./use-copilot-chat";
-import { Message } from "@copilotkit/runtime-client-gql";
+import { loadMessagesFromJsonRepresentation, Message } from "@copilotkit/runtime-client-gql";
 import { flushSync } from "react-dom";
 import { useAsyncCallback } from "../components/error-boundary/error-utils";
 import { useToast } from "../components/toast/toast-provider";
+import { useCopilotRuntimeClient } from "./use-copilot-runtime-client";
 
 interface WithInternalStateManagementAndInitial<T> {
   /**
@@ -230,7 +231,8 @@ export function useCoAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentRe
 
   const messagesContext = useCopilotMessagesContext();
   const context = { ...generalContext, ...messagesContext };
-  const { coagentStates, coagentStatesRef, setCoagentStatesWithRef } = context;
+  const { coagentStates, coagentStatesRef, setCoagentStatesWithRef, threadId, copilotApiConfig } =
+    context;
   const { appendMessage, runChatCompletion } = useCopilotChat();
 
   const getCoagentState = (coagentStates: Record<string, CoagentState>, name: string) => {
@@ -249,6 +251,12 @@ export function useCoAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentRe
     }
   };
 
+  const runtimeClient = useCopilotRuntimeClient({
+    url: copilotApiConfig.chatApiEndpoint,
+    publicApiKey: copilotApiConfig.publicApiKey,
+    credentials: copilotApiConfig.credentials,
+  });
+
   // if we manage state internally, we need to provide a function to set the state
   const setState = (newState: T | ((prevState: T | undefined) => T)) => {
     let coagentState: CoagentState = getCoagentState(coagentStatesRef.current || {}, name);
@@ -264,9 +272,27 @@ export function useCoAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentRe
     });
   };
 
-  const coagentState = getCoagentState(coagentStates, name);
+  useEffect(() => {
+    const fetchAgentState = async () => {
+      const result = await runtimeClient.loadAgentState({
+        threadId: threadId,
+        agentName: name,
+      });
+      if (
+        result.data?.loadAgentState?.threadExists &&
+        result.data?.loadAgentState?.state &&
+        result.data?.loadAgentState?.state != "{}"
+      ) {
+        const fetchedState = JSON.parse(result.data?.loadAgentState?.state);
+        isExternalStateManagement(options)
+          ? options.setState(fetchedState)
+          : setState(fetchedState);
+      }
+    };
+    void fetchAgentState();
+  }, [threadId]);
 
-  const state = isExternalStateManagement(options) ? options.state : coagentState.state;
+  const coagentState = getCoagentState(coagentStates, name);
 
   // Sync internal state with external state if state management is external
   useEffect(() => {

CopilotKit/packages/react-core/src/hooks/use-copilot-chat.ts

@@ -47,7 +47,7 @@ import { Message, Role, TextMessage } from "@copilotkit/runtime-client-gql";
 import { SystemMessageFunction } from "../types";
 import { useChat, AppendMessageOptions } from "./use-chat";
 import { defaultCopilotContextCategories } from "../components";
-import { CoAgentStateRenderHandlerArguments } from "@copilotkit/shared";
+import { CoAgentStateRenderHandlerArguments, randomId } from "@copilotkit/shared";
 import { useCopilotMessagesContext } from "../context";
 import { useAsyncCallback } from "../components/error-boundary/error-utils";
 
@@ -110,6 +110,8 @@ export function useCopilotChat({
     runId,
     setRunId,
     chatAbortControllerRef,
+    extensions,
+    setExtensions,
   } = useCopilotContext();
   const { messages, setMessages } = useCopilotMessagesContext();
 
@@ -175,6 +177,8 @@ export function useCopilotChat({
     setRunId,
     chatAbortControllerRef,
     agentLock,
+    extensions,
+    setExtensions,
   });
 
   // this is a workaround born out of a bug that Athena incessantly ran into.
@@ -227,7 +231,6 @@ export function useCopilotChat({
   const reset = useCallback(() => {
     latestStopFunc();
     setMessages([]);
-    setThreadId(null);
     setRunId(null);
     setCoagentStatesWithRef({});
     let initialAgentSession: AgentSession | null = null;