@copilotkit/react-core 1.57.3 → 1.58.0-canary.thread-id-propagation

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

@@ -1,1209 +0,0 @@
-import { vi } from "vitest";
-import React from "react";
-import { act, render, screen, waitFor, within } from "@testing-library/react";
-import { useCoAgentStateRender } from "../use-coagent-state-render";
-import { CoAgentStateRenderBridge } from "../use-coagent-state-render-bridge";
-import { useCopilotChatInternal } from "../use-copilot-chat_internal";
-import {
-  CoAgentStateRendersProvider,
-  CopilotContext,
-  useCoAgentStateRenders,
-} from "../../context";
-import type { Claim } from "../use-coagent-state-render-bridge.helpers";
-import { createTestCopilotContext } from "../../test-helpers/copilot-context";
-import { useRenderCustomMessages } from "../../v2";
-
-type TestMessage = {
-  id: string;
-  role: "user" | "assistant" | "system" | string;
-  content?: string;
-};
-
-type TestAgentSubscriber = {
-  onStateChanged?: (args?: { state?: Record<string, unknown> }) => void;
-  onStepStartedEvent?: (args: { event: { stepName: string } }) => void;
-  onStepFinishedEvent?: (args: { event: { stepName: string } }) => void;
-};
-
-const mockAgent = {
-  messages: [] as TestMessage[],
-  state: {},
-  isRunning: true,
-  subscribe: vi.fn(),
-  setMessages: vi.fn(),
-  setState: vi.fn(),
-  addMessage: vi.fn(),
-  abortRun: vi.fn(),
-  runAgent: vi.fn(),
-  detachActiveRun: vi.fn().mockResolvedValue(undefined),
-};
-
-let lastSubscriber: TestAgentSubscriber | null = null;
-
-vi.mock("../../v2", () => ({
-  useAgent: vi.fn(() => ({ agent: mockAgent })),
-  useCopilotKit: vi.fn(() => ({
-    copilotkit: {
-      connectAgent: vi.fn(),
-      getRunIdForMessage: vi.fn(),
-      runAgent: vi.fn(),
-      clearSuggestions: vi.fn(),
-      addSuggestionsConfig: vi.fn(),
-      reloadSuggestions: vi.fn(),
-      interruptElement: null,
-      subscribe: vi.fn(() => ({ unsubscribe: vi.fn() })),
-    },
-  })),
-  useCopilotChatConfiguration: vi.fn(() => ({ agentId: "test-agent" })),
-  useRenderCustomMessages: vi.fn(() => undefined),
-  useSuggestions: vi.fn(() => ({ suggestions: [], isLoading: false })),
-}));
-
-vi.mock("../../components/toast/toast-provider", () => ({
-  useToast: () => ({
-    setBannerError: vi.fn(),
-    addToast: vi.fn(),
-  }),
-}));
-
-vi.mock("../../components/error-boundary/error-utils", () => ({
-  useAsyncCallback: <T extends (...args: unknown[]) => unknown>(fn: T) => fn,
-}));
-
-vi.mock("../use-lazy-tool-renderer", () => ({
-  useLazyToolRenderer: vi.fn(() => () => null),
-}));
-
-function TestHarness({ snapshot }: { snapshot: string }) {
-  useCoAgentStateRender<{ current_step?: string }>({
-    name: "test-agent",
-    render: ({ state }) => (
-      <div data-testid="state">{state.current_step ?? "none"}</div>
-    ),
-  });
-
-  return (
-    <CoAgentStateRenderBridge
-      agentId="test-agent"
-      message={{ id: "msg-1", role: "assistant" }}
-      position="after"
-      runId="run-1"
-      messageIndex={0}
-      messageIndexInRun={0}
-      numberOfMessagesInRun={1}
-      stateSnapshot={snapshot}
-    />
-  );
-}
-
-function SnapshotHarness({
-  snapshot,
-  message,
-}: {
-  snapshot: string;
-  message: TestMessage;
-}) {
-  useCoAgentStateRender<{ current_step?: string }>({
-    name: "test-agent",
-    render: ({ state }) => (
-      <div data-testid="state">{state.current_step ?? "none"}</div>
-    ),
-  });
-
-  return (
-    <CoAgentStateRenderBridge
-      agentId="test-agent"
-      message={message}
-      position="after"
-      runId="run-early"
-      messageIndex={0}
-      messageIndexInRun={0}
-      numberOfMessagesInRun={1}
-      stateSnapshot={snapshot}
-    />
-  );
-}
-
-function LiveStateHarness({
-  message,
-}: {
-  message: Pick<TestMessage, "id" | "role">;
-}) {
-  useCoAgentStateRender<{ current_step?: string }>({
-    name: "test-agent",
-    render: ({ state }) => (
-      <div data-testid="state">{state.current_step ?? "none"}</div>
-    ),
-  });
-
-  return (
-    <CoAgentStateRenderBridge
-      agentId="test-agent"
-      message={message}
-      position="after"
-      runId="run-live"
-      messageIndex={0}
-      messageIndexInRun={0}
-      numberOfMessagesInRun={1}
-      stateSnapshot={undefined}
-    />
-  );
-}
-
-function NonFirstMessageHarness({ snapshot }: { snapshot: string }) {
-  useCoAgentStateRender<{ current_step?: string }>({
-    name: "test-agent",
-    render: ({ state }) => (
-      <div data-testid="state">{state.current_step ?? "none"}</div>
-    ),
-  });
-
-  return (
-    <CoAgentStateRenderBridge
-      agentId="test-agent"
-      message={{ id: "msg-second", role: "assistant", content: "" }}
-      position="after"
-      runId="run-non-first"
-      messageIndex={1}
-      messageIndexInRun={1}
-      numberOfMessagesInRun={2}
-      stateSnapshot={snapshot}
-    />
-  );
-}
-
-function MultiRunHarness({
-  snapshot,
-  runId,
-  messageId,
-  messageIndex,
-}: {
-  snapshot: string;
-  runId: string;
-  messageId: string;
-  messageIndex: number;
-}) {
-  useCoAgentStateRender<{ current_step?: string }>({
-    name: "test-agent",
-    render: ({ state }) => (
-      <div data-testid="state">{state.current_step ?? "none"}</div>
-    ),
-  });
-
-  return (
-    <CoAgentStateRenderBridge
-      agentId="test-agent"
-      message={{ id: messageId, role: "assistant", content: "" }}
-      position="after"
-      runId={runId}
-      messageIndex={messageIndex}
-      messageIndexInRun={0}
-      numberOfMessagesInRun={1}
-      stateSnapshot={snapshot}
-    />
-  );
-}
-
-function ClaimsObserver({
-  onChange,
-}: {
-  onChange: (claims: Record<string, Claim>) => void;
-}) {
-  const { claimsRef } = useCoAgentStateRenders();
-  React.useEffect(() => {
-    onChange(claimsRef.current as Record<string, Claim>);
-  });
-  return null;
-}
-
-function ChatHarness({ tick }: { tick: number }) {
-  useCoAgentStateRender<{ current_step?: string }>({
-    name: "test-agent",
-    render: ({ state }) => (
-      <div data-testid="state">{state.current_step ?? "none"}</div>
-    ),
-  });
-
-  const { messages } = useCopilotChatInternal();
-
-  return (
-    <>
-      {messages.map((message) =>
-        message.generativeUI ? (
-          <div key={message.id} data-testid={`message-${message.id}`}>
-            {message.generativeUI()}
-          </div>
-        ) : null,
-      )}
-      <div data-testid="tick">{tick}</div>
-    </>
-  );
-}
-
-describe("useCoAgentStateRender", () => {
-  beforeEach(() => {
-    lastSubscriber = null;
-    mockAgent.state = {};
-    mockAgent.messages = [];
-    mockAgent.subscribe.mockImplementation(
-      (subscriber: TestAgentSubscriber) => {
-        lastSubscriber = subscriber;
-        return { unsubscribe: vi.fn() };
-      },
-    );
-  });
-
-  it("re-renders when state snapshots change", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <TestHarness
-            snapshot={JSON.stringify({ current_step: "Processing..." })}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Processing...");
-    });
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <TestHarness
-            snapshot={JSON.stringify({ current_step: "Thinking..." })}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Thinking...");
-    });
-  });
-
-  it("renders snapshots that arrive before assistant text", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    const baseMessage = { id: "msg-early", role: "assistant", content: "" };
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <SnapshotHarness
-            message={baseMessage}
-            snapshot={JSON.stringify({ current_step: "Processing..." })}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Processing...");
-    });
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <SnapshotHarness
-            message={baseMessage}
-            snapshot={JSON.stringify({ current_step: "Thinking..." })}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Thinking...");
-    });
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <SnapshotHarness
-            message={{ ...baseMessage, content: "Hello" }}
-            snapshot={JSON.stringify({ current_step: "Finalizing..." })}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Finalizing...");
-    });
-  });
-
-  it("re-renders from live agent state updates", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <LiveStateHarness message={{ id: "msg-live", role: "assistant" }} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    mockAgent.state = { current_step: "Processing..." };
-    act(() => {
-      lastSubscriber?.onStateChanged?.({ state: mockAgent.state });
-    });
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Processing...");
-    });
-
-    mockAgent.state = { current_step: "Thinking..." };
-    act(() => {
-      lastSubscriber?.onStateChanged?.({ state: mockAgent.state });
-    });
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Thinking...");
-    });
-  });
-
-  it("renders state snapshots before any assistant message exists", async () => {
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    mockAgent.messages = [{ id: "msg-user-1", role: "user", content: "Hi" }];
-    mockAgent.isRunning = false;
-    mockAgent.state = {};
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    const placeholderTestId =
-      "message-coagent-state-render-test-agent-pending:msg-user-1";
-    expect(screen.queryByTestId(placeholderTestId)).toBeNull();
-
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "Processing..." };
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={1} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Processing...");
-    });
-
-    expect(screen.getByTestId(placeholderTestId)).toBeTruthy();
-
-    const placeholderStep = screen.getByTestId("state").textContent;
-
-    mockAgent.isRunning = false;
-    mockAgent.state = {};
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-    ];
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={2} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.queryByTestId(placeholderTestId)).toBeNull();
-      expect(screen.getByTestId("message-msg-assistant-1")).toBeTruthy();
-      expect(screen.getByTestId("state").textContent).toBe("none");
-    });
-  });
-
-  it("does not render placeholder until the agent is running or has state", async () => {
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    mockAgent.messages = [{ id: "msg-user-1", role: "user", content: "Hi" }];
-    mockAgent.isRunning = false;
-    mockAgent.state = {};
-
-    render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(
-        screen.queryByTestId(
-          "message-coagent-state-render-test-agent-pending:msg-user-1",
-        ),
-      ).toBeNull();
-    });
-  });
-
-  it("renders for non-first messages in a run", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <NonFirstMessageHarness
-            snapshot={JSON.stringify({ current_step: "Processing..." })}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Processing...");
-    });
-  });
-
-  it("falls back to legacy renderer when renderCustomMessages throws", async () => {
-    vi.mocked(useRenderCustomMessages).mockImplementationOnce(() => () => {
-      throw new Error("boom");
-    });
-
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "Processing..." };
-
-    render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Processing...");
-    });
-  });
-
-  it("prefers legacy renderer over renderCustomMessages when both exist", async () => {
-    const renderCustomSpy = vi.fn(() => null);
-    vi.mocked(useRenderCustomMessages).mockImplementationOnce(
-      () => renderCustomSpy,
-    );
-
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "Processing..." };
-
-    render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Processing...");
-    });
-
-    expect(renderCustomSpy).not.toHaveBeenCalled();
-  });
-
-  it("renders empty state when agent state clears after completion", async () => {
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "Processing..." };
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Processing...");
-    });
-
-    mockAgent.isRunning = false;
-    mockAgent.state = {};
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={1} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("none");
-    });
-  });
-
-  it("allows independent renders across different runs with the same snapshot", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    const snapshot = JSON.stringify({ current_step: "Processing..." });
-
-    render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <MultiRunHarness
-            snapshot={snapshot}
-            runId="run-1"
-            messageId="msg-1"
-            messageIndex={0}
-          />
-          <MultiRunHarness
-            snapshot={snapshot}
-            runId="run-2"
-            messageId="msg-2"
-            messageIndex={1}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const nodes = screen.getAllByTestId("state");
-      expect(nodes).toHaveLength(2);
-      expect(nodes[0].textContent).toBe("Processing...");
-      expect(nodes[1].textContent).toBe("Processing...");
-    });
-  });
-
-  it("lets newer assistant messages claim even if snapshot matches", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    const snapshot = JSON.stringify({ current_step: "Processing..." });
-
-    render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <MultiRunHarness
-            snapshot={snapshot}
-            runId="run-1"
-            messageId="msg-1"
-            messageIndex={0}
-          />
-          <MultiRunHarness
-            snapshot={snapshot}
-            runId="run-1"
-            messageId="msg-2"
-            messageIndex={1}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const nodes = screen.getAllByTestId("state");
-      expect(nodes).toHaveLength(2);
-      expect(nodes[0].textContent).toBe("Processing...");
-      expect(nodes[1].textContent).toBe("Processing...");
-    });
-  });
-
-  it("locks older claims when a newer message claims", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    const snapshot = JSON.stringify({ current_step: "Processing..." });
-    let latestClaims: Record<string, Claim> = {};
-
-    render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ClaimsObserver onChange={(claims) => (latestClaims = claims)} />
-          <MultiRunHarness
-            snapshot={snapshot}
-            runId="run-1"
-            messageId="msg-1"
-            messageIndex={0}
-          />
-          <MultiRunHarness
-            snapshot={snapshot}
-            runId="run-1"
-            messageId="msg-2"
-            messageIndex={1}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(latestClaims["msg-1"]).toBeTruthy();
-      expect(latestClaims["msg-2"]).toBeTruthy();
-      expect(latestClaims["msg-1"].locked).toBe(true);
-    });
-  });
-
-  it("keeps older message snapshots stable when a new run starts", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    const snapshot = JSON.stringify({ current_step: "Processing..." });
-    let latestClaims: Record<string, Claim> = {};
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ClaimsObserver onChange={(claims) => (latestClaims = claims)} />
-          <MultiRunHarness
-            snapshot={snapshot}
-            runId="run-1"
-            messageId="msg-1"
-            messageIndex={0}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(latestClaims["msg-1"]?.stateSnapshot?.current_step).toBe(
-        "Processing...",
-      );
-    });
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ClaimsObserver onChange={(claims) => (latestClaims = claims)} />
-          <MultiRunHarness
-            snapshot={snapshot}
-            runId="run-1"
-            messageId="msg-1"
-            messageIndex={0}
-          />
-          <MultiRunHarness
-            snapshot={snapshot}
-            runId="run-2"
-            messageId="msg-2"
-            messageIndex={1}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(latestClaims["msg-1"]?.stateSnapshot?.current_step).toBe(
-        "Processing...",
-      );
-      expect(latestClaims["msg-2"]?.stateSnapshot?.current_step).toBe(
-        "Processing...",
-      );
-    });
-  });
-
-  it("does not overwrite a previous run snapshot when a new run has different state", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    const snapshotRun1 = JSON.stringify({ current_step: "Processing..." });
-    const snapshotRun2 = JSON.stringify({ current_step: "Finalizing..." });
-    let latestClaims: Record<string, Claim> = {};
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ClaimsObserver onChange={(claims) => (latestClaims = claims)} />
-          <MultiRunHarness
-            snapshot={snapshotRun1}
-            runId="run-1"
-            messageId="msg-1"
-            messageIndex={0}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(latestClaims["msg-1"]?.stateSnapshot?.current_step).toBe(
-        "Processing...",
-      );
-    });
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ClaimsObserver onChange={(claims) => (latestClaims = claims)} />
-          <MultiRunHarness
-            snapshot={snapshotRun1}
-            runId="run-1"
-            messageId="msg-1"
-            messageIndex={0}
-          />
-          <MultiRunHarness
-            snapshot={snapshotRun2}
-            runId="run-2"
-            messageId="msg-2"
-            messageIndex={1}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(latestClaims["msg-1"]?.stateSnapshot?.current_step).toBe(
-        "Processing...",
-      );
-      expect(latestClaims["msg-2"]?.stateSnapshot?.current_step).toBe(
-        "Finalizing...",
-      );
-    });
-  });
-
-  it("locks a claim when newer agent messages exist", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    const snapshot = JSON.stringify({ current_step: "Processing..." });
-    let latestClaims: Record<string, Claim> = {};
-
-    mockAgent.messages = [
-      { id: "msg-1", role: "assistant", content: "" },
-      { id: "msg-2", role: "assistant", content: "" },
-    ];
-
-    render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ClaimsObserver onChange={(claims) => (latestClaims = claims)} />
-          <MultiRunHarness
-            snapshot={snapshot}
-            runId="run-1"
-            messageId="msg-1"
-            messageIndex={0}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(latestClaims["msg-1"]).toBeTruthy();
-      expect(latestClaims["msg-1"].locked).toBe(true);
-    });
-  });
-
-  it("does not update a locked claim from new agent state", async () => {
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    let latestClaims: Record<string, Claim> = {};
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "First" };
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ClaimsObserver onChange={(claims) => (latestClaims = claims)} />
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(latestClaims["msg-assistant-1"]?.stateSnapshot?.current_step).toBe(
-        "First",
-      );
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-      { id: "msg-user-2", role: "user", content: "Next" },
-    ];
-    mockAgent.state = { current_step: "Second" };
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ClaimsObserver onChange={(claims) => (latestClaims = claims)} />
-          <ChatHarness tick={1} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(latestClaims["msg-assistant-1"]?.stateSnapshot?.current_step).toBe(
-        "First",
-      );
-    });
-  });
-
-  it("renders new run snapshots instead of reusing the previous run state", async () => {
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "First run" };
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("First run");
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-      { id: "msg-user-2", role: "user", content: "Next" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "Second run" };
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={1} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const stateNodes = screen.getAllByTestId("state");
-      const values = stateNodes.map((node) => node.textContent);
-      expect(values).toContain("Second run");
-    });
-  });
-
-  it("does not reuse latest cached snapshot for a new run placeholder before snapshots arrive", async () => {
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "First run" };
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("First run");
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-      { id: "msg-user-2", role: "user", content: "Next" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = {};
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={1} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const placeholder = screen.getByTestId(
-        "message-coagent-state-render-test-agent-pending:msg-user-2",
-      );
-      expect(placeholder.textContent).toContain("none");
-    });
-  });
-
-  it("does not show the previous snapshot for a new assistant message before its snapshot arrives", async () => {
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "First run" };
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const message = screen.getByTestId("message-msg-assistant-1");
-      expect(within(message).getByTestId("state").textContent).toBe(
-        "First run",
-      );
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-      { id: "msg-user-2", role: "user", content: "Next" },
-      { id: "msg-assistant-2", role: "assistant", content: "" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = {};
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={1} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const message = screen.getByTestId("message-msg-assistant-2");
-      expect(within(message).getByTestId("state").textContent).toBe("none");
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-      { id: "msg-user-2", role: "user", content: "Next" },
-      {
-        id: "msg-assistant-2",
-        role: "assistant",
-        content: "",
-        state: '{"current_step":"Second run"}',
-      },
-    ];
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={2} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const message = screen.getByTestId("message-msg-assistant-2");
-      expect(within(message).getByTestId("state").textContent).toBe(
-        "Second run",
-      );
-    });
-  });
-
-  it("does not show previous live state for a new run before the run updates state", async () => {
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "First run" };
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const message = screen.getByTestId("message-msg-assistant-1");
-      expect(within(message).getByTestId("state").textContent).toBe(
-        "First run",
-      );
-    });
-
-    mockAgent.messages = [
-      { id: "msg-user-1", role: "user", content: "Hi" },
-      { id: "msg-assistant-1", role: "assistant", content: "" },
-      { id: "msg-user-2", role: "user", content: "Next" },
-      { id: "msg-assistant-2", role: "assistant", content: "" },
-    ];
-    mockAgent.isRunning = true;
-    // live state still shows the previous run
-    mockAgent.state = { current_step: "First run" };
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={1} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const message = screen.getByTestId("message-msg-assistant-2");
-      expect(within(message).getByTestId("state").textContent).toBe("none");
-    });
-
-    mockAgent.state = { current_step: "Second run" };
-    act(() => {
-      lastSubscriber?.onStateChanged?.({ state: mockAgent.state });
-    });
-
-    await waitFor(() => {
-      const message = screen.getByTestId("message-msg-assistant-2");
-      expect(within(message).getByTestId("state").textContent).toBe(
-        "Second run",
-      );
-    });
-  });
-
-  it("renders an explicit empty state snapshot", async () => {
-    const copilotContextValue = createTestCopilotContext({
-      threadId: "thread-1",
-      agentSession: null,
-    });
-
-    mockAgent.messages = [
-      {
-        id: "msg-user-1",
-        role: "user",
-        content: "Hi",
-      },
-      {
-        id: "msg-assistant-1",
-        role: "assistant",
-        content: "",
-        state: '{"current_step":"Processing..."}',
-      },
-    ];
-    mockAgent.isRunning = true;
-    mockAgent.state = { current_step: "Processing..." };
-
-    const { rerender } = render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={0} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const message = screen.getByTestId("message-msg-assistant-1");
-      expect(within(message).getByTestId("state").textContent).toBe(
-        "Processing...",
-      );
-    });
-
-    mockAgent.messages = [
-      {
-        id: "msg-user-1",
-        role: "user",
-        content: "Hi",
-      },
-      {
-        id: "msg-assistant-1",
-        role: "assistant",
-        content: "",
-        state: "{}",
-      },
-    ];
-
-    rerender(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <ChatHarness tick={1} />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    await waitFor(() => {
-      const message = screen.getByTestId("message-msg-assistant-1");
-      expect(within(message).getByTestId("state").textContent).toBe("none");
-    });
-  });
-
-  it("renders an empty live state update", async () => {
-    const copilotContextValue = createTestCopilotContext();
-
-    render(
-      <CopilotContext.Provider value={copilotContextValue}>
-        <CoAgentStateRendersProvider>
-          <LiveStateHarness
-            message={{ id: "msg-live-empty", role: "assistant" }}
-          />
-        </CoAgentStateRendersProvider>
-      </CopilotContext.Provider>,
-    );
-
-    mockAgent.state = { current_step: "Processing..." };
-    act(() => {
-      lastSubscriber?.onStateChanged?.({ state: mockAgent.state });
-    });
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("Processing...");
-    });
-
-    mockAgent.state = {};
-    act(() => {
-      lastSubscriber?.onStateChanged?.({ state: mockAgent.state });
-    });
-
-    await waitFor(() => {
-      expect(screen.getByTestId("state").textContent).toBe("none");
-    });
-  });
-});