implement renderAndWait

Author: mme

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

@@ -288,6 +288,10 @@ export function useChat(options: UseChatOptions): UseChatHelpers {
                 }
                 // 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 result = await onFunctionCall({
                     messages: previousMessages,
                     name: message.name,

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

@@ -1,5 +1,5 @@
 import { useRef, useEffect } from "react";
-import { FrontendAction } from "../types/frontend-action";
+import { FrontendAction, ActionRenderProps, ActionRenderPropsWait } from "../types/frontend-action";
 import { useCopilotContext } from "../context/copilot-context";
 import { Parameter, randomId } from "@copilotkit/shared";
 
@@ -18,6 +18,44 @@ export function useCopilotAction<const T extends Parameter[] | [] = []>(
 ): void {
   const { setAction, removeAction, actions, chatComponentsCache } = useCopilotContext();
   const idRef = useRef<string>(randomId());
+  const renderAndWaitRef = useRef<RenderAndWait | null>(null);
+
+  // clone the action to avoid mutating the original object
+  action = { ...action };
+
+  // If the developer provides a renderAndWait function, we transform the action
+  // to use a promise internally, so that we can treat it like a normal action.
+  if (action.renderAndWait) {
+    const renderAndWait = action.renderAndWait!;
+
+    // remove the renderAndWait function from the action
+    action.renderAndWait = undefined;
+
+    // add a handler that will be called when the action is executed
+    action.handler = (async () => {
+      // we create a new promise when the handler is called
+      let resolve: (result: any) => void;
+      let reject: (error: any) => void;
+      const promise = new Promise<any>((resolvePromise, rejectPromise) => {
+        resolve = resolvePromise;
+        reject = rejectPromise;
+      });
+      renderAndWaitRef.current = { promise, resolve: resolve!, reject: reject! };
+      // then we await the promise (it will be resolved in the original renderAndWait function)
+      return await promise;
+    }) as any;
+
+    // add a render function that will be called when the action is rendered
+    action.render = ((props: ActionRenderProps<any>): React.ReactElement => {
+      const waitProps: ActionRenderPropsWait<any> = {
+        status: props.status as any,
+        args: props.args,
+        result: props.result,
+        handler: props.status === "executing" ? renderAndWaitRef.current!.resolve : undefined,
+      };
+      return renderAndWait(waitProps);
+    }) as any;
+  }
 
   // If the developer doesn't provide dependencies, we assume they want to
   // update handler and render function when the action object changes.
@@ -62,6 +100,52 @@ export function useCopilotAction<const T extends Parameter[] | [] = []>(
   ]);
 }
 
+interface RenderAndWait {
+  promise: Promise<any>;
+  resolve: (result: any) => void;
+  reject: (error: any) => void;
+}
+
+function transformRenderAndWaitAction<T extends Parameter[] | []>(
+  action: FrontendAction<T>,
+): FrontendAction<T> {
+  console.log("transformRenderAndWaitAction", action);
+
+  let resolvePromise: (result: any) => void;
+  const promise = new Promise<any>((resolve, reject) => {
+    resolvePromise = function (result: any) {
+      console.log("resolvePromise", result);
+      resolve(result);
+    };
+  });
+  let handler = async () => {
+    try {
+      console.log("awaiting promise");
+      const result = await promise;
+      console.log("result", result);
+      return result;
+    } catch (error) {
+      console.error("Error in handler:", error);
+      throw error;
+    }
+  };
+  const render = (props: ActionRenderProps<any>): React.ReactElement => {
+    const waitProps: ActionRenderPropsWait<any> = {
+      status: props.status as any,
+      args: props.args,
+      result: props.result,
+      handler: props.status === "executing" ? resolvePromise : undefined,
+    };
+    return action.renderAndWait!(waitProps);
+  };
+  return {
+    ...action,
+    render: render as any,
+    handler: handler as any,
+    renderAndWait: undefined,
+  };
+}
+
 // Usage Example:
 // useCopilotAction({
 //   name: "myAction",

CopilotKit/packages/react-core/src/types/frontend-action.ts

@@ -37,6 +37,48 @@ interface CompleteStateNoArgs<T extends Parameter[] | [] = []> {
   result: any;
 }
 
+interface InProgressStateWait<T extends Parameter[] | [] = []> {
+  status: "inProgress";
+  args: Partial<MappedParameterTypes<T>>;
+  handler: undefined;
+  result: undefined;
+}
+
+interface ExecutingStateWait<T extends Parameter[] | [] = []> {
+  status: "executing";
+  args: MappedParameterTypes<T>;
+  handler: (result: any) => void;
+  result: undefined;
+}
+
+interface CompleteStateWait<T extends Parameter[] | [] = []> {
+  status: "complete";
+  args: MappedParameterTypes<T>;
+  handler: undefined;
+  result: any;
+}
+
+interface InProgressStateNoArgsWait<T extends Parameter[] | [] = []> {
+  status: "inProgress";
+  args: Partial<MappedParameterTypes<T>>;
+  handler: undefined;
+  result: undefined;
+}
+
+interface ExecutingStateNoArgsWait<T extends Parameter[] | [] = []> {
+  status: "executing";
+  args: MappedParameterTypes<T>;
+  handler: (result: any) => void;
+  result: undefined;
+}
+
+interface CompleteStateNoArgsWait<T extends Parameter[] | [] = []> {
+  status: "complete";
+  args: MappedParameterTypes<T>;
+  handler: undefined;
+  result: any;
+}
+
 export type ActionRenderProps<T extends Parameter[] | [] = []> =
   | CompleteState<T>
   | ExecutingState<T>
@@ -47,13 +89,34 @@ export type ActionRenderPropsNoArgs<T extends Parameter[] | [] = []> =
   | ExecutingStateNoArgs<T>
   | InProgressStateNoArgs<T>;
 
+export type ActionRenderPropsWait<T extends Parameter[] | [] = []> =
+  | CompleteStateWait<T>
+  | ExecutingStateWait<T>
+  | InProgressStateWait<T>;
+
+export type ActionRenderPropsNoArgsWait<T extends Parameter[] | [] = []> =
+  | CompleteStateNoArgsWait<T>
+  | ExecutingStateNoArgsWait<T>
+  | InProgressStateNoArgsWait<T>;
+
 export type FrontendAction<T extends Parameter[] | [] = []> = Action<T> & {
   disabled?: boolean;
-  render?:
-    | string
-    | (T extends []
-        ? (props: ActionRenderPropsNoArgs<T>) => string | React.ReactElement
-        : (props: ActionRenderProps<T>) => string | React.ReactElement);
-};
+} & (
+    | {
+        render?:
+          | string
+          | (T extends []
+              ? (props: ActionRenderPropsNoArgs<T>) => string | React.ReactElement
+              : (props: ActionRenderProps<T>) => string | React.ReactElement);
+        renderAndWait?: never;
+      }
+    | {
+        render?: never;
+        renderAndWait: T extends []
+          ? (props: ActionRenderPropsNoArgsWait<T>) => React.ReactElement
+          : (props: ActionRenderPropsWait<T>) => React.ReactElement;
+        handler?: never;
+      }
+  );
 
 export type RenderFunctionStatus = ActionRenderProps<any>["status"];