@copilotkit/react-core 1.50.0-beta.0 → 1.50.0-beta.10

package/src/hooks/__tests__/use-coagent-config.test.ts

@@ -1,8 +1,43 @@
-import { renderHook } from "@testing-library/react";
+import { renderHook, waitFor } from "@testing-library/react";
 import { useCoAgent } from "../use-coagent";
+import type { AgentSubscriber } from "@ag-ui/client";
+
+// Mock functions for @copilotkitnext/react
+const mockSetState = jest.fn();
+const mockRunAgent = jest.fn();
+const mockAbortRun = jest.fn();
+const mockSubscribe = jest.fn();
+const mockSetProperties = jest.fn();
+
+// Store the last subscriber for triggering events
+let lastSubscriber: AgentSubscriber | null = null;
+
+const mockAgent = {
+  agentId: "test-agent",
+  state: { count: 0 },
+  isRunning: false,
+  threadId: "thread-123",
+  setState: mockSetState,
+  runAgent: mockRunAgent,
+  abortRun: mockAbortRun,
+  subscribe: mockSubscribe.mockImplementation((subscriber: AgentSubscriber) => {
+    lastSubscriber = subscriber;
+    return {
+      unsubscribe: jest.fn(),
+    };
+  }),
+};
 
-// Mock the dependencies
-const mockSetCoagentStatesWithRef = jest.fn();
+jest.mock("@copilotkitnext/react", () => ({
+  useAgent: jest.fn(() => ({ agent: mockAgent })),
+  useCopilotKit: jest.fn(() => ({
+    copilotkit: {
+      setProperties: mockSetProperties,
+    },
+  })),
+}));
+
+// Mock other dependencies
 const mockAppendMessage = jest.fn();
 const mockRunChatCompletion = jest.fn();
 
@@ -27,7 +62,6 @@ jest.mock("../../context", () => ({
     availableAgents: [],
     coagentStates: {},
     coagentStatesRef: { current: {} },
-    setCoagentStatesWithRef: mockSetCoagentStatesWithRef,
     threadId: "test-thread",
     copilotApiConfig: {
       headers: {},
@@ -61,14 +95,10 @@ jest.mock("../../components/copilot-provider/copilot-messages", () => ({
 describe("useCoAgent config synchronization", () => {
   beforeEach(() => {
     jest.clearAllMocks();
-    mockSetCoagentStatesWithRef.mockImplementation((updater) => {
-      if (typeof updater === "function") {
-        updater({});
-      }
-    });
+    lastSubscriber = null;
   });
 
-  it("should call setCoagentStatesWithRef when config changes", () => {
+  it("should call setProperties when config changes", async () => {
     const { rerender } = renderHook(
       ({ config }) =>
         useCoAgent({
@@ -82,16 +112,20 @@ describe("useCoAgent config synchronization", () => {
     );
 
     // Clear the initial calls
-    mockSetCoagentStatesWithRef.mockClear();
+    mockSetProperties.mockClear();
 
     // Change config
     rerender({ config: { configurable: { model: "gpt-4o" } } });
 
-    // Should have called setCoagentStatesWithRef with new config
-    expect(mockSetCoagentStatesWithRef).toHaveBeenCalledWith(expect.any(Function));
+    // Wait for effect to complete and verify setProperties was called
+    await waitFor(() => {
+      expect(mockSetProperties).toHaveBeenCalledWith({
+        configurable: { model: "gpt-4o" },
+      });
+    });
   });
 
-  it("should not call setCoagentStatesWithRef when config is unchanged", () => {
+  it("should not call setProperties when config is unchanged", () => {
     const config = { configurable: { model: "gpt-4" } };
 
     const { rerender } = renderHook(
@@ -107,16 +141,16 @@ describe("useCoAgent config synchronization", () => {
     );
 
     // Clear the initial calls
-    mockSetCoagentStatesWithRef.mockClear();
+    mockSetProperties.mockClear();
 
-    // Re-render with same config
+    // Re-render with same config reference
     rerender({ config });
 
-    // Should not have called setCoagentStatesWithRef
-    expect(mockSetCoagentStatesWithRef).not.toHaveBeenCalled();
+    // Should not have called setProperties
+    expect(mockSetProperties).not.toHaveBeenCalled();
   });
 
-  it("should handle backward compatibility with configurable prop", () => {
+  it("should handle backward compatibility with configurable prop", async () => {
     const { rerender } = renderHook(
       ({ configurable }) =>
         useCoAgent({
@@ -130,23 +164,32 @@ describe("useCoAgent config synchronization", () => {
     );
 
     // Clear the initial calls
-    mockSetCoagentStatesWithRef.mockClear();
+    mockSetProperties.mockClear();
 
     // Change configurable prop
     rerender({ configurable: { model: "gpt-4o" } });
 
-    // Should have called setCoagentStatesWithRef
-    expect(mockSetCoagentStatesWithRef).toHaveBeenCalledWith(expect.any(Function));
+    // Wait for effect to complete and verify setProperties was called
+    await waitFor(() => {
+      expect(mockSetProperties).toHaveBeenCalledWith({
+        configurable: { model: "gpt-4o" },
+      });
+    });
   });
 
-  it("should update config while preserving other state properties", () => {
-    let capturedUpdater: any = null;
+  it("should not call setProperties when both config and configurable are undefined", () => {
+    renderHook(() =>
+      useCoAgent({
+        name: "test-agent",
+        initialState: { count: 0 },
+      }),
+    );
 
-    mockSetCoagentStatesWithRef.mockImplementation((updater) => {
-      capturedUpdater = updater;
-      return updater;
-    });
+    // Should not have called setProperties
+    expect(mockSetProperties).not.toHaveBeenCalled();
+  });
 
+  it("should handle deeply nested config changes", async () => {
     const { rerender } = renderHook(
       ({ config }) =>
         useCoAgent({
@@ -155,137 +198,154 @@ describe("useCoAgent config synchronization", () => {
           config,
         }),
       {
-        initialProps: { config: { configurable: { model: "gpt-4" } } },
+        initialProps: {
+          config: {
+            configurable: {
+              model: "gpt-4",
+              settings: { temperature: 0.5 },
+            },
+          },
+        },
       },
     );
 
-    // Clear the initial calls and reset captured updater
-    mockSetCoagentStatesWithRef.mockClear();
-    capturedUpdater = null;
-
-    // Change config
-    rerender({ config: { configurable: { model: "gpt-4o" } } });
-
-    // Should have called setCoagentStatesWithRef
-    expect(mockSetCoagentStatesWithRef).toHaveBeenCalledWith(expect.any(Function));
-    expect(capturedUpdater).toBeTruthy();
+    // Clear the initial calls
+    mockSetProperties.mockClear();
 
-    // Test the updater function behavior
-    const prevState = {
-      "test-agent": {
-        name: "test-agent",
-        state: { count: 5 },
-        config: { configurable: { model: "gpt-4" } },
-        running: true,
-        active: true,
-        threadId: "thread-123",
-        nodeName: "test-node",
-        runId: "run-456",
+    // Change nested config
+    rerender({
+      config: {
+        configurable: {
+          model: "gpt-4",
+          settings: { temperature: 0.7 }, // Only nested property changed
+        },
       },
-    };
-
-    const newState = capturedUpdater(prevState);
+    });
 
-    // Verify the state is updated correctly
-    expect(newState).toEqual({
-      "test-agent": {
-        name: "test-agent",
-        state: { count: 5 }, // State preserved
-        config: { configurable: { model: "gpt-4o" } }, // Config updated
-        running: true, // Other properties preserved
-        active: true,
-        threadId: "thread-123",
-        nodeName: "test-node",
-        runId: "run-456",
-      },
+    // Wait for effect to complete and verify setProperties was called
+    await waitFor(() => {
+      expect(mockSetProperties).toHaveBeenCalledWith({
+        configurable: {
+          model: "gpt-4",
+          settings: { temperature: 0.7 },
+        },
+      });
     });
   });
 
-  it("should create new agent state when agent doesn't exist", () => {
-    let capturedUpdater: any = null;
+  describe("State Management", () => {
+    // Helper to create mock subscriber params
+    const createMockParams = (stateOverride: any = {}) => ({
+      messages: [],
+      state: stateOverride,
+      agent: mockAgent as any,
+      input: {} as any,
+    });
+
+    it("should initialize agent state via onRunInitialized event", () => {
+      // Set agent to have NO state (empty object)
+      mockAgent.state = {} as any;
+
+      renderHook(() =>
+        useCoAgent({
+          name: "test-agent",
+          initialState: { count: 42 },
+        }),
+      );
+
+      // Verify subscription was created
+      expect(mockSubscribe).toHaveBeenCalled();
+      expect(lastSubscriber).toBeTruthy();
+
+      // Clear any initial state calls
+      mockSetState.mockClear();
+
+      // Trigger onRunInitialized with no run state
+      lastSubscriber?.onRunInitialized?.(createMockParams({}));
+
+      // Should set state to initialState
+      expect(mockSetState).toHaveBeenCalledWith({ count: 42 });
 
-    mockSetCoagentStatesWithRef.mockImplementation((updater) => {
-      capturedUpdater = updater;
-      return updater;
+      // Reset agent state
+      mockAgent.state = { count: 0 };
     });
 
-    const { rerender } = renderHook(
-      ({ config }) =>
+    it("should preserve existing agent state on onRunInitialized", () => {
+      // Set agent to have existing state
+      mockAgent.state = { count: 100 };
+
+      renderHook(() =>
         useCoAgent({
           name: "test-agent",
-          initialState: { count: 0 },
-          config,
+          initialState: { count: 42 },
         }),
-      {
-        initialProps: { config: { configurable: { model: "gpt-4" } } },
-      },
-    );
+      );
 
-    // Clear the initial calls and reset captured updater
-    mockSetCoagentStatesWithRef.mockClear();
-    capturedUpdater = null;
+      // Clear any initial state calls
+      mockSetState.mockClear();
 
-    // Change config
-    rerender({ config: { configurable: { model: "gpt-4o" } } });
+      // Trigger onRunInitialized with no new state
+      lastSubscriber?.onRunInitialized?.(createMockParams({}));
+
+      // Should NOT override existing state
+      expect(mockSetState).not.toHaveBeenCalled();
+
+      // Reset agent state
+      mockAgent.state = { count: 0 };
+    });
 
-    // Should have called setCoagentStatesWithRef
-    expect(mockSetCoagentStatesWithRef).toHaveBeenCalledWith(expect.any(Function));
-    expect(capturedUpdater).toBeTruthy();
+    it("should prioritize run state over initialState", () => {
+      renderHook(() =>
+        useCoAgent({
+          name: "test-agent",
+          initialState: { count: 42 },
+        }),
+      );
 
-    // Test the updater function behavior with empty previous state
-    const prevState = {}; // No existing agent state
+      // Clear any initial state calls
+      mockSetState.mockClear();
 
-    const newState = capturedUpdater(prevState);
+      // Trigger onRunInitialized with state from the run
+      lastSubscriber?.onRunInitialized?.(createMockParams({ count: 999 }));
 
-    // Verify the state creates a new agent state with default values
-    expect(newState).toEqual({
-      "test-agent": {
-        name: "test-agent",
-        state: { count: 0 }, // Uses initialState
-        config: { configurable: { model: "gpt-4o" } }, // New config
-        running: false, // Default values
-        active: false,
-        threadId: undefined,
-        nodeName: undefined,
-        runId: undefined,
-      },
+      // Should use run state, not initialState
+      expect(mockSetState).toHaveBeenCalledWith({ count: 999 });
     });
-  });
 
-  it("should handle deeply nested config changes", () => {
-    const { rerender } = renderHook(
-      ({ config }) =>
+    it("should handle setState with object updates", () => {
+      const { result } = renderHook(() =>
         useCoAgent({
           name: "test-agent",
           initialState: { count: 0 },
-          config,
         }),
-      {
-        initialProps: {
-          config: {
-            configurable: {
-              model: "gpt-4",
-              settings: { temperature: 0.5 },
-            },
-          },
-        },
-      },
-    );
+      );
 
-    // Clear the initial calls
-    mockSetCoagentStatesWithRef.mockClear();
+      // Update state with object
+      result.current.setState({ count: 5 });
 
-    // Change nested config
-    rerender({
-      config: {
-        configurable: {
-          model: "gpt-4",
-          settings: { temperature: 0.7 }, // Only nested property changed
-        },
-      },
+      // Should merge with existing state
+      expect(mockSetState).toHaveBeenCalledWith({ count: 5 });
     });
 
-    // Should detect the nested change
-    expect(mockSetCoagentStatesWithRef).toHaveBeenCalledWith(expect.any(Function));
+    it("should handle setState with function updaters", () => {
+      // Set current agent state
+      mockAgent.state = { count: 10 };
+
+      const { result } = renderHook(() =>
+        useCoAgent({
+          name: "test-agent",
+          initialState: { count: 0 },
+        }),
+      );
+
+      // Update state with function
+      result.current.setState((prev) => ({ count: (prev?.count || 0) + 5 }));
+
+      // Should call function with current state and set result
+      expect(mockSetState).toHaveBeenCalledWith({ count: 15 });
+
+      // Reset agent state
+      mockAgent.state = { count: 0 };
+    });
   });
 });

package/src/hooks/index.ts

@@ -24,6 +24,6 @@ export { useRenderToolCall } from "./use-render-tool-call";
 export { useDefaultTool } from "./use-default-tool";
 export { useLazyToolRenderer } from "./use-lazy-tool-renderer";
 export {
-  useConfigureChatSuggestions,
+  useCopilotChatSuggestions,
   type UseCopilotChatSuggestionsConfiguration,
-} from "./use-configure-chat-suggestions";
+} from "./use-copilot-chat-suggestions";

package/src/hooks/use-agent-nodename.ts

@@ -0,0 +1,30 @@
+import { useEffect, useRef } from "react";
+import type { AgentSubscriber } from "@ag-ui/client";
+import { useAgent } from "@copilotkitnext/react";
+
+export function useAgentNodeName(agentName?: string) {
+  const { agent } = useAgent({ agentId: agentName });
+  const nodeNameRef = useRef<string>("start");
+
+  useEffect(() => {
+    if (!agent) return;
+    const subscriber: AgentSubscriber = {
+      onStepStartedEvent: ({ event }) => {
+        nodeNameRef.current = event.stepName;
+      },
+      onRunStartedEvent: () => {
+        nodeNameRef.current = "start";
+      },
+      onRunFinishedEvent: () => {
+        nodeNameRef.current = "end";
+      },
+    };
+
+    const subscription = agent.subscribe(subscriber);
+    return () => {
+      subscription.unsubscribe();
+    };
+  }, [agent]);
+
+  return nodeNameRef.current;
+}

package/src/hooks/use-coagent-state-render-bridge.tsx

@@ -5,6 +5,7 @@ import { useCoAgentStateRenders } from "../context";
 import { dataToUUID, parseJson } from "@copilotkit/shared";
 
 function getStateWithoutConstantKeys(state: any) {
+  if (!state) return {};
   const { messages, tools, copilotkit, ...stateWithoutConstantKeys } = state;
   return stateWithoutConstantKeys;
 }
@@ -148,23 +149,17 @@ export function useCoagentStateRenderBridge(agentId: string, props: CoAgentState
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [agentId, nodeName]);
 
-  if (messageIndexInRun !== 0) {
-    return null;
-  }
-
   const getStateRender = useCallback(
     (messageId: string) => {
-    return Object.entries(coAgentStateRenders).find(([stateRenderId, stateRender]) => {
-      if (claimsRef.current[messageId]) {
-        return stateRenderId === claimsRef.current[messageId].stateRenderId;
-      }
-      const matchingAgentName = stateRender.name === agentId;
-      const matchesNodeContext = stateRender.nodeName
-        ? stateRender.nodeName === nodeName
-        : true;
-      return matchingAgentName && matchesNodeContext;
-    });
-  },
+      return Object.entries(coAgentStateRenders).find(([stateRenderId, stateRender]) => {
+        if (claimsRef.current[messageId]) {
+          return stateRenderId === claimsRef.current[messageId].stateRenderId;
+        }
+        const matchingAgentName = stateRender.name === agentId;
+        const matchesNodeContext = stateRender.nodeName ? stateRender.nodeName === nodeName : true;
+        return matchingAgentName && matchesNodeContext;
+      });
+    },
     [coAgentStateRenders, nodeName, agentId],
   );
 
@@ -200,8 +195,8 @@ export function useCoagentStateRenderBridge(agentId: string, props: CoAgentState
     const renderClaimedByOtherMessage = Object.values(claimsRef.current).find(
       (c) =>
         c.stateRenderId === stateRenderId &&
-        dataToUUID(JSON.stringify(getStateWithoutConstantKeys(c.stateSnapshot))) ===
-          dataToUUID(JSON.stringify(getStateWithoutConstantKeys(renderSnapshot))),
+        dataToUUID(getStateWithoutConstantKeys(c.stateSnapshot)) ===
+          dataToUUID(getStateWithoutConstantKeys(renderSnapshot)),
     );
     if (renderClaimedByOtherMessage) {
       // If:
@@ -228,6 +223,10 @@ export function useCoagentStateRenderBridge(agentId: string, props: CoAgentState
   };
 
   return useMemo(() => {
+    if (messageIndexInRun !== 0) {
+      return null;
+    }
+
     const [stateRenderId, stateRender] = getStateRender(message.id) ?? [];
 
     if (!stateRender || !stateRenderId) {
@@ -270,11 +269,11 @@ export function useCoagentStateRenderBridge(agentId: string, props: CoAgentState
 
       if (typeof stateRender.render === "string") return stateRender.render;
 
-    return stateRender.render({
-      status,
-      // Always use state from claim, to make sure the state does not seem "wiped" for a fraction of a second
-      state: claimsRef.current[message.id].stateSnapshot ?? {},
-      nodeName: nodeName ?? "",
+      return stateRender.render({
+        status,
+        // Always use state from claim, to make sure the state does not seem "wiped" for a fraction of a second
+        state: claimsRef.current[message.id].stateSnapshot ?? {},
+        nodeName: nodeName ?? "",
       });
     }
   }, [
@@ -285,6 +284,7 @@ export function useCoagentStateRenderBridge(agentId: string, props: CoAgentState
     nodeName,
     effectiveRunId,
     message.id,
+    messageIndexInRun,
   ]);
 }
 

package/src/hooks/use-coagent.ts

@@ -90,8 +90,9 @@
 
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { Message } from "@copilotkit/shared";
-import { useAgent } from "@copilotkitnext/react";
+import { useAgent, useCopilotKit } from "@copilotkitnext/react";
 import { type AgentSubscriber } from "@ag-ui/client";
+import { useAgentNodeName } from "./use-agent-nodename";
 
 interface UseCoagentOptionsBase {
   /**
@@ -203,7 +204,8 @@ export type HintFunction = (params: HintFunctionParams) => Message | undefined;
  */
 export function useCoAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentReturnType<T> {
   const { agent } = useAgent({ agentId: options.name });
-  const nodeNameRef = useRef<string>("start");
+  const { copilotkit } = useCopilotKit();
+  const nodeName = useAgentNodeName(options.name);
 
   const handleStateUpdate = useCallback(
     (newState: T | ((prevState: T | undefined) => T)) => {
@@ -219,6 +221,22 @@ export function useCoAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentRe
     [agent?.state, agent?.setState],
   );
 
+  useEffect(() => {
+    if (!options.config && !options.configurable) return;
+
+    let config = options.config ?? {};
+    if (options.configurable) {
+      config = {
+        ...config,
+        configurable: {
+          ...options.configurable,
+          ...config.configurable,
+        },
+      };
+    }
+    copilotkit.setProperties(config);
+  }, [options.config, options.configurable]);
+
   const externalStateStr = useMemo(
     () => (isExternalStateManagement(options) ? JSON.stringify(options.state) : undefined),
     [isExternalStateManagement(options) ? JSON.stringify(options.state) : undefined],
@@ -284,15 +302,6 @@ export function useCoAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentRe
           handleStateUpdate(initialStateRef.current);
         }
       },
-      onStepStartedEvent: ({ event }) => {
-        nodeNameRef.current = event.stepName;
-      },
-      onRunStartedEvent: () => {
-        nodeNameRef.current = "start";
-      },
-      onRunFinishedEvent: () => {
-        nodeNameRef.current = "end";
-      },
     };
 
     const subscription = agent.subscribe(subscriber);
@@ -314,7 +323,7 @@ export function useCoAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentRe
         ({} as T);
       return {
         name: options.name,
-        nodeName: nodeNameRef.current,
+        nodeName,
         threadId: undefined,
         running: false,
         state: initialState as T,
@@ -327,7 +336,7 @@ export function useCoAgent<T = any>(options: UseCoagentOptions<T>): UseCoagentRe
 
     return {
       name: agent?.agentId ?? options.name,
-      nodeName: nodeNameRef.current,
+      nodeName,
       threadId: agent.threadId,
       running: agent.isRunning,
       state: agent.state,