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

package/src/v2/runtime/runner/__tests__/in-memory-runner.e2e.test.ts

@@ -1,775 +0,0 @@
-import { describe, it, expect } from "vitest";
-import { InMemoryAgentRunner } from "../in-memory";
-import {
-  AbstractAgent,
-  BaseEvent,
-  EventType,
-  Message,
-  RunAgentInput,
-  RunErrorEvent,
-  RunFinishedEvent,
-  RunStartedEvent,
-} from "@ag-ui/client";
-import { EMPTY, Subscription, firstValueFrom, from } from "rxjs";
-import { toArray } from "rxjs/operators";
-
-type RunCallbacks = {
-  onEvent: (event: { event: BaseEvent }) => void;
-  onNewMessage?: (args: { message: Message }) => void;
-  onRunStartedEvent?: () => void;
-};
-
-interface EmitAgentOptions {
-  events?: BaseEvent[];
-  emitDefaultRunStarted?: boolean;
-  includeRunFinished?: boolean;
-  runFinishedEvent?: RunFinishedEvent;
-  afterEvent?: (args: {
-    event: BaseEvent;
-    index: number;
-  }) => void | Promise<void>;
-}
-
-class EmitAgent extends AbstractAgent {
-  constructor(private readonly options: EmitAgentOptions = {}) {
-    super();
-  }
-
-  async runAgent(input: RunAgentInput, callbacks: RunCallbacks): Promise<void> {
-    const {
-      emitDefaultRunStarted = true,
-      includeRunFinished = true,
-      runFinishedEvent,
-      afterEvent,
-    } = this.options;
-    const scriptedEvents = this.options.events ?? [];
-
-    let index = 0;
-    const emit = async (event: BaseEvent) => {
-      callbacks.onEvent({ event });
-      if (event.type === EventType.RUN_STARTED) {
-        callbacks.onRunStartedEvent?.();
-      }
-      await afterEvent?.({ event, index });
-      index += 1;
-    };
-
-    if (emitDefaultRunStarted) {
-      const runStarted: RunStartedEvent = {
-        type: EventType.RUN_STARTED,
-        threadId: input.threadId,
-        runId: input.runId,
-        parentRunId: input.parentRunId,
-      };
-      await emit(runStarted);
-    }
-
-    for (const event of scriptedEvents) {
-      await emit(event);
-    }
-
-    const hasRunFinishedEvent =
-      scriptedEvents.some((event) => event.type === EventType.RUN_FINISHED) ||
-      runFinishedEvent?.type === EventType.RUN_FINISHED;
-
-    if (includeRunFinished && !hasRunFinishedEvent) {
-      const finishEvent: RunFinishedEvent = runFinishedEvent ?? {
-        type: EventType.RUN_FINISHED,
-        threadId: input.threadId,
-        runId: input.runId,
-      };
-      await emit(finishEvent);
-    }
-  }
-
-  clone(): AbstractAgent {
-    return new EmitAgent({
-      ...this.options,
-      events: this.options.events ? [...this.options.events] : undefined,
-    });
-  }
-
-  protected run(): ReturnType<AbstractAgent["run"]> {
-    return EMPTY;
-  }
-
-  protected connect(): ReturnType<AbstractAgent["connect"]> {
-    return EMPTY;
-  }
-}
-
-class ReplayAgent extends AbstractAgent {
-  constructor(
-    private readonly replayEvents: BaseEvent[],
-    threadId: string,
-  ) {
-    super({ threadId });
-  }
-
-  async runAgent(): Promise<void> {
-    throw new Error("not used");
-  }
-
-  protected run(): ReturnType<AbstractAgent["run"]> {
-    return EMPTY;
-  }
-
-  protected connect(): ReturnType<AbstractAgent["connect"]> {
-    return from(this.replayEvents);
-  }
-}
-
-class RunnerConnectAgent extends AbstractAgent {
-  constructor(
-    private readonly runner: InMemoryAgentRunner,
-    threadId: string,
-  ) {
-    super({ threadId });
-  }
-
-  async runAgent(): Promise<void> {
-    throw new Error("not used");
-  }
-
-  protected run(): ReturnType<AbstractAgent["run"]> {
-    return EMPTY;
-  }
-
-  protected connect(
-    input: RunAgentInput,
-  ): ReturnType<AbstractAgent["connect"]> {
-    return this.runner.connect({ threadId: input.threadId });
-  }
-}
-
-type Deferred<T> = {
-  promise: Promise<T>;
-  resolve: (value: T) => void;
-  reject: (reason?: unknown) => void;
-};
-
-function createDeferred<T>(): Deferred<T> {
-  let resolve!: (value: T) => void;
-  let reject!: (reason?: unknown) => void;
-  const promise = new Promise<T>((res, rej) => {
-    resolve = res;
-    reject = rej;
-  });
-  return { promise, resolve, reject };
-}
-
-async function collectEvents(
-  observable:
-    | ReturnType<InMemoryAgentRunner["run"]>
-    | ReturnType<InMemoryAgentRunner["connect"]>,
-) {
-  return firstValueFrom(observable.pipe(toArray()));
-}
-
-function createRunInput({
-  threadId,
-  runId,
-  messages,
-  state,
-  parentRunId,
-}: {
-  threadId: string;
-  runId: string;
-  messages: Message[];
-  state?: Record<string, unknown>;
-  parentRunId?: string | null;
-}): RunAgentInput {
-  return {
-    threadId,
-    runId,
-    parentRunId: parentRunId ?? undefined,
-    state: state ?? {},
-    messages,
-    tools: [],
-    context: [],
-    forwardedProps: undefined,
-  };
-}
-
-function expectRunStartedEvent(event: BaseEvent, expectedMessages: Message[]) {
-  expect(event.type).toBe(EventType.RUN_STARTED);
-  const runStarted = event as RunStartedEvent;
-  expect(runStarted.input?.messages).toEqual(expectedMessages);
-}
-
-function createTextMessageEvents({
-  messageId,
-  role = "assistant",
-  content,
-}: {
-  messageId: string;
-  role?: "assistant" | "developer" | "system" | "user";
-  content: string;
-}): BaseEvent[] {
-  return [
-    {
-      type: EventType.TEXT_MESSAGE_START,
-      messageId,
-      role,
-    },
-    {
-      type: EventType.TEXT_MESSAGE_CONTENT,
-      messageId,
-      delta: content,
-    },
-    {
-      type: EventType.TEXT_MESSAGE_END,
-      messageId,
-    },
-  ] as BaseEvent[];
-}
-
-function createToolCallEvents({
-  toolCallId,
-  parentMessageId,
-  toolName,
-  argsJson,
-  resultMessageId,
-  resultContent,
-}: {
-  toolCallId: string;
-  parentMessageId: string;
-  toolName: string;
-  argsJson: string;
-  resultMessageId: string;
-  resultContent: string;
-}): BaseEvent[] {
-  return [
-    {
-      type: EventType.TOOL_CALL_START,
-      toolCallId,
-      toolCallName: toolName,
-      parentMessageId,
-    },
-    {
-      type: EventType.TOOL_CALL_ARGS,
-      toolCallId,
-      delta: argsJson,
-    },
-    {
-      type: EventType.TOOL_CALL_END,
-      toolCallId,
-    },
-    {
-      type: EventType.TOOL_CALL_RESULT,
-      toolCallId,
-      messageId: resultMessageId,
-      content: resultContent,
-      role: "tool",
-    },
-  ] as BaseEvent[];
-}
-
-describe("InMemoryAgentRunner e2e", () => {
-  describe("Fresh Replay After Single Run", () => {
-    it("replays sanitized message history on connectAgent", async () => {
-      const runner = new InMemoryAgentRunner();
-      const threadId = "thread-fresh-replay";
-      const existingMessage: Message = {
-        id: "message-existing",
-        role: "user",
-        content: "Hello there",
-      };
-
-      const runEvents = await collectEvents(
-        runner.run({
-          threadId,
-          agent: new EmitAgent(),
-          input: createRunInput({
-            threadId,
-            runId: "run-0",
-            messages: [existingMessage],
-          }),
-        }),
-      );
-
-      expectRunStartedEvent(runEvents[0], [existingMessage]);
-      expect(runEvents.at(-1)?.type).toBe(EventType.RUN_FINISHED);
-
-      const replayEvents = await collectEvents(runner.connect({ threadId }));
-      const replayAgent = new ReplayAgent(replayEvents, threadId);
-      await replayAgent.connectAgent({ runId: "replay-run" });
-
-      expect(replayAgent.messages).toEqual([existingMessage]);
-    });
-  });
-
-  describe("New Messages on Subsequent Runs", () => {
-    it("merges new message IDs without duplicating history", async () => {
-      const runner = new InMemoryAgentRunner();
-      const threadId = "thread-subsequent-runs";
-      const existingMessage: Message = {
-        id: "msg-existing",
-        role: "user",
-        content: "First turn",
-      };
-
-      const initialRunEvents = await collectEvents(
-        runner.run({
-          threadId,
-          agent: new EmitAgent(),
-          input: createRunInput({
-            threadId,
-            runId: "run-0",
-            messages: [existingMessage],
-          }),
-        }),
-      );
-      expectRunStartedEvent(initialRunEvents[0], [existingMessage]);
-
-      const newMessage: Message = {
-        id: "msg-new",
-        role: "user",
-        content: "Second turn",
-      };
-
-      const secondRunEvents = await collectEvents(
-        runner.run({
-          threadId,
-          agent: new EmitAgent(),
-          input: createRunInput({
-            threadId,
-            runId: "run-1",
-            messages: [existingMessage, newMessage],
-          }),
-        }),
-      );
-
-      expectRunStartedEvent(secondRunEvents[0], [newMessage]);
-
-      const replayEvents = await collectEvents(runner.connect({ threadId }));
-      const replayAgent = new ReplayAgent(replayEvents, threadId);
-      await replayAgent.connectAgent({ runId: "replay-run" });
-
-      expect(replayAgent.messages).toEqual([existingMessage, newMessage]);
-      expect(
-        new Set(replayAgent.messages.map((message) => message.id)).size,
-      ).toBe(replayAgent.messages.length);
-    });
-  });
-
-  describe("Fresh Agent Connection After Prior Runs", () => {
-    it("hydrates a brand-new agent via connect()", async () => {
-      const runner = new InMemoryAgentRunner();
-      const threadId = "thread-new-agent-connection";
-      const existingMessage: Message = {
-        id: "existing-connection",
-        role: "user",
-        content: "Persist me",
-      };
-
-      const runEvents = await collectEvents(
-        runner.run({
-          threadId,
-          agent: new EmitAgent(),
-          input: createRunInput({
-            threadId,
-            runId: "run-0",
-            messages: [existingMessage],
-          }),
-        }),
-      );
-      expectRunStartedEvent(runEvents[0], [existingMessage]);
-
-      const connectingAgent = new RunnerConnectAgent(runner, threadId);
-      await connectingAgent.connectAgent({ runId: "connect-run" });
-
-      expect(connectingAgent.messages).toEqual([existingMessage]);
-    });
-  });
-
-  describe("Mixed Roles and Tool Results", () => {
-    it("preserves agent-emitted tool events alongside heterogeneous inputs", async () => {
-      const runner = new InMemoryAgentRunner();
-      const threadId = "thread-mixed-roles";
-
-      const systemMessage: Message = {
-        id: "sys-1",
-        role: "system",
-        content: "Global directive",
-      };
-      const developerMessage: Message = {
-        id: "dev-1",
-        role: "developer",
-        content: "Internal guidance",
-      };
-      const userMessage: Message = {
-        id: "user-1",
-        role: "user",
-        content: "Need the weather",
-      };
-      const baseMessages = [systemMessage, developerMessage, userMessage];
-
-      const assistantMessageId = "assistant-1";
-      const toolCallId = "tool-call-1";
-      const toolMessageId = "tool-msg-1";
-
-      const agentEvents: BaseEvent[] = [
-        ...createTextMessageEvents({
-          messageId: assistantMessageId,
-          content: "Calling the weather tool",
-        }),
-        ...createToolCallEvents({
-          toolCallId,
-          parentMessageId: assistantMessageId,
-          toolName: "getWeather",
-          argsJson: '{"location":"NYC"}',
-          resultMessageId: toolMessageId,
-          resultContent: '{"temp":72}',
-        }),
-      ];
-
-      const runEvents = await collectEvents(
-        runner.run({
-          threadId,
-          agent: new EmitAgent({ events: agentEvents }),
-          input: createRunInput({
-            threadId,
-            runId: "run-0",
-            messages: baseMessages,
-          }),
-        }),
-      );
-
-      expectRunStartedEvent(runEvents[0], baseMessages);
-      expect(
-        runEvents.filter((event) => event.type === EventType.TOOL_CALL_RESULT),
-      ).toHaveLength(1);
-
-      const replayEvents = await collectEvents(runner.connect({ threadId }));
-      const replayAgent = new ReplayAgent(replayEvents, threadId);
-      await replayAgent.connectAgent({ runId: "replay-run" });
-
-      expect(replayAgent.messages).toEqual([
-        systemMessage,
-        developerMessage,
-        userMessage,
-        {
-          id: assistantMessageId,
-          role: "assistant",
-          content: "Calling the weather tool",
-          toolCalls: [
-            {
-              id: toolCallId,
-              type: "function",
-              function: {
-                name: "getWeather",
-                arguments: '{"location":"NYC"}',
-              },
-            },
-          ],
-        },
-        {
-          id: toolMessageId,
-          role: "tool",
-          content: '{"temp":72}',
-          toolCallId,
-        },
-      ]);
-      expect(
-        replayAgent.messages.filter((message) => message.role === "tool"),
-      ).toHaveLength(1);
-    });
-  });
-
-  describe("Multiple Consecutive Runs with Agent Output", () => {
-    it("deduplicates input history while emitting each agent message once", async () => {
-      const runner = new InMemoryAgentRunner();
-      const threadId = "thread-multi-runs";
-      const systemMessage: Message = {
-        id: "system-shared",
-        role: "system",
-        content: "System context",
-      };
-      const userMessages: Message[] = [];
-
-      for (let index = 0; index < 3; index += 1) {
-        const userMessage: Message = {
-          id: `user-${index + 1}`,
-          role: "user",
-          content: `User message ${index + 1}`,
-        };
-        userMessages.push(userMessage);
-
-        const messagesForRun = [systemMessage, ...userMessages];
-        const assistantId = `assistant-${index + 1}`;
-        const toolCallId = `tool-call-${index + 1}`;
-        const toolMessageId = `tool-msg-${index + 1}`;
-
-        const events: BaseEvent[] = [
-          ...createTextMessageEvents({
-            messageId: assistantId,
-            content: `Assistant reply ${index + 1}`,
-          }),
-          ...createToolCallEvents({
-            toolCallId,
-            parentMessageId: assistantId,
-            toolName: `tool-${index + 1}`,
-            argsJson: `{"step":${index + 1}}`,
-            resultMessageId: toolMessageId,
-            resultContent: `{"ok":${index + 1}}`,
-          }),
-        ];
-
-        const runEvents = await collectEvents(
-          runner.run({
-            threadId,
-            agent: new EmitAgent({ events }),
-            input: createRunInput({
-              threadId,
-              runId: `run-${index}`,
-              messages: messagesForRun,
-            }),
-          }),
-        );
-
-        if (index === 0) {
-          expectRunStartedEvent(runEvents[0], messagesForRun);
-        } else {
-          expectRunStartedEvent(runEvents[0], [userMessage]);
-        }
-        expect(runEvents.at(-1)?.type).toBe(EventType.RUN_FINISHED);
-      }
-
-      const replayEvents = await collectEvents(runner.connect({ threadId }));
-      const replayAgent = new ReplayAgent(replayEvents, threadId);
-      await replayAgent.connectAgent({ runId: "replay-final" });
-
-      const finalMessages = replayAgent.messages;
-      expect(new Set(finalMessages.map((message) => message.id)).size).toBe(
-        finalMessages.length,
-      );
-      const roleCounts = finalMessages.reduce<Record<string, number>>(
-        (counts, message) => {
-          counts[message.role] = (counts[message.role] ?? 0) + 1;
-          return counts;
-        },
-        {},
-      );
-      expect(roleCounts.system).toBe(1);
-      expect(roleCounts.user).toBe(3);
-      expect(roleCounts.assistant).toBe(3);
-      expect(roleCounts.tool).toBe(3);
-    });
-  });
-
-  describe("Agent-Provided RUN_STARTED input", () => {
-    it("forwards the agent-specified payload without sanitizing", async () => {
-      const runner = new InMemoryAgentRunner();
-      const threadId = "thread-custom-run-started";
-      const runId = "run-0";
-
-      const customMessages: Message[] = [
-        {
-          id: "custom-user",
-          role: "user",
-          content: "Pre-sent content",
-        },
-      ];
-      const customInput: RunAgentInput = {
-        threadId,
-        runId,
-        parentRunId: undefined,
-        state: { injected: true },
-        messages: customMessages,
-        tools: [],
-        context: [],
-        forwardedProps: { source: "agent" },
-      };
-      const customRunStarted: RunStartedEvent = {
-        type: EventType.RUN_STARTED,
-        threadId,
-        runId,
-        parentRunId: null,
-        input: customInput,
-      };
-
-      const agentEvents: BaseEvent[] = [
-        customRunStarted,
-        ...createTextMessageEvents({
-          messageId: "agent-message",
-          content: "Custom start acknowledged",
-        }),
-      ];
-
-      const runEvents = await collectEvents(
-        runner.run({
-          threadId,
-          agent: new EmitAgent({
-            events: agentEvents,
-            emitDefaultRunStarted: false,
-          }),
-          input: createRunInput({
-            threadId,
-            runId,
-            messages: [],
-          }),
-        }),
-      );
-
-      expect(runEvents[0]).toEqual(customRunStarted);
-      expect(
-        runEvents.filter((event) => event.type === EventType.RUN_FINISHED),
-      ).toHaveLength(1);
-
-      const replayEvents = await collectEvents(runner.connect({ threadId }));
-      const replayAgent = new ReplayAgent(replayEvents, threadId);
-      await replayAgent.connectAgent({ runId: "replay-run" });
-      expect(
-        replayAgent.messages.find((message) => message.id === "custom-user"),
-      ).toEqual(customMessages[0]);
-    });
-  });
-
-  describe("Concurrent Connections During Run", () => {
-    it("streams in-flight events to live subscribers and persists final history", async () => {
-      const runner = new InMemoryAgentRunner();
-      const threadId = "thread-concurrency";
-      const runId = "run-live";
-      const initialMessage: Message = {
-        id: "initial-user",
-        role: "user",
-        content: "Start run",
-      };
-
-      const runStartedSignal = createDeferred<void>();
-      const resumeSignal = createDeferred<void>();
-
-      const agent = new EmitAgent({
-        events: [
-          ...createTextMessageEvents({
-            messageId: "assistant-live",
-            content: "Streaming content",
-          }),
-        ],
-        afterEvent: async ({ event }) => {
-          if (event.type === EventType.RUN_STARTED) {
-            runStartedSignal.resolve();
-            await resumeSignal.promise;
-          }
-        },
-      });
-
-      const runEvents: BaseEvent[] = [];
-      const run$ = runner.run({
-        threadId,
-        agent,
-        input: createRunInput({
-          threadId,
-          runId,
-          messages: [initialMessage],
-        }),
-      });
-
-      let runSubscription: Subscription;
-      const runCompletion = new Promise<void>((resolve, reject) => {
-        runSubscription = run$.subscribe({
-          next: (event) => runEvents.push(event),
-          error: (error) => {
-            runSubscription.unsubscribe();
-            reject(error);
-          },
-          complete: () => {
-            runSubscription.unsubscribe();
-            resolve();
-          },
-        });
-      });
-
-      await runStartedSignal.promise;
-
-      const liveEvents: BaseEvent[] = [];
-      const connect$ = runner.connect({ threadId });
-      let connectSubscription: Subscription;
-      const connectCompletion = new Promise<void>((resolve, reject) => {
-        connectSubscription = connect$.subscribe({
-          next: (event) => liveEvents.push(event),
-          error: (error) => {
-            connectSubscription.unsubscribe();
-            reject(error);
-          },
-          complete: () => {
-            connectSubscription.unsubscribe();
-            resolve();
-          },
-        });
-      });
-
-      resumeSignal.resolve();
-
-      await Promise.all([runCompletion, connectCompletion]);
-
-      expectRunStartedEvent(runEvents[0], [initialMessage]);
-      expect(runEvents.at(-1)?.type).toBe(EventType.RUN_FINISHED);
-      expect(liveEvents).toEqual(runEvents);
-
-      const persistedEvents = await collectEvents(runner.connect({ threadId }));
-      expect(persistedEvents).toEqual(runEvents);
-
-      const replayAgent = new ReplayAgent(persistedEvents, threadId);
-      await replayAgent.connectAgent({ runId: "replay-run" });
-      expect(replayAgent.messages.map((message) => message.id)).toEqual([
-        initialMessage.id,
-        "assistant-live",
-      ]);
-    });
-  });
-
-  describe("Error Handling", () => {
-    it("propagates RUN_ERROR while retaining input history", async () => {
-      const runner = new InMemoryAgentRunner();
-      const threadId = "thread-run-error";
-      const userMessage: Message = {
-        id: "error-user",
-        role: "user",
-        content: "Trigger error",
-      };
-
-      const runErrorEvent: RunErrorEvent = {
-        type: EventType.RUN_ERROR,
-        message: "Agent failure",
-      };
-
-      const runEvents = await collectEvents(
-        runner.run({
-          threadId,
-          agent: new EmitAgent({
-            events: [runErrorEvent],
-            includeRunFinished: false,
-          }),
-          input: createRunInput({
-            threadId,
-            runId: "run-error",
-            messages: [userMessage],
-          }),
-        }),
-      );
-
-      expectRunStartedEvent(runEvents[0], [userMessage]);
-      expect(runEvents.at(-1)).toEqual(runErrorEvent);
-
-      const replayEvents = await collectEvents(runner.connect({ threadId }));
-      const replayAgent = new ReplayAgent(replayEvents, threadId);
-      const capturedRunErrors: RunErrorEvent[] = [];
-      const result = await replayAgent.connectAgent(
-        { runId: "replay-run" },
-        {
-          onRunErrorEvent: ({ event }) => {
-            capturedRunErrors.push(event);
-          },
-        },
-      );
-
-      expect(runEvents.at(-1)?.type).toBe(EventType.RUN_ERROR);
-      expect(capturedRunErrors).toHaveLength(1);
-      expect(capturedRunErrors[0]).toMatchObject(runErrorEvent);
-      expect(result.newMessages).toEqual([userMessage]);
-      expect(replayAgent.messages).toEqual([userMessage]);
-    });
-  });
-});