@copilotkit/react-core 1.51.3-next.6 → 1.51.3-next.7

package/src/hooks/__tests__/use-coagent-state-render.test.tsx

@@ -0,0 +1,334 @@
+import React, { type ReactNode } from "react";
+import { render, renderHook, waitFor } from "@testing-library/react";
+import { useCoAgentStateRender } from "../use-coagent-state-render";
+import type { CoAgentStateRender } from "../../types/coagent-action";
+import {
+  CoAgentStateRendersProvider,
+  CopilotContext,
+  useCoAgentStateRenders,
+} from "../../context";
+import { CopilotKitAgentDiscoveryError, randomId } from "@copilotkit/shared";
+import { createTestCopilotContext } from "../../test-helpers/copilot-context";
+
+const addToast = jest.fn();
+const setBannerError = jest.fn();
+
+jest.mock("../../components/toast/toast-provider", () => ({
+  useToast: () => ({
+    addToast,
+    setBannerError,
+  }),
+}));
+
+function createWrapper(copilotContextValue: ReturnType<typeof createTestCopilotContext>) {
+  return function Wrapper({ children }: { children: ReactNode }) {
+    return (
+      <CopilotContext.Provider value={copilotContextValue}>
+        <CoAgentStateRendersProvider>{children}</CoAgentStateRendersProvider>
+      </CopilotContext.Provider>
+    );
+  };
+}
+
+function useHarness<T>(action: Parameters<typeof useCoAgentStateRender<T>>[0], deps?: unknown[]) {
+  useCoAgentStateRender(action, deps);
+  return useCoAgentStateRenders();
+}
+
+function HookUser<T>({
+  action,
+  deps,
+}: {
+  action: CoAgentStateRender<T>;
+  deps?: unknown[];
+}) {
+  useCoAgentStateRender(action, deps);
+  return null;
+}
+
+function getSingleEntry<T>(renders: Record<string, T>) {
+  const entries = Object.entries(renders);
+  expect(entries).toHaveLength(1);
+  return entries[0];
+}
+
+describe("useCoAgentStateRender (hook behaviors)", () => {
+  let idCounter = 0;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    idCounter = 0;
+    (randomId as jest.Mock).mockImplementation(() => `test-random-id-${++idCounter}`);
+  });
+
+  it("registers state render and writes to the render cache", async () => {
+    const chatComponentsCache = { current: { actions: {}, coAgentStateRenders: {} } };
+    const wrapper = createWrapper(
+      createTestCopilotContext({
+        chatComponentsCache,
+      }),
+    );
+
+    const renderFn = jest.fn(() => null);
+
+    const { result } = renderHook(
+      () =>
+        useHarness({
+          name: "agent-a",
+          nodeName: "node-1",
+          render: renderFn,
+        }),
+      { wrapper },
+    );
+
+    await waitFor(() => {
+      expect(Object.keys(result.current.coAgentStateRenders)).toHaveLength(1);
+    });
+
+    expect(chatComponentsCache.current.coAgentStateRenders["agent-a-node-1"]).toBe(renderFn);
+  });
+
+  it("mutates handler + cache in place when dependencies are omitted", async () => {
+    const chatComponentsCache = { current: { actions: {}, coAgentStateRenders: {} } };
+    const wrapper = createWrapper(
+      createTestCopilotContext({
+        chatComponentsCache,
+      }),
+    );
+
+    const handlerOne = jest.fn();
+    const handlerTwo = jest.fn();
+    const renderOne = jest.fn(() => null);
+    const renderTwo = jest.fn(() => null);
+
+    const { result, rerender } = renderHook(
+      ({ handler, renderFn }) =>
+        useHarness({
+          name: "agent-b",
+          handler,
+          render: renderFn,
+        }),
+      {
+        wrapper,
+        initialProps: { handler: handlerOne, renderFn: renderOne },
+      },
+    );
+
+    await waitFor(() => {
+      expect(Object.keys(result.current.coAgentStateRenders)).toHaveLength(1);
+    });
+
+    const initialRenders = result.current.coAgentStateRenders;
+    const [id, initialRender] = getSingleEntry(initialRenders);
+
+    expect(initialRender.handler).toBe(handlerOne);
+    expect(chatComponentsCache.current.coAgentStateRenders["agent-b-global"]).toBe(renderOne);
+
+    rerender({ handler: handlerTwo, renderFn: renderTwo });
+
+    expect(result.current.coAgentStateRenders).toBe(initialRenders);
+    expect(result.current.coAgentStateRenders[id].handler).toBe(handlerTwo);
+    expect(chatComponentsCache.current.coAgentStateRenders["agent-b-global"]).toBe(renderTwo);
+  });
+
+  it("re-registers when dependencies change", async () => {
+    const wrapper = createWrapper(createTestCopilotContext());
+
+    const handlerOne = jest.fn();
+    const handlerTwo = jest.fn();
+
+    const { result, rerender } = renderHook(
+      ({ deps, handler }) =>
+        useHarness(
+          {
+            name: "agent-c",
+            handler,
+          },
+          deps,
+        ),
+      {
+        wrapper,
+        initialProps: { deps: [0], handler: handlerOne },
+      },
+    );
+
+    await waitFor(() => {
+      expect(Object.keys(result.current.coAgentStateRenders)).toHaveLength(1);
+    });
+
+    const initialRenders = result.current.coAgentStateRenders;
+    const [id] = Object.keys(initialRenders);
+
+    rerender({ deps: [1], handler: handlerTwo });
+
+    await waitFor(() => {
+      expect(result.current.coAgentStateRenders).not.toBe(initialRenders);
+    });
+
+    expect(result.current.coAgentStateRenders[id].handler).toBe(handlerTwo);
+  });
+
+  it("re-registers when string render changes", async () => {
+    const chatComponentsCache = { current: { actions: {}, coAgentStateRenders: {} } };
+    const wrapper = createWrapper(
+      createTestCopilotContext({
+        chatComponentsCache,
+      }),
+    );
+
+    const { result, rerender } = renderHook(
+      ({ renderValue }) =>
+        useHarness({
+          name: "agent-d",
+          render: renderValue,
+        }),
+      {
+        wrapper,
+        initialProps: { renderValue: "Step 1" },
+      },
+    );
+
+    await waitFor(() => {
+      expect(Object.keys(result.current.coAgentStateRenders)).toHaveLength(1);
+    });
+
+    const initialRenders = result.current.coAgentStateRenders;
+    rerender({ renderValue: "Step 2" });
+
+    await waitFor(() => {
+      expect(result.current.coAgentStateRenders).not.toBe(initialRenders);
+    });
+
+    expect(chatComponentsCache.current.coAgentStateRenders["agent-d-global"]).toBe("Step 2");
+  });
+
+  it("warns when duplicate registrations target the same agent + node", async () => {
+    const copilotContextValue = createTestCopilotContext();
+
+    function DuplicateHarness() {
+      return (
+        <>
+          <HookUser
+            action={{
+              name: "agent-dup",
+              nodeName: "node-x",
+              handler: jest.fn(),
+            }}
+          />
+          <HookUser
+            action={{
+              name: "agent-dup",
+              nodeName: "node-x",
+              handler: jest.fn(),
+            }}
+          />
+        </>
+      );
+    }
+
+    render(
+      <CopilotContext.Provider value={copilotContextValue}>
+        <CoAgentStateRendersProvider>
+          <DuplicateHarness />
+        </CoAgentStateRendersProvider>
+      </CopilotContext.Provider>,
+    );
+
+    await waitFor(() => {
+      expect(addToast).toHaveBeenCalled();
+    });
+
+    expect(addToast).toHaveBeenCalledWith(
+      expect.objectContaining({
+        type: "warning",
+        message:
+          "Found multiple state renders for agent agent-dup and node node-x. State renders might get overridden",
+      }),
+    );
+  });
+
+  it("does not warn when duplicate agents target different nodes", async () => {
+    const copilotContextValue = createTestCopilotContext();
+
+    function NonDuplicateHarness() {
+      return (
+        <>
+          <HookUser
+            action={{
+              name: "agent-ok",
+              nodeName: "node-a",
+              handler: jest.fn(),
+            }}
+          />
+          <HookUser
+            action={{
+              name: "agent-ok",
+              nodeName: "node-b",
+              handler: jest.fn(),
+            }}
+          />
+        </>
+      );
+    }
+
+    render(
+      <CopilotContext.Provider value={copilotContextValue}>
+        <CoAgentStateRendersProvider>
+          <NonDuplicateHarness />
+        </CoAgentStateRendersProvider>
+      </CopilotContext.Provider>,
+    );
+
+    await waitFor(() => {
+      expect(addToast).not.toHaveBeenCalled();
+    });
+  });
+
+  it("surfaces missing agents in the banner error state", async () => {
+    const availableAgents = [{ name: "known-agent", id: "agent-1" }];
+    const wrapper = createWrapper(
+      createTestCopilotContext({
+        availableAgents,
+      }),
+    );
+
+    renderHook(
+      () =>
+        useHarness({
+          name: "missing-agent",
+          handler: jest.fn(),
+        }),
+      { wrapper },
+    );
+
+    await waitFor(() => {
+      expect(CopilotKitAgentDiscoveryError).toHaveBeenCalledWith({
+        agentName: "missing-agent",
+        availableAgents: [{ name: "known-agent", id: "agent-1" }],
+      });
+      expect(setBannerError).toHaveBeenCalled();
+    });
+  });
+
+  it("does not surface banner errors when agent is available", async () => {
+    const availableAgents = [{ name: "agent-present", id: "agent-2" }];
+    const wrapper = createWrapper(
+      createTestCopilotContext({
+        availableAgents,
+      }),
+    );
+
+    renderHook(
+      () =>
+        useHarness({
+          name: "agent-present",
+          handler: jest.fn(),
+        }),
+      { wrapper },
+    );
+
+    await waitFor(() => {
+      expect(setBannerError).not.toHaveBeenCalled();
+    });
+  });
+});

package/src/hooks/use-coagent-state-render-bridge.helpers.ts

@@ -0,0 +1,311 @@
+import { dataToUUID, parseJson } from "@copilotkit/shared";
+
+export enum RenderStatus {
+  InProgress = "inProgress",
+  Complete = "complete",
+}
+
+export enum ClaimAction {
+  Create = "create",
+  Override = "override",
+  Existing = "existing",
+  Block = "block",
+}
+
+export interface StateRenderContext {
+  agentId: string;
+  stateRenderId: string;
+  messageId: string;
+  runId: string;
+  messageIndex?: number;
+}
+
+export interface Claim {
+  stateRenderId: string;
+  runId?: string;
+  stateSnapshot?: any;
+  locked?: boolean;
+  messageIndex?: number;
+}
+
+export type ClaimsByMessageId = Record<string, Claim>;
+
+export interface ClaimResolution {
+  canRender: boolean;
+  action: ClaimAction;
+  nextClaim?: Claim;
+  lockOthers?: boolean;
+  updateRunId?: string;
+}
+
+export interface SnapshotCaches {
+  byStateRenderAndRun: Record<string, any>;
+  byMessageId: Record<string, any>;
+}
+
+export interface SnapshotSelectionInput {
+  messageId: string;
+  messageName?: string;
+  allowLiveState?: boolean;
+  skipLatestCache?: boolean;
+  stateRenderId?: string;
+  effectiveRunId: string;
+  stateSnapshotProp?: any;
+  agentState?: any;
+  agentMessages?: Array<{ id: string; role?: string }>;
+  existingClaim?: Claim;
+  caches: SnapshotCaches;
+}
+
+export interface SnapshotSelectionResult {
+  snapshot?: any;
+  hasSnapshotKeys: boolean;
+  cachedSnapshot?: any;
+  allowEmptySnapshot?: boolean;
+  snapshotForClaim?: any;
+}
+
+function getStateWithoutConstantKeys(state: any) {
+  if (!state) return {};
+  const { messages, tools, copilotkit, ...stateWithoutConstantKeys } = state;
+  return stateWithoutConstantKeys;
+}
+
+// Function that compares states, without the constant keys
+export function areStatesEquals(a: any, b: any) {
+  if ((a && !b) || (!a && b)) return false;
+  const { messages, tools, copilotkit, ...aWithoutConstantKeys } = a;
+  const {
+    messages: bMessages,
+    tools: bTools,
+    copilotkit: bCopilotkit,
+    ...bWithoutConstantKeys
+  } = b;
+
+  return JSON.stringify(aWithoutConstantKeys) === JSON.stringify(bWithoutConstantKeys);
+}
+
+export function isPlaceholderMessageId(messageId: string | undefined) {
+  return !!messageId && messageId.startsWith("coagent-state-render-");
+}
+
+export function isPlaceholderMessageName(messageName: string | undefined) {
+  return messageName === "coagent-state-render";
+}
+
+export function readCachedMessageEntry(entry: any): { snapshot?: any; runId?: string } {
+  if (!entry || typeof entry !== "object") {
+    return { snapshot: entry, runId: undefined };
+  }
+  const snapshot = "snapshot" in entry ? entry.snapshot : entry;
+  const runId = "runId" in entry ? entry.runId : undefined;
+  return { snapshot, runId };
+}
+
+export function getEffectiveRunId({
+  existingClaimRunId,
+  cachedMessageRunId,
+  runId,
+}: {
+  existingClaimRunId?: string;
+  cachedMessageRunId?: string;
+  runId?: string;
+}) {
+  return existingClaimRunId || cachedMessageRunId || runId || "pending";
+}
+
+/**
+ * Resolve whether a message can claim a render slot.
+ * This is a pure decision function; the caller applies claim mutations.
+ */
+export function resolveClaim({
+  claims,
+  context,
+  stateSnapshot,
+}: {
+  claims: ClaimsByMessageId;
+  context: StateRenderContext;
+  stateSnapshot?: any;
+}): ClaimResolution {
+  const { messageId, stateRenderId, runId, messageIndex } = context;
+  const existing = claims[messageId];
+
+  if (existing) {
+    const canRender = existing.stateRenderId === stateRenderId;
+    const shouldUpdateRunId =
+      canRender && runId && (!existing.runId || existing.runId === "pending");
+    return {
+      canRender,
+      action: canRender ? ClaimAction.Existing : ClaimAction.Block,
+      updateRunId: shouldUpdateRunId ? runId : undefined,
+    };
+  }
+
+  const normalizedRunId = runId ?? "pending";
+  const renderClaimedByOtherMessageEntry = Object.entries(claims).find(
+    ([, claim]) =>
+      claim.stateRenderId === stateRenderId &&
+      (claim.runId ?? "pending") === normalizedRunId &&
+      dataToUUID(getStateWithoutConstantKeys(claim.stateSnapshot)) ===
+        dataToUUID(getStateWithoutConstantKeys(stateSnapshot)),
+  );
+
+  const renderClaimedByOtherMessage = renderClaimedByOtherMessageEntry?.[1];
+  const claimedMessageId = renderClaimedByOtherMessageEntry?.[0];
+
+  if (renderClaimedByOtherMessage) {
+    if (
+      messageIndex !== undefined &&
+      renderClaimedByOtherMessage.messageIndex !== undefined &&
+      messageIndex > renderClaimedByOtherMessage.messageIndex
+    ) {
+      return {
+        canRender: true,
+        action: ClaimAction.Override,
+        nextClaim: { stateRenderId, runId, messageIndex },
+        lockOthers:
+          runId === renderClaimedByOtherMessage.runId || isPlaceholderMessageId(claimedMessageId),
+      };
+    }
+
+    if (runId && renderClaimedByOtherMessage.runId && runId !== renderClaimedByOtherMessage.runId) {
+      return {
+        canRender: true,
+        action: ClaimAction.Override,
+        nextClaim: { stateRenderId, runId, messageIndex },
+        lockOthers: isPlaceholderMessageId(claimedMessageId),
+      };
+    }
+
+    if (isPlaceholderMessageId(claimedMessageId)) {
+      return {
+        canRender: true,
+        action: ClaimAction.Override,
+        nextClaim: { stateRenderId, runId, messageIndex },
+        lockOthers: true,
+      };
+    }
+
+    if (
+      stateSnapshot &&
+      renderClaimedByOtherMessage.stateSnapshot &&
+      !areStatesEquals(renderClaimedByOtherMessage.stateSnapshot, stateSnapshot)
+    ) {
+      return {
+        canRender: true,
+        action: ClaimAction.Override,
+        nextClaim: { stateRenderId, runId },
+      };
+    }
+
+    return { canRender: false, action: ClaimAction.Block };
+  }
+
+  if (!runId) {
+    return { canRender: false, action: ClaimAction.Block };
+  }
+
+  return {
+    canRender: true,
+    action: ClaimAction.Create,
+    nextClaim: { stateRenderId, runId, messageIndex },
+  };
+}
+
+/**
+ * Select the best snapshot to render for this message.
+ * Priority order is:
+ * 1) explicit message snapshot
+ * 2) live agent state (latest assistant only)
+ * 3) cached snapshot for message
+ * 4) cached snapshot for stateRenderId+runId
+ * 5) last cached snapshot for stateRenderId
+ */
+export function selectSnapshot({
+  messageId,
+  messageName,
+  allowLiveState,
+  skipLatestCache,
+  stateRenderId,
+  effectiveRunId,
+  stateSnapshotProp,
+  agentState,
+  agentMessages,
+  existingClaim,
+  caches,
+}: SnapshotSelectionInput): SnapshotSelectionResult {
+  const lastAssistantId = agentMessages
+    ? [...agentMessages].reverse().find((msg) => msg.role === "assistant")?.id
+    : undefined;
+  const latestSnapshot =
+    stateRenderId !== undefined ? caches.byStateRenderAndRun[`${stateRenderId}::latest`] : undefined;
+  const messageIndex = agentMessages
+    ? agentMessages.findIndex((msg) => msg.id === messageId)
+    : -1;
+  const messageRole =
+    messageIndex >= 0 && agentMessages ? agentMessages[messageIndex]?.role : undefined;
+  let previousUserMessageId: string | undefined;
+  if (messageIndex > 0 && agentMessages) {
+    for (let i = messageIndex - 1; i >= 0; i -= 1) {
+      if (agentMessages[i]?.role === "user") {
+        previousUserMessageId = agentMessages[i]?.id;
+        break;
+      }
+    }
+  }
+  const liveStateIsStale =
+    stateSnapshotProp === undefined &&
+    latestSnapshot !== undefined &&
+    agentState !== undefined &&
+    areStatesEquals(latestSnapshot, agentState);
+  const shouldUseLiveState =
+    (Boolean(allowLiveState) || !lastAssistantId || messageId === lastAssistantId) &&
+    !liveStateIsStale;
+  const snapshot = stateSnapshotProp
+    ? parseJson(stateSnapshotProp, stateSnapshotProp)
+    : shouldUseLiveState
+      ? agentState
+      : undefined;
+  const hasSnapshotKeys = !!(snapshot && Object.keys(snapshot).length > 0);
+  const allowEmptySnapshot =
+    snapshot !== undefined &&
+    !hasSnapshotKeys &&
+    (stateSnapshotProp !== undefined || shouldUseLiveState);
+
+  const messageCacheEntry = caches.byMessageId[messageId];
+  const cachedMessageSnapshot = readCachedMessageEntry(messageCacheEntry).snapshot;
+  const cacheKey =
+    stateRenderId !== undefined ? `${stateRenderId}::${effectiveRunId}` : undefined;
+  let cachedSnapshot = cachedMessageSnapshot ?? caches.byMessageId[messageId];
+  if (cachedSnapshot === undefined && cacheKey && caches.byStateRenderAndRun[cacheKey] !== undefined) {
+    cachedSnapshot = caches.byStateRenderAndRun[cacheKey];
+  }
+  if (
+    cachedSnapshot === undefined &&
+    stateRenderId &&
+    previousUserMessageId &&
+    caches.byStateRenderAndRun[`${stateRenderId}::pending:${previousUserMessageId}`] !==
+      undefined
+  ) {
+    cachedSnapshot =
+      caches.byStateRenderAndRun[`${stateRenderId}::pending:${previousUserMessageId}`];
+  }
+  if (
+    cachedSnapshot === undefined &&
+    !skipLatestCache &&
+    stateRenderId &&
+    messageRole !== "assistant" &&
+    (stateSnapshotProp !== undefined ||
+      (agentState && Object.keys(agentState).length > 0))
+  ) {
+    cachedSnapshot = caches.byStateRenderAndRun[`${stateRenderId}::latest`];
+  }
+
+  const snapshotForClaim = existingClaim?.locked
+    ? existingClaim.stateSnapshot ?? cachedSnapshot
+    : hasSnapshotKeys
+      ? snapshot
+      : existingClaim?.stateSnapshot ?? cachedSnapshot;
+
+  return { snapshot, hasSnapshotKeys, cachedSnapshot, allowEmptySnapshot, snapshotForClaim };
+}