@copilotkit/react-core 1.55.0-next.8 → 1.55.0-next.9

package/src/v2/hooks/__tests__/use-agent-thread-isolation.test.tsx

@@ -0,0 +1,327 @@
+import React from "react";
+import { render } from "@testing-library/react";
+import { renderHook } from "@testing-library/react";
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+import {
+  AbstractAgent,
+  type BaseEvent,
+  type RunAgentInput,
+} from "@ag-ui/client";
+import { useCopilotKit } from "../../providers/CopilotKitProvider";
+import { useAgent } from "../use-agent";
+import { CopilotKitCoreRuntimeConnectionStatus } from "@copilotkit/core";
+import { Observable } from "rxjs";
+
+vi.mock("../../providers/CopilotKitProvider", () => ({
+  useCopilotKit: vi.fn(),
+}));
+
+const mockUseCopilotKit = useCopilotKit as ReturnType<typeof vi.fn>;
+
+/**
+ * A minimal mock agent whose clone() returns a NEW instance and copies
+ * messages from the source. This is essential for testing per-thread
+ * isolation — each clone must be a distinct object that starts with the
+ * source's state so that cloneForThread's setMessages([]) / setState({})
+ * calls are meaningful (not vacuously true on an already-empty clone).
+ */
+class CloneableAgent extends AbstractAgent {
+  clone(): CloneableAgent {
+    const cloned = new CloneableAgent();
+    cloned.agentId = this.agentId;
+    // Copy messages so cloneForThread's setMessages([]) actually clears state
+    cloned.setMessages([...this.messages]);
+    return cloned;
+  }
+
+  run(_input: RunAgentInput): Observable<BaseEvent> {
+    return new Observable();
+  }
+}
+
+describe("useAgent thread isolation", () => {
+  let mockCopilotkit: {
+    getAgent: ReturnType<typeof vi.fn>;
+    runtimeUrl: string | undefined;
+    runtimeConnectionStatus: CopilotKitCoreRuntimeConnectionStatus;
+    runtimeTransport: string;
+    headers: Record<string, string>;
+    agents: Record<string, AbstractAgent>;
+  };
+
+  let registeredAgent: CloneableAgent;
+
+  beforeEach(() => {
+    registeredAgent = new CloneableAgent();
+    registeredAgent.agentId = "my-agent";
+
+    mockCopilotkit = {
+      getAgent: vi.fn((id: string) =>
+        id === "my-agent" ? registeredAgent : undefined,
+      ),
+      runtimeUrl: "http://localhost:3000/api/copilotkit",
+      runtimeConnectionStatus: CopilotKitCoreRuntimeConnectionStatus.Connected,
+      runtimeTransport: "rest",
+      headers: {},
+      agents: { "my-agent": registeredAgent },
+    };
+
+    mockUseCopilotKit.mockReturnValue({
+      copilotkit: mockCopilotkit,
+      executingToolCallIds: new Set(),
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("returns different agent instances for different threadIds with the same agentId", () => {
+    const agents: Record<string, AbstractAgent> = {};
+
+    function TrackerA() {
+      const { agent } = useAgent({ agentId: "my-agent", threadId: "thread-a" });
+      agents["a"] = agent;
+      return null;
+    }
+
+    function TrackerB() {
+      const { agent } = useAgent({ agentId: "my-agent", threadId: "thread-b" });
+      agents["b"] = agent;
+      return null;
+    }
+
+    render(
+      <>
+        <TrackerA />
+        <TrackerB />
+      </>,
+    );
+
+    expect(agents["a"]).toBeDefined();
+    expect(agents["b"]).toBeDefined();
+    expect(agents["a"]).not.toBe(agents["b"]);
+  });
+
+  it("returns the same cached instance for the same (agentId, threadId) across re-renders", () => {
+    const instances: AbstractAgent[] = [];
+
+    function Tracker() {
+      const { agent } = useAgent({ agentId: "my-agent", threadId: "thread-x" });
+      instances.push(agent);
+      return null;
+    }
+
+    const { rerender } = render(<Tracker />);
+    rerender(<Tracker />);
+
+    expect(instances.length).toBe(2);
+    expect(instances[0]).toBe(instances[1]);
+  });
+
+  it("returns the shared registry agent when no threadId is provided (backward compat)", () => {
+    let captured: AbstractAgent | undefined;
+
+    function Tracker() {
+      const { agent } = useAgent({ agentId: "my-agent" });
+      captured = agent;
+      return null;
+    }
+
+    render(<Tracker />);
+    expect(captured).toBe(registeredAgent);
+  });
+
+  it("isolates messages between thread-specific agents", () => {
+    // Pre-populate the source agent so CloneableAgent.clone() copies the
+    // message into each clone — this makes cloneForThread's setMessages([])
+    // meaningful rather than vacuously true on an already-empty clone.
+    registeredAgent.addMessage({
+      id: "source-msg",
+      role: "user",
+      content: "pre-existing on source",
+    });
+
+    const agents: Record<string, AbstractAgent> = {};
+
+    function TrackerA() {
+      const { agent } = useAgent({ agentId: "my-agent", threadId: "thread-a" });
+      agents["a"] = agent;
+      return null;
+    }
+
+    function TrackerB() {
+      const { agent } = useAgent({ agentId: "my-agent", threadId: "thread-b" });
+      agents["b"] = agent;
+      return null;
+    }
+
+    render(
+      <>
+        <TrackerA />
+        <TrackerB />
+      </>,
+    );
+
+    // Both clones should start empty even though the source had a message —
+    // cloneForThread must have called setMessages([]) on each clone.
+    expect(agents["a"]!.messages).toHaveLength(0);
+    expect(agents["b"]!.messages).toHaveLength(0);
+
+    // Adding a message to thread A must not affect thread B
+    agents["a"]!.addMessage({
+      id: "msg-1",
+      role: "user",
+      content: "hello from thread A",
+    });
+
+    expect(agents["a"]!.messages).toHaveLength(1);
+    expect(agents["b"]!.messages).toHaveLength(0);
+  });
+
+  it("sets threadId on cloned agents", () => {
+    const agents: Record<string, AbstractAgent> = {};
+
+    function TrackerA() {
+      const { agent } = useAgent({ agentId: "my-agent", threadId: "thread-a" });
+      agents["a"] = agent;
+      return null;
+    }
+
+    function TrackerB() {
+      const { agent } = useAgent({ agentId: "my-agent", threadId: "thread-b" });
+      agents["b"] = agent;
+      return null;
+    }
+
+    render(
+      <>
+        <TrackerA />
+        <TrackerB />
+      </>,
+    );
+
+    expect(agents["a"]!.threadId).toBe("thread-a");
+    expect(agents["b"]!.threadId).toBe("thread-b");
+  });
+
+  it("invalidates stale clone when the registry agent is replaced", () => {
+    // Simulates reconnect / hot-reload: copilotkit.agents holds a new object.
+    const { result, rerender } = renderHook(
+      ({ tid }: { tid: string }) =>
+        useAgent({ agentId: "my-agent", threadId: tid }),
+      { initialProps: { tid: "thread-a" } },
+    );
+
+    const firstClone = result.current.agent;
+    expect(firstClone).not.toBe(registeredAgent); // it's a clone
+
+    // Replace the registry agent
+    const replacementAgent = new CloneableAgent();
+    replacementAgent.agentId = "my-agent";
+
+    mockCopilotkit.agents = { "my-agent": replacementAgent };
+    mockCopilotkit.getAgent.mockImplementation((id: string) =>
+      id === "my-agent" ? replacementAgent : undefined,
+    );
+    mockUseCopilotKit.mockReturnValue({
+      copilotkit: { ...mockCopilotkit },
+      executingToolCallIds: new Set(),
+    });
+
+    rerender({ tid: "thread-a" });
+
+    const secondClone = result.current.agent;
+    expect(secondClone).not.toBe(firstClone); // stale clone was invalidated
+    expect(secondClone).not.toBe(replacementAgent); // still a clone, not the source
+  });
+
+  it("switching threadId returns a fresh clone; switching back returns the cached one", () => {
+    const { result, rerender } = renderHook(
+      ({ tid }: { tid: string }) =>
+        useAgent({ agentId: "my-agent", threadId: tid }),
+      { initialProps: { tid: "thread-a" } },
+    );
+
+    const cloneA = result.current.agent;
+
+    rerender({ tid: "thread-b" });
+    const cloneB = result.current.agent;
+    expect(cloneB).not.toBe(cloneA);
+
+    // Switching back to thread-a should return the originally cached clone
+    rerender({ tid: "thread-a" });
+    expect(result.current.agent).toBe(cloneA);
+  });
+
+  it("uses a fresh clone with correct threadId when provisional transitions to real agent", () => {
+    // Start in Disconnected state — a provisional is created
+    mockCopilotkit.runtimeConnectionStatus =
+      CopilotKitCoreRuntimeConnectionStatus.Disconnected;
+    mockCopilotkit.getAgent.mockReturnValue(undefined);
+    mockCopilotkit.agents = {};
+    mockUseCopilotKit.mockReturnValue({
+      copilotkit: { ...mockCopilotkit },
+      executingToolCallIds: new Set(),
+    });
+
+    const { result, rerender } = renderHook(() =>
+      useAgent({ agentId: "my-agent", threadId: "thread-a" }),
+    );
+
+    const provisional = result.current.agent;
+    expect(provisional.threadId).toBe("thread-a");
+
+    // Real agent appears (runtime connected and agent registered)
+    mockCopilotkit.runtimeConnectionStatus =
+      CopilotKitCoreRuntimeConnectionStatus.Connected;
+    mockCopilotkit.getAgent.mockImplementation((id: string) =>
+      id === "my-agent" ? registeredAgent : undefined,
+    );
+    mockCopilotkit.agents = { "my-agent": registeredAgent };
+    mockUseCopilotKit.mockReturnValue({
+      copilotkit: { ...mockCopilotkit },
+      executingToolCallIds: new Set(),
+    });
+
+    rerender();
+
+    const realClone = result.current.agent;
+    expect(realClone).not.toBe(provisional); // provisional replaced by real clone
+    expect(realClone).not.toBe(registeredAgent); // it's a clone, not the source
+    expect(realClone.threadId).toBe("thread-a");
+  });
+
+  it("uses composite key for provisional agents when threadId is provided", () => {
+    // Put runtime in Disconnected state so provisionals are created
+    mockCopilotkit.runtimeConnectionStatus =
+      CopilotKitCoreRuntimeConnectionStatus.Disconnected;
+    mockCopilotkit.getAgent.mockReturnValue(undefined);
+    mockCopilotkit.agents = {};
+
+    const agents: Record<string, AbstractAgent> = {};
+
+    function TrackerA() {
+      const { agent } = useAgent({ agentId: "my-agent", threadId: "thread-a" });
+      agents["a"] = agent;
+      return null;
+    }
+
+    function TrackerB() {
+      const { agent } = useAgent({ agentId: "my-agent", threadId: "thread-b" });
+      agents["b"] = agent;
+      return null;
+    }
+
+    render(
+      <>
+        <TrackerA />
+        <TrackerB />
+      </>,
+    );
+
+    expect(agents["a"]).not.toBe(agents["b"]);
+    expect(agents["a"]!.threadId).toBe("thread-a");
+    expect(agents["b"]!.threadId).toBe("thread-b");
+  });
+});

package/src/v2/hooks/__tests__/use-agent.e2e.test.tsx

@@ -19,10 +19,21 @@ import { CopilotChat } from "../../components/chat/CopilotChat";
  * Mock agent that captures RunAgentInput to verify state is passed correctly
  */
 class StateCapturingMockAgent extends MockStepwiseAgent {
-  public lastRunInput?: RunAgentInput;
+  // Shared via a container so the clone and original both see the same value
+  private _capture: { lastRunInput?: RunAgentInput } = {};
+
+  get lastRunInput(): RunAgentInput | undefined {
+    return this._capture.lastRunInput;
+  }
+
+  clone(): this {
+    const cloned = super.clone();
+    (cloned as unknown as StateCapturingMockAgent)._capture = this._capture;
+    return cloned;
+  }
 
   run(input: RunAgentInput): Observable<BaseEvent> {
-    this.lastRunInput = input;
+    this._capture.lastRunInput = input;
     return super.run(input);
   }
 }

package/src/v2/hooks/__tests__/use-frontend-tool.e2e.test.tsx

@@ -502,13 +502,24 @@ describe("useFrontendTool E2E - Dynamic Registration", () => {
   describe("Agent input plumbing", () => {
     it("forwards registered frontend tools to runAgent input", async () => {
       class InstrumentedMockAgent extends MockStepwiseAgent {
-        public lastRunParameters?: RunAgentParameters;
+        // Shared so the clone and original both see the captured parameters
+        private _capture: { lastRunParameters?: RunAgentParameters } = {};
+
+        get lastRunParameters(): RunAgentParameters | undefined {
+          return this._capture.lastRunParameters;
+        }
+
+        clone(): this {
+          const cloned = super.clone();
+          (cloned as unknown as InstrumentedMockAgent)._capture = this._capture;
+          return cloned;
+        }
 
         async runAgent(
           parameters?: RunAgentParameters,
           subscriber?: AgentSubscriber,
         ) {
-          this.lastRunParameters = parameters;
+          this._capture.lastRunParameters = parameters;
           return super.runAgent(parameters, subscriber);
         }
       }
@@ -568,8 +579,16 @@ describe("useFrontendTool E2E - Dynamic Registration", () => {
       class OneShotToolCallAgent extends AbstractAgent {
         private runCount = 0;
         clone(): OneShotToolCallAgent {
-          // Keep state across runs so the second run emits different args
-          return this;
+          const cloned = new OneShotToolCallAgent();
+          cloned.agentId = this.agentId;
+          // Share runCount via reference so the second run emits different args
+          Object.defineProperty(cloned, "runCount", {
+            get: () => this.runCount,
+            set: (v: number) => {
+              this.runCount = v;
+            },
+          });
+          return cloned;
         }
         async detachActiveRun(): Promise<void> {}
         run(_input: RunAgentInput): Observable<BaseEvent> {

package/src/v2/hooks/use-agent.tsx

@@ -1,7 +1,8 @@
 import { useCopilotKit } from "../providers/CopilotKitProvider";
+import { useCopilotChatConfiguration } from "../providers/CopilotChatConfigurationProvider";
 import { useMemo, useEffect, useReducer, useRef } from "react";
 import { DEFAULT_AGENT_ID } from "@copilotkit/shared";
-import { AbstractAgent } from "@ag-ui/client";
+import { AbstractAgent, HttpAgent } from "@ag-ui/client";
 import {
   ProxiedCopilotRuntimeAgent,
   CopilotKitCoreRuntimeConnectionStatus,
@@ -21,13 +22,88 @@ const ALL_UPDATES: UseAgentUpdate[] = [
 
 export interface UseAgentProps {
   agentId?: string;
+  threadId?: string;
   updates?: UseAgentUpdate[];
 }
 
-export function useAgent({ agentId, updates }: UseAgentProps = {}) {
+/**
+ * Clone a registry agent for per-thread isolation.
+ * Copies agent configuration (transport, headers, etc.) but resets conversation
+ * state (messages, threadId, state) so each thread starts fresh.
+ */
+function cloneForThread(
+  source: AbstractAgent,
+  threadId: string,
+  headers: Record<string, string>,
+): AbstractAgent {
+  const clone = source.clone();
+  if (clone === source) {
+    throw new Error(
+      `useAgent: ${source.constructor.name}.clone() returned the same instance. ` +
+        `clone() must return a new, independent object.`,
+    );
+  }
+  clone.threadId = threadId;
+  clone.setMessages([]);
+  clone.setState({});
+  if (clone instanceof HttpAgent) {
+    clone.headers = { ...headers };
+  }
+  return clone;
+}
+
+/**
+ * Module-level WeakMap: registryAgent → (threadId → clone).
+ * Shared across all useAgent() calls so that every component using the same
+ * (agentId, threadId) pair receives the same agent instance. Using WeakMap
+ * ensures the clone map is garbage-collected when the registry agent is
+ * replaced (e.g. after reconnect or hot-reload).
+ */
+export const globalThreadCloneMap = new WeakMap<
+  AbstractAgent,
+  Map<string, AbstractAgent>
+>();
+
+/**
+ * Look up an existing per-thread clone without creating one.
+ * Returns undefined when no clone has been created yet for this pair.
+ */
+export function getThreadClone(
+  registryAgent: AbstractAgent | undefined | null,
+  threadId: string | undefined | null,
+): AbstractAgent | undefined {
+  if (!registryAgent || !threadId) return undefined;
+  return globalThreadCloneMap.get(registryAgent)?.get(threadId);
+}
+
+function getOrCreateThreadClone(
+  existing: AbstractAgent,
+  threadId: string,
+  headers: Record<string, string>,
+): AbstractAgent {
+  let byThread = globalThreadCloneMap.get(existing);
+  if (!byThread) {
+    byThread = new Map();
+    globalThreadCloneMap.set(existing, byThread);
+  }
+  const cached = byThread.get(threadId);
+  if (cached) return cached;
+
+  const clone = cloneForThread(existing, threadId, headers);
+  byThread.set(threadId, clone);
+  return clone;
+}
+
+export function useAgent({ agentId, threadId, updates }: UseAgentProps = {}) {
   agentId ??= DEFAULT_AGENT_ID;
 
   const { copilotkit } = useCopilotKit();
+  // Fall back to the enclosing CopilotChatConfigurationProvider's threadId so
+  // that useAgent() called without explicit threadId (e.g. inside a custom
+  // message renderer) automatically uses the same per-thread clone as the
+  // CopilotChat component it lives within.
+  const chatConfig = useCopilotChatConfiguration();
+  threadId ??= chatConfig?.threadId;
   const [, forceUpdate] = useReducer((x) => x + 1, 0);
 
   const updateFlags = useMemo(
@@ -43,11 +119,29 @@ export function useAgent({ agentId, updates }: UseAgentProps = {}) {
   );
 
   const agent: AbstractAgent = useMemo(() => {
+    // Use a composite key when threadId is provided so that different threads
+    // for the same agent get independent instances.
+    const cacheKey = threadId ? `${agentId}:${threadId}` : agentId;
+
     const existing = copilotkit.getAgent(agentId);
     if (existing) {
-      // Real agent found — clear any cached provisional for this ID
+      // Real agent found — clear any cached provisionals for this key and the
+      // bare agentId key (handles the case where a provisional was created
+      // before threadId was available, then the component re-renders with one).
+      provisionalAgentCache.current.delete(cacheKey);
       provisionalAgentCache.current.delete(agentId);
-      return existing;
+
+      if (!threadId) {
+        // No threadId — return the shared registry agent (original behavior)
+        return existing;
+      }
+
+      // threadId provided — return the shared per-thread clone.
+      // The global WeakMap ensures all components using the same
+      // (registryAgent, threadId) pair receive the same instance, so state
+      // mutations (addMessage, setState) are visible everywhere. The WeakMap
+      // entry is GC-collected automatically when the registry agent is replaced.
+      return getOrCreateThreadClone(existing, threadId, copilotkit.headers);
     }
 
     const isRuntimeConfigured = copilotkit.runtimeUrl !== undefined;
@@ -60,7 +154,7 @@ export function useAgent({ agentId, updates }: UseAgentProps = {}) {
         status === CopilotKitCoreRuntimeConnectionStatus.Connecting)
     ) {
       // Return cached provisional if available (keeps reference stable)
-      const cached = provisionalAgentCache.current.get(agentId);
+      const cached = provisionalAgentCache.current.get(cacheKey);
       if (cached) {
         // Update headers on the cached agent in case they changed
         cached.headers = { ...copilotkit.headers };
@@ -75,7 +169,10 @@ export function useAgent({ agentId, updates }: UseAgentProps = {}) {
       });
       // Apply current headers so runs/connects inherit them
       provisional.headers = { ...copilotkit.headers };
-      provisionalAgentCache.current.set(agentId, provisional);
+      if (threadId) {
+        provisional.threadId = threadId;
+      }
+      provisionalAgentCache.current.set(cacheKey, provisional);
       return provisional;
     }
 
@@ -88,6 +185,14 @@ export function useAgent({ agentId, updates }: UseAgentProps = {}) {
       isRuntimeConfigured &&
       status === CopilotKitCoreRuntimeConnectionStatus.Error
     ) {
+      // Cache the provisional so that dep changes while in Error state (e.g.
+      // headers update) return the same agent reference, matching the
+      // Disconnected/Connecting path and preventing spurious re-subscriptions.
+      const cached = provisionalAgentCache.current.get(cacheKey);
+      if (cached) {
+        cached.headers = { ...copilotkit.headers };
+        return cached;
+      }
       const provisional = new ProxiedCopilotRuntimeAgent({
         runtimeUrl: copilotkit.runtimeUrl,
         agentId,
@@ -95,6 +200,10 @@ export function useAgent({ agentId, updates }: UseAgentProps = {}) {
         runtimeMode: "pending",
       });
       provisional.headers = { ...copilotkit.headers };
+      if (threadId) {
+        provisional.threadId = threadId;
+      }
+      provisionalAgentCache.current.set(cacheKey, provisional);
       return provisional;
     }
 
@@ -113,6 +222,7 @@ export function useAgent({ agentId, updates }: UseAgentProps = {}) {
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [
     agentId,
+    threadId,
     copilotkit.agents,
     copilotkit.runtimeConnectionStatus,
     copilotkit.runtimeUrl,
@@ -149,6 +259,16 @@ export function useAgent({ agentId, updates }: UseAgentProps = {}) {
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [agent, forceUpdate, JSON.stringify(updateFlags)]);
 
+  // Keep HttpAgent headers fresh without mutating inside useMemo, which is
+  // unsafe in concurrent mode (React may invoke useMemo multiple times and
+  // discard intermediate results, but mutations always land).
+  useEffect(() => {
+    if (agent instanceof HttpAgent) {
+      agent.headers = { ...copilotkit.headers };
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [agent, JSON.stringify(copilotkit.headers)]);
+
   return {
     agent,
   };

package/src/v2/hooks/use-render-custom-messages.tsx

@@ -1,4 +1,5 @@
 import { useCopilotChatConfiguration, useCopilotKit } from "../providers";
+import { getThreadClone } from "./use-agent";
 import { ReactCustomMessageRendererPosition } from "../types/react-custom-message-renderer";
 import { Message } from "@ag-ui/core";
 
@@ -38,7 +39,11 @@ export function useRenderCustomMessages() {
       copilotkit.getRunIdForMessage(agentId, threadId, message.id) ??
       copilotkit.getRunIdsForThread(agentId, threadId).slice(-1)[0];
     const runId = resolvedRunId ?? `missing-run-id:${message.id}`;
-    const agent = copilotkit.getAgent(agentId);
+    // Prefer the per-thread clone so that agent.messages reflects the actual
+    // conversation state (messages live on the clone, not the registry agent).
+    // Fall back to the registry agent when no clone exists (no threadId).
+    const registryAgent = copilotkit.getAgent(agentId);
+    const agent = getThreadClone(registryAgent, threadId) ?? registryAgent;
     if (!agent) {
       throw new Error("Agent not found");
     }