refactor: remove per-thread agent cloning, restore single registry agent per id

Reverts the cloning design from #3525 (useAgent per-thread clones, getThreadClone,
globalThreadCloneMap, cloneForThread) and #3630 (clone routing in activity renderers),
plus the inspector machinery that existed only to handle clones (onAgentRunStarted
subscriber + run-handler emissions from #3869, the connect-time emission from #3872,
and the agentRunThreadId map that read from it).

State-manager isClone composite-key path and SuggestionEngine consumerAgent param —
both added in #3525 to keep clones visible to bookkeeping — are gone too.

Restores agent.threadId = resolvedThreadId in CopilotChat (pre-#3525 behavior) and
swaps the inspector's agentRunThreadId map for a direct agent.threadId read.

Removes the DemoButtonAgent and /a2ui-demo page from the demo (added by #3630 as a
clone-fix repro).

Re-opens the original issue #2957 (CPK-7155): two CopilotChat instances with the same
agentId and different threadIds will share message state again. The follow-up is a
public registerProxiedAgent API so callers can opt into multiple frontend agents
proxying to the same runtime agent, without implicit per-thread cloning.
This commit is contained in:
Markus Ecker
2026-05-04 13:27:04 +02:00
parent db595f8d8b
commit 762370a4e5
18 changed files with 42 additions and 880 deletions
@@ -1,21 +0,0 @@
"use client";
import { CopilotChat, CopilotKitProvider } from "@copilotkit/react-core/v2";
export const dynamic = "force-dynamic";
export default function A2UIDemoPage() {
return (
<CopilotKitProvider
runtimeUrl="/api/copilotkit"
a2ui={{}}
showDevConsole="auto"
>
<div
style={{ height: "100vh", margin: 0, padding: 0, overflow: "hidden" }}
>
<CopilotChat agentId="demo-button" threadId="a2ui-demo-thread" />
</div>
</CopilotKitProvider>
);
}
@@ -7,13 +7,6 @@ import {
import { TranscriptionServiceOpenAI } from "@copilotkit/voice";
import { handle } from "hono/vercel";
import OpenAI from "openai";
import {
AbstractAgent,
EventType,
type RunAgentInput,
type BaseEvent,
} from "@ag-ui/client";
import { Observable } from "rxjs";
const determineModel = () => {
if (process.env.OPENAI_API_KEY?.trim()) {
@@ -42,141 +35,6 @@ const builtInAgent = new BuiltInAgent({
},
});
/**
* Minimal demo agent for reproducing the A2UI thread-clone bug.
*
* First run → emits an A2UI surface with a "Confirm" button.
* Button click → runAgent fires again with forwardedProps.a2uiAction set,
* and the agent replies with a text message.
*
* Bug (before fix): the response appears on the registry agent's messages,
* not the per-thread clone, so CopilotChat never re-renders.
* Fix: useRenderActivityMessage now passes the clone to ReactSurfaceHost.
*/
class DemoButtonAgent extends AbstractAgent {
run(input: RunAgentInput): Observable<BaseEvent> {
return new Observable((observer) => {
const emit = (event: BaseEvent) => observer.next(event);
emit({
type: EventType.RUN_STARTED,
threadId: input.threadId,
runId: input.runId,
} as BaseEvent);
const a2uiAction = (input.forwardedProps as Record<string, any>)
?.a2uiAction;
if (a2uiAction) {
// Button was clicked — respond with a text message.
// Without the fix this message lands on the registry agent and never
// shows in chat. With the fix it lands on the per-thread clone.
const msgId = crypto.randomUUID();
emit({
type: EventType.TEXT_MESSAGE_START,
messageId: msgId,
role: "assistant",
} as BaseEvent);
emit({
type: EventType.TEXT_MESSAGE_CONTENT,
messageId: msgId,
delta: `✅ Confirmed! (thread: ${input.threadId}) — if you can read this, the fix is working.`,
} as BaseEvent);
emit({
type: EventType.TEXT_MESSAGE_END,
messageId: msgId,
} as BaseEvent);
} else {
// First run — render an A2UI surface with a Confirm button.
const activityId = crypto.randomUUID();
emit({
type: EventType.ACTIVITY_SNAPSHOT,
messageId: activityId,
activityType: "a2ui-surface",
content: {
operations: [
{
beginRendering: {
surfaceId: "demo-surface",
root: "container",
},
},
{
surfaceUpdate: {
surfaceId: "demo-surface",
components: [
{
id: "container",
component: {
Column: {
children: {
explicitList: ["prompt-text", "confirm-btn"],
},
},
},
},
{
id: "prompt-text",
component: {
Text: {
text: {
literalString:
"Click the button to trigger a response:",
},
},
},
},
{
id: "btn-label",
component: {
Text: { text: { literalString: "Confirm" } },
},
},
{
id: "confirm-btn",
component: {
Button: {
child: "btn-label",
action: { name: "confirm" },
primary: true,
},
},
},
],
},
},
],
},
} as BaseEvent);
const msgId = crypto.randomUUID();
emit({
type: EventType.TEXT_MESSAGE_START,
messageId: msgId,
role: "assistant",
} as BaseEvent);
emit({
type: EventType.TEXT_MESSAGE_CONTENT,
messageId: msgId,
delta:
"Click the button above. Without the fix, the confirmation message won't appear here.",
} as BaseEvent);
emit({
type: EventType.TEXT_MESSAGE_END,
messageId: msgId,
} as BaseEvent);
}
emit({ type: EventType.RUN_FINISHED } as BaseEvent);
observer.complete();
});
}
clone(): AbstractAgent {
return new DemoButtonAgent();
}
}
// Set up transcription service if OpenAI API key is available
const transcriptionService = process.env.OPENAI_API_KEY?.trim()
? new TranscriptionServiceOpenAI({
@@ -187,7 +45,6 @@ const transcriptionService = process.env.OPENAI_API_KEY?.trim()
const honoRuntime = new CopilotRuntime({
agents: {
default: builtInAgent,
"demo-button": new DemoButtonAgent(),
},
runner: new InMemoryAgentRunner(),
transcriptionService,
-22
View File
@@ -185,16 +185,6 @@ export interface CopilotKitCoreSubscriber {
agentId: string;
prevStore: ɵThreadStore;
}) => void | Promise<void>;
/**
* Fired immediately before each agent run, including per-thread clones that
* are not in the agent registry and therefore not surfaced via onAgentsChanged.
* Subscribers that track agent events (e.g. the web inspector) can use this
* to subscribe to the concrete agent instance that will emit events.
*/
onAgentRunStarted?: (event: {
copilotkit: CopilotKitCore;
agent: AbstractAgent;
}) => void | Promise<void>;
}
// Subscription object returned by subscribe() and subscribeToAgentWithOptions()
@@ -322,12 +312,6 @@ export interface CopilotKitCoreFriendsAccess {
*/
waitForPendingFrameworkUpdates(): Promise<void>;
/**
* Subscribe the state manager to an agent (including per-thread clones).
* Called by RunHandler before executing an agent so that events from
* clones are tracked in stateByRun/messageToRun.
*/
subscribeAgentToStateManager(agent: AbstractAgent): void;
}
export class CopilotKitCore {
@@ -1009,12 +993,6 @@ export class CopilotKitCore {
return this.stateManager.getRunIdsForThread(agentId, threadId);
}
subscribeAgentToStateManager(agent: AbstractAgent): void {
// isClone: true — use composite agentId:threadId key, keeping the clone's
// subscription independent of the registry agent's bare-agentId subscription.
this.stateManager.subscribeToAgent(agent, { isClone: true });
}
/**
* Internal method used by RunHandler to build frontend tools
*/
-28
View File
@@ -206,18 +206,6 @@ export class RunHandler {
};
}
// Notify subscribers (e.g. the inspector) about the agent that is about
// to run. This is critical for per-thread clones that are not present in
// the agent registry and would otherwise be invisible to subscribers.
await this._internal.notifySubscribers(
(subscriber) =>
subscriber.onAgentRunStarted?.({
copilotkit: this.core,
agent,
}),
"Subscriber onAgentRunStarted error:",
);
const runAgentResult = await agent.connectAgent(
{
forwardedProps: this._internal.properties,
@@ -288,22 +276,6 @@ export class RunHandler {
await agent.detachActiveRun();
}
// Ensure the state manager is subscribed to this agent (handles per-thread
// clones that are not in the registry and therefore not subscribed via
// onAgentsChanged). The composite-key logic in StateManager means this
// does not overwrite the registry agent's subscription.
this._internal.subscribeAgentToStateManager(agent);
// Notify subscribers (e.g. the web inspector) that a run is about to start
// on this specific agent instance. Must be awaited so that subscribers can
// call agent.subscribe() before agent.runAgent() captures its subscriber
// snapshot — agent.runAgent() snapshots [this.subscribers] synchronously.
await this._internal.notifySubscribers(
(subscriber) =>
subscriber.onAgentRunStarted?.({ copilotkit: this.core, agent }),
"Subscriber onAgentRunStarted error:",
);
// Set up abort controller and agent.abortRun() intercept only for the
// top-level call. Recursive follow-up calls from processAgentResult
// reuse the same controller.
+8 -23
View File
@@ -39,31 +39,19 @@ export class StateManager {
/**
* Subscribe to an agent's events to track state and messages.
*
* Registry agents (subscribed via `onAgentsChanged`) use the bare `agentId`
* key so that `unsubscribeFromAgent(agentId)` can remove them when they
* are replaced. Per-thread clones (subscribed via `subscribeAgentToStateManager`)
* pass `{ isClone: true }` to use a composite `agentId:threadId` key, keeping
* their subscription independent of the registry agent's.
*/
subscribeToAgent(
agent: AbstractAgent,
/** @param isClone When true, uses a composite `agentId:threadId` key for per-thread isolation. */
{ isClone = false }: { isClone?: boolean } = {},
): void {
subscribeToAgent(agent: AbstractAgent): void {
if (!agent.agentId) {
return; // Skip agents without IDs
}
const agentId = agent.agentId;
const subscriptionKey =
isClone && agent.threadId ? `${agentId}:${agent.threadId}` : agentId;
// Unsubscribe existing subscription for this key only
const existingUnsubscribe = this.agentSubscriptions.get(subscriptionKey);
// Unsubscribe existing subscription for this agent only
const existingUnsubscribe = this.agentSubscriptions.get(agentId);
if (existingUnsubscribe) {
existingUnsubscribe();
this.agentSubscriptions.delete(subscriptionKey);
this.agentSubscriptions.delete(agentId);
}
// Subscribe to agent events.
@@ -79,8 +67,8 @@ export class StateManager {
// 2. Run isolation within one subscription: in tests (and edge cases), a new
// run's events can arrive through the same subscription before the new
// pipeline is set up. Concretely: the test emits RUN_STARTED for run2
// before copilotkit.runAgent() has called subscribeAgentToStateManager for
// run2. At that point S1 is still active and sees run2's events with
// before copilotkit.runAgent() has had a chance to set up the new
// pipeline. At that point S1 is still active and sees run2's events with
// input1.runId. To prevent both runs from sharing the same runId key, we
// detect the "seen RUN_FINISHED, then RUN_STARTED again" pattern and
// generate a fresh runId for the second logical run.
@@ -148,17 +136,14 @@ export class StateManager {
},
});
this.agentSubscriptions.set(subscriptionKey, () => {
this.agentSubscriptions.set(agentId, () => {
revoked = true;
unsubscribe();
});
}
/**
* Unsubscribe a registry agent's subscription (bare `agentId` key).
* Per-thread clone subscriptions use composite `agentId:threadId` keys and
* are replaced (not removed) by subsequent subscribeToAgent(agent, { isClone: true })
* calls for the same (agentId, threadId) pair.
* Unsubscribe an agent's subscription.
*/
unsubscribeFromAgent(agentId: string): void {
const unsubscribe = this.agentSubscriptions.get(agentId);
+8 -24
View File
@@ -52,28 +52,17 @@ export class SuggestionEngine {
/**
* Reload suggestions for a specific agent
* This triggers generation of new suggestions based on current configs
*
* @param agentId - The consumer agent ID
* @param consumerAgent - Optional: the specific agent instance whose messages should be used
* for availability filtering and context. When running with per-thread clones, the thread
* clone holds the conversation messages; passing it here ensures dynamic suggestions fire
* after the first message even though the registry agent has an empty message list.
*/
public reloadSuggestions(
agentId: string,
consumerAgent?: AbstractAgent,
): void {
public reloadSuggestions(agentId: string): void {
this.clearSuggestions(agentId);
// Use the provided agent instance when available; fall back to the registry agent.
// Per-thread clones hold the actual conversation messages; the registry agent does not.
// The agent may legitimately be missing here (e.g. runtime info still loading,
// or the consumer agent is configured but not yet registered) — static suggestions
// don't need it, only dynamic generation does. Treat that case as "no messages yet"
// and process static configs anyway.
const agent =
consumerAgent ??
(this.core as unknown as CopilotKitCoreFriendsAccess).getAgent(agentId);
const agent = (
this.core as unknown as CopilotKitCoreFriendsAccess
).getAgent(agentId);
const messageCount = agent?.messages?.length ?? 0;
let hasAnySuggestions = false;
@@ -106,7 +95,7 @@ export class SuggestionEngine {
hasAnySuggestions = true;
void this.notifySuggestionsStartedLoading(agentId);
}
void this.generateSuggestions(suggestionId, config, agentId, agent);
void this.generateSuggestions(suggestionId, config, agentId);
} else if (isStaticSuggestionsConfig(config)) {
this.addStaticSuggestions(suggestionId, config, agentId);
}
@@ -145,7 +134,6 @@ export class SuggestionEngine {
suggestionId: string,
config: DynamicSuggestionsConfig,
consumerAgentId: string,
consumerAgent?: AbstractAgent,
): Promise<void> {
let agent: AbstractAgent | undefined = undefined;
try {
@@ -157,13 +145,9 @@ export class SuggestionEngine {
`Suggestions provider agent not found: ${config.providerAgentId}`,
);
}
// Use the provided consumer agent when available (per-thread clone with actual messages);
// fall back to the registry agent for non-threaded use.
const suggestionsConsumerAgent =
consumerAgent ??
(this.core as unknown as CopilotKitCoreFriendsAccess).getAgent(
consumerAgentId,
);
const suggestionsConsumerAgent = (
this.core as unknown as CopilotKitCoreFriendsAccess
).getAgent(consumerAgentId);
if (!suggestionsConsumerAgent) {
throw new Error(
`Suggestions consumer agent not found: ${consumerAgentId}`,
@@ -179,10 +179,9 @@ describe("useCopilotChatInternal connectAgent guard", () => {
});
it("does not call connectAgent when threadId matches (same agent instance, no re-render)", async () => {
// useAgent now returns a per-thread clone, so the wrapper guards via
// lastConnectedAgentRef: connect fires once per agent instance, not once
// per render. After the first connect, further re-renders with the same
// agent do not trigger another connect.
// The wrapper guards via lastConnectedAgentRef: connect fires once per
// agent instance, not once per render. After the first connect, further
// re-renders with the same agent do not trigger another connect.
mockRuntimeConnectionStatus =
CopilotKitCoreRuntimeConnectionStatus.Connected;
mockAgent.threadId = "config-thread-id";
@@ -230,13 +229,13 @@ describe("useCopilotChatInternal connectAgent guard", () => {
});
});
it("passes config threadId to useAgent", () => {
it("passes resolved agentId to useAgent", () => {
applyMocks();
renderHook(() => useCopilotChatInternal(), { wrapper: createWrapper() });
expect(vi.mocked(useAgent)).toHaveBeenCalledWith(
expect.objectContaining({ threadId: "config-thread-id" }),
expect.objectContaining({ agentId: "test-agent" }),
);
});
});
@@ -339,7 +339,6 @@ export function useCopilotChatInternal({
const resolvedAgentId = existingConfig?.agentId ?? "default";
const { agent } = useAgent({
agentId: resolvedAgentId,
threadId: existingConfig?.threadId,
});
// Track the last agent instance we called connect() on. Without this,
@@ -118,7 +118,6 @@ export function CopilotChat({
const { agent } = useAgent({
agentId: resolvedAgentId,
threadId: resolvedThreadId,
throttleMs,
});
const { copilotkit } = useCopilotKit();
@@ -235,6 +234,8 @@ export function CopilotChat({
agent.abortController = connectAbortController;
}
agent.threadId = resolvedThreadId;
const connect = async (agent: AbstractAgent) => {
try {
await copilotkit.connectAgent({ agent });
@@ -22,7 +22,6 @@ import {
} from "@ag-ui/core";
import { twMerge } from "tailwind-merge";
import { useRenderActivityMessage, useRenderCustomMessages } from "../../hooks";
import { getThreadClone } from "../../hooks/use-agent";
import { useCopilotKit } from "../../providers/CopilotKitProvider";
import { useCopilotChatConfiguration } from "../../providers/CopilotChatConfigurationProvider";
@@ -392,18 +391,14 @@ export function CopilotChatMessageView({
// Subscribe to state changes so custom message renderers re-render when state updates.
useEffect(() => {
if (!config?.agentId) return;
const registryAgent = copilotkit.getAgent(config.agentId);
// Prefer the per-thread clone so that state changes from the running agent
// (which is the clone, not the registry) trigger re-renders.
const agent =
getThreadClone(registryAgent, config.threadId) ?? registryAgent;
const agent = copilotkit.getAgent(config.agentId);
if (!agent) return;
const subscription = agent.subscribe({
onStateChanged: forceUpdate,
});
return () => subscription.unsubscribe();
}, [config?.agentId, config?.threadId, copilotkit, forceUpdate]);
}, [config?.agentId, copilotkit, forceUpdate]);
// Subscribe to interrupt element changes for in-chat rendering.
const [interruptElement, setInterruptElement] =
@@ -9,8 +9,7 @@ import { DEFAULT_AGENT_ID } from "@copilotkit/shared";
/**
* Mock agent that records every connectAgent() invocation and resolves
* immediately with an empty run result. Tracking lives on the class so
* per-thread clones (from useAgent's WeakMap) share the counter.
* immediately with an empty run result.
*/
class TrackingAgent extends MockStepwiseAgent {
static connectCalls: Array<{
@@ -107,7 +106,6 @@ describe("CopilotChat welcome / connect integration", () => {
expect(TrackingAgent.connectCalls.length).toBeGreaterThan(0);
});
// The per-thread clone carries threadId; agentId is the default.
expect(
TrackingAgent.connectCalls.some((c) => c.threadId === "real-thread"),
).toBe(true);
@@ -17,9 +17,7 @@ import {
CopilotKitProvider,
useCopilotKit,
} from "../../../providers";
import type { AbstractAgent } from "@ag-ui/client";
import { IntelligenceAgent } from "@copilotkit/core";
import { getThreadClone } from "../../../hooks/use-agent";
import { createA2UIMessageRenderer } from "../../../a2ui/A2UIMessageRenderer";
import type { Theme } from "@copilotkit/a2ui-renderer";
import { CopilotChat } from "..";
@@ -306,65 +304,6 @@ describe("CopilotChat activity message rendering", () => {
expect(capturedCopilotkit).toBeDefined();
});
it("passes the per-thread clone (not the registry agent) to activity message renderers", async () => {
// Regression test for: A2UI button clicks firing runAgent on the registry
// agent instead of the per-thread clone that CopilotChat renders from.
// Caused by useRenderActivityMessage calling copilotkit.getAgent() directly
// instead of getThreadClone(registryAgent, threadId) ?? registryAgent.
const agent = new MockStepwiseAgent();
const agentId = "action-agent";
agent.agentId = agentId;
const threadId = "thread-for-action-test";
let capturedAgent: AbstractAgent | undefined;
const activityRenderer: ReactActivityMessageRenderer<{ label: string }> = {
activityType: "button-action",
content: z.object({ label: z.string() }),
render: ({ content, agent: renderedAgent }) => {
capturedAgent = renderedAgent;
return <button data-testid="action-button">{content.label}</button>;
},
};
renderWithCopilotKit({
agents: { [agentId]: agent },
agentId,
threadId,
renderActivityMessages: [activityRenderer],
});
const input = await screen.findByRole("textbox");
fireEvent.change(input, { target: { value: "show me buttons" } });
fireEvent.keyDown(input, { key: "Enter", code: "Enter" });
await waitFor(() => {
expect(screen.getByText("show me buttons")).toBeDefined();
});
agent.emit(runStartedEvent());
agent.emit(
activitySnapshotEvent({
messageId: testId("activity-action"),
activityType: "button-action",
content: { label: "Click Me" },
}),
);
agent.emit(runFinishedEvent());
await waitFor(() => {
expect(screen.getByTestId("action-button")).toBeDefined();
});
// CopilotChat creates a per-thread clone via useAgent. The activity renderer
// must receive that clone so that handleAction → runAgent targets the same
// instance chat is rendering from.
const clone = getThreadClone(agent, threadId);
expect(clone).toBeDefined();
expect(capturedAgent).toBe(clone);
expect(capturedAgent).not.toBe(agent); // must NOT be the registry agent
});
it("restores a completed A2UI surface after reconnect from an event-native baseline", async () => {
const agent = new MockReconnectableAgent();
const threadId = testId("a2ui-thread");
@@ -49,6 +49,11 @@ class MockMCPProxyAgent extends AbstractAgent {
this.runAgentResponses.set(method, response);
}
addMessage(msg: Parameters<AbstractAgent["addMessage"]>[0]) {
this.addMessageCalls.push(msg as any);
return super.addMessage(msg);
}
emit(event: BaseEvent) {
if (event.type === EventType.RUN_STARTED) {
this.isRunning = true;
@@ -70,66 +75,6 @@ class MockMCPProxyAgent extends AbstractAgent {
});
}
clone(): MockMCPProxyAgent {
const cloned = new MockMCPProxyAgent();
cloned.agentId = this.agentId;
type Internal = {
subject: Subject<BaseEvent>;
runAgentCalls: Array<{ input: Partial<RunAgentInput> }>;
addMessageCalls: Array<{ id: string; role: string; content: string }>;
runAgentResponses: Map<string, unknown>;
};
(cloned as unknown as Internal).subject = (
this as unknown as Internal
).subject;
(cloned as unknown as Internal).runAgentCalls = (
this as unknown as Internal
).runAgentCalls;
(cloned as unknown as Internal).addMessageCalls = (
this as unknown as Internal
).addMessageCalls;
(cloned as unknown as Internal).runAgentResponses = (
this as unknown as Internal
).runAgentResponses;
const registry = this;
Object.defineProperty(cloned, "isRunning", {
get() {
return registry.isRunning;
},
set(v: boolean) {
registry.isRunning = v;
},
configurable: true,
enumerable: true,
});
const proto = MockMCPProxyAgent.prototype;
cloned.runAgent = async function (
input?: Partial<RunAgentInput>,
): Promise<RunAgentResult> {
const proxiedRequest = input?.forwardedProps?.__proxiedMCPRequest;
if (proxiedRequest) {
return registry.runAgent(input);
}
return proto.runAgent.call(cloned, input);
};
// Track addMessage calls on the clone (the component uses the clone)
const origAddMessage = cloned.addMessage.bind(cloned);
cloned.addMessage = function (msg: Parameters<typeof origAddMessage>[0]) {
registry.addMessageCalls.push(msg as any);
return origAddMessage(msg);
};
// Proxy run() calls so spies on the registry's run() see clone invocations
cloned.run = function (input: RunAgentInput): Observable<BaseEvent> {
return registry.run(input);
};
return cloned;
}
async detachActiveRun(): Promise<void> {}
run(_input: RunAgentInput): Observable<BaseEvent> {
@@ -1,333 +0,0 @@
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>;
subscribeToAgentWithOptions: (
agent: AbstractAgent,
subscriber: any,
) => { unsubscribe: () => void };
};
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 },
subscribeToAgentWithOptions: (agent, subscriber) =>
agent.subscribe(subscriber),
};
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");
});
});
+6 -111
View File
@@ -1,5 +1,4 @@
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, HttpAgent } from "@ag-ui/client";
@@ -23,7 +22,6 @@ const ALL_UPDATES: UseAgentUpdate[] = [
export interface UseAgentProps {
agentId?: string;
threadId?: string;
updates?: UseAgentUpdate[];
/**
* Throttle interval (in milliseconds) for re-renders triggered by
@@ -50,77 +48,8 @@ export interface UseAgentProps {
throttleMs?: number;
}
/**
* 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,
throttleMs,
}: UseAgentProps = {}) {
@@ -131,12 +60,6 @@ export function useAgent({
// subscribeToAgentWithOptions reads it from the core instance, but React needs the dep
// to know when to re-subscribe.
const providerThrottleMs = copilotkit.defaultThrottleMs;
// 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);
@@ -153,29 +76,11 @@ export function useAgent({
);
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 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);
// Real agent found — clear any cached provisional for this ID
provisionalAgentCache.current.delete(agentId);
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);
return existing;
}
const isRuntimeConfigured = copilotkit.runtimeUrl !== undefined;
@@ -188,7 +93,7 @@ export function useAgent({
status === CopilotKitCoreRuntimeConnectionStatus.Connecting)
) {
// Return cached provisional if available (keeps reference stable)
const cached = provisionalAgentCache.current.get(cacheKey);
const cached = provisionalAgentCache.current.get(agentId);
if (cached) {
// Update headers on the cached agent in case they changed
cached.headers = { ...copilotkit.headers };
@@ -203,10 +108,7 @@ export function useAgent({
});
// Apply current headers so runs/connects inherit them
provisional.headers = { ...copilotkit.headers };
if (threadId) {
provisional.threadId = threadId;
}
provisionalAgentCache.current.set(cacheKey, provisional);
provisionalAgentCache.current.set(agentId, provisional);
return provisional;
}
@@ -219,10 +121,7 @@ export function useAgent({
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);
const cached = provisionalAgentCache.current.get(agentId);
if (cached) {
cached.headers = { ...copilotkit.headers };
return cached;
@@ -234,10 +133,7 @@ export function useAgent({
runtimeMode: "pending",
});
provisional.headers = { ...copilotkit.headers };
if (threadId) {
provisional.threadId = threadId;
}
provisionalAgentCache.current.set(cacheKey, provisional);
provisionalAgentCache.current.set(agentId, provisional);
return provisional;
}
@@ -256,7 +152,6 @@ export function useAgent({
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [
agentId,
threadId,
copilotkit.agents,
copilotkit.runtimeConnectionStatus,
copilotkit.runtimeUrl,
@@ -3,12 +3,10 @@ import { DEFAULT_AGENT_ID } from "@copilotkit/shared";
import { useCopilotKit, useCopilotChatConfiguration } from "../providers";
import { useCallback, useMemo } from "react";
import { ReactActivityMessageRenderer } from "../types";
import { getThreadClone } from "./use-agent";
export function useRenderActivityMessage() {
const { copilotkit } = useCopilotKit();
const config = useCopilotChatConfiguration();
const agentId = config?.agentId ?? DEFAULT_AGENT_ID;
const agentId = useCopilotChatConfiguration()?.agentId ?? DEFAULT_AGENT_ID;
const renderers = copilotkit.renderActivityMessages;
@@ -52,13 +50,7 @@ export function useRenderActivityMessage() {
}
const Component = renderer.render;
// Prefer the per-thread clone so that handleAction in ReactSurfaceHost
// calls runAgent on the same agent instance that CopilotChat renders from.
// Without this, button clicks accumulate messages on the registry agent
// while CopilotChat displays from the clone — responses appear to vanish.
const registryAgent = copilotkit.getAgent(agentId);
const agent =
getThreadClone(registryAgent, config?.threadId) ?? registryAgent;
const agent = copilotkit.getAgent(agentId);
return (
<Component
@@ -70,7 +62,7 @@ export function useRenderActivityMessage() {
/>
);
},
[agentId, config?.threadId, copilotkit, findRenderer],
[agentId, copilotkit, findRenderer],
);
return useMemo(
@@ -1,5 +1,4 @@
import { useCopilotChatConfiguration, useCopilotKit } from "../providers";
import { getThreadClone } from "./use-agent";
import { ReactCustomMessageRendererPosition } from "../types/react-custom-message-renderer";
import { Message } from "@ag-ui/core";
@@ -39,11 +38,7 @@ export function useRenderCustomMessages() {
copilotkit.getRunIdForMessage(agentId, threadId, message.id) ??
copilotkit.getRunIdsForThread(agentId, threadId).slice(-1)[0];
const runId = resolvedRunId ?? `missing-run-id:${message.id}`;
// 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;
const agent = copilotkit.getAgent(agentId);
if (!agent) {
return null;
}
+1 -19
View File
@@ -2345,11 +2345,6 @@ export class WebInspectorElement extends LitElement {
private agentSubscriptions: Map<string, () => void> = new Map();
private agentEvents: Map<string, InspectorEvent[]> = new Map();
private agentMessages: Map<string, InspectorMessage[]> = new Map();
// Tracks which thread each agent is currently running on. Populated from the
// agent instance handed to `onAgentRunStarted` so that, when that agent
// subsequently emits message changes, we can bump the per-thread live
// version (below) only for the thread the messages actually belong to.
private agentRunThreadId: Map<string, string> = new Map();
// Per-thread monotonic version that ticks every time an agent currently
// running on that thread emits a message change. `cpk-thread-details`
// watches this prop and re-fetches `/threads/:id/messages` when it changes,
@@ -2655,19 +2650,6 @@ export class WebInspectorElement extends LitElement {
this._threads = Array.from(this._threadsByAgent.values()).flat();
this.requestUpdate();
},
onAgentRunStarted: ({ agent }) => {
// Subscribe to the concrete agent instance about to run. This handles
// per-thread clones that are not in core.agents and therefore not
// reachable via onAgentsChanged. Replacing an existing subscription for
// the same agentId is safe: the previous instance emits no more events
// once a new run starts on a fresh clone.
this.subscribeToAgent(agent);
const runThreadId = (agent as { threadId?: string }).threadId;
if (agent.agentId && runThreadId) {
this.agentRunThreadId.set(agent.agentId, runThreadId);
}
this.requestUpdate();
},
} satisfies CopilotKitCoreSubscriber;
this.coreUnsubscribe = core.subscribe(this.coreSubscriber).unsubscribe;
@@ -3019,7 +3001,7 @@ export class WebInspectorElement extends LitElement {
// selected thread and re-fetches `/threads/:id/messages` when it ticks,
// so the conversation view stays in sync with the streaming agent
// without the parent re-implementing AG-UI → ConversationItem mapping.
const runThreadId = this.agentRunThreadId.get(agent.agentId);
const runThreadId = (agent as { threadId?: string }).threadId;
if (runThreadId) {
this.liveMessageVersion.set(
runThreadId,