@copilotkit/runtime 1.55.1 → 1.55.2-next.1

package/src/agent/__tests__/agent.test.ts

@@ -0,0 +1,593 @@
+import { describe, it, expect } from "vitest";
+import { EventType, type BaseEvent } from "@ag-ui/client";
+import {
+  BuiltInAgent,
+  type AgentFactoryContext,
+  type BuiltInAgentFactoryConfig,
+  createDefaultInput,
+  createAgent,
+  createThrowingAgent,
+  createMidStreamErrorAgent,
+  collectEvents,
+  collectEventsIncludingErrors,
+  expectLifecycleWrapped,
+  eventField,
+  textDelta,
+  finish,
+  tanstackTextChunk,
+  type AgentType,
+  type MockStreamEvent,
+} from "./agent-test-helpers";
+
+// ---------------------------------------------------------------------------
+// Local helpers for parameterized tests
+// ---------------------------------------------------------------------------
+
+const allTypes: AgentType[] = ["aisdk", "tanstack", "custom"];
+
+function minimalStreamData(
+  type: AgentType,
+): MockStreamEvent[] | Record<string, unknown>[] | BaseEvent[] {
+  switch (type) {
+    case "aisdk":
+      return [textDelta("hi"), finish()];
+    case "tanstack":
+      return [tanstackTextChunk("hi")];
+    case "custom":
+      return [
+        {
+          type: EventType.TEXT_MESSAGE_CHUNK,
+          role: "assistant",
+          delta: "hi",
+        } as BaseEvent,
+      ];
+  }
+}
+
+function emptyStreamData(
+  type: AgentType,
+): MockStreamEvent[] | Record<string, unknown>[] | BaseEvent[] {
+  switch (type) {
+    case "aisdk":
+      return [finish()];
+    case "tanstack":
+      return [];
+    case "custom":
+      return [];
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Parameterized test suites
+// ---------------------------------------------------------------------------
+
+describe.each(allTypes)("Agent [%s]", (type) => {
+  // -------------------------------------------------------------------------
+  // Lifecycle
+  // -------------------------------------------------------------------------
+  describe("lifecycle", () => {
+    it("emits RUN_STARTED as the first event with correct threadId/runId", async () => {
+      const agent = createAgent(type, minimalStreamData(type));
+      const input = createDefaultInput({ threadId: "t1", runId: "r1" });
+      const events = await collectEvents(agent.run(input));
+
+      expect(events.length).toBeGreaterThanOrEqual(2);
+      const first = events[0];
+      expect(first.type).toBe(EventType.RUN_STARTED);
+      expect(eventField<string>(first, "threadId")).toBe("t1");
+      expect(eventField<string>(first, "runId")).toBe("r1");
+    });
+
+    it("emits RUN_FINISHED as the last event with correct threadId/runId", async () => {
+      const agent = createAgent(type, minimalStreamData(type));
+      const input = createDefaultInput({ threadId: "t2", runId: "r2" });
+      const events = await collectEvents(agent.run(input));
+
+      const last = events[events.length - 1];
+      expect(last.type).toBe(EventType.RUN_FINISHED);
+      expect(eventField<string>(last, "threadId")).toBe("t2");
+      expect(eventField<string>(last, "runId")).toBe("r2");
+    });
+
+    it("emits RUN_FINISHED for an empty stream", async () => {
+      const agent = createAgent(type, emptyStreamData(type));
+      const input = createDefaultInput();
+      const events = await collectEvents(agent.run(input));
+
+      expect(events.length).toBeGreaterThanOrEqual(2);
+      expect(events[0].type).toBe(EventType.RUN_STARTED);
+      expect(events[events.length - 1].type).toBe(EventType.RUN_FINISHED);
+    });
+
+    it("wraps content with lifecycle events", async () => {
+      const agent = createAgent(type, minimalStreamData(type));
+      const input = createDefaultInput({ threadId: "wrap-t", runId: "wrap-r" });
+      const events = await collectEvents(agent.run(input));
+
+      expectLifecycleWrapped(events, "wrap-t", "wrap-r");
+
+      // There should be content events between the lifecycle bookends
+      const contentEvents = events.slice(1, -1);
+      expect(contentEvents.length).toBeGreaterThan(0);
+      for (const e of contentEvents) {
+        expect(e.type).not.toBe(EventType.RUN_STARTED);
+        expect(e.type).not.toBe(EventType.RUN_FINISHED);
+      }
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // RUN_ERROR
+  // -------------------------------------------------------------------------
+  describe("RUN_ERROR", () => {
+    it("emits RUN_ERROR when factory throws", async () => {
+      const agent = createThrowingAgent(type, "factory-boom");
+      const input = createDefaultInput();
+      const { events, errored } = await collectEventsIncludingErrors(
+        agent.run(input),
+      );
+
+      expect(errored).toBe(true);
+      const errorEvents = events.filter((e) => e.type === EventType.RUN_ERROR);
+      expect(errorEvents.length).toBe(1);
+      expect(eventField<string>(errorEvents[0], "message")).toBe(
+        "factory-boom",
+      );
+    });
+
+    it("emits RUN_ERROR when stream throws mid-iteration", async () => {
+      const agent = createMidStreamErrorAgent(type, "mid-stream-boom");
+      const input = createDefaultInput();
+      const { events, errored } = await collectEventsIncludingErrors(
+        agent.run(input),
+      );
+
+      expect(errored).toBe(true);
+      const errorEvents = events.filter((e) => e.type === EventType.RUN_ERROR);
+      expect(errorEvents.length).toBe(1);
+      expect(eventField<string>(errorEvents[0], "message")).toBe(
+        "mid-stream-boom",
+      );
+    });
+
+    it("does not emit RUN_FINISHED after RUN_ERROR", async () => {
+      const agent = createThrowingAgent(type, "no-finish");
+      const input = createDefaultInput();
+      const { events } = await collectEventsIncludingErrors(agent.run(input));
+
+      const errorIdx = events.findIndex((e) => e.type === EventType.RUN_ERROR);
+      expect(errorIdx).toBeGreaterThanOrEqual(0);
+
+      const eventsAfterError = events.slice(errorIdx + 1);
+      const finishAfterError = eventsAfterError.filter(
+        (e) => e.type === EventType.RUN_FINISHED,
+      );
+      expect(finishAfterError.length).toBe(0);
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // Abort
+  // -------------------------------------------------------------------------
+  describe("abort", () => {
+    it("completes without error after abortRun()", async () => {
+      // Use a signal to synchronize: abort after the first chunk is emitted
+      let emittedFirstChunk: () => void;
+      const firstChunkEmitted = new Promise<void>(
+        (r) => (emittedFirstChunk = r),
+      );
+
+      let config: BuiltInAgentFactoryConfig;
+      switch (type) {
+        case "aisdk":
+          config = {
+            type: "aisdk",
+            factory: ({ abortSignal }: AgentFactoryContext) => ({
+              fullStream: (async function* () {
+                yield { type: "text-delta", text: "tick" };
+                emittedFirstChunk();
+                // Wait for abort — use a promise that resolves on abort
+                await new Promise<void>((r) => {
+                  if (abortSignal.aborted) return r();
+                  abortSignal.addEventListener("abort", () => r(), {
+                    once: true,
+                  });
+                });
+              })(),
+            }),
+          };
+          break;
+        case "tanstack":
+          config = {
+            type: "tanstack",
+            factory: ({ abortSignal }: AgentFactoryContext) => ({
+              [Symbol.asyncIterator]: async function* () {
+                yield { type: "TEXT_MESSAGE_CONTENT", delta: "tick" };
+                emittedFirstChunk();
+                await new Promise<void>((r) => {
+                  if (abortSignal.aborted) return r();
+                  abortSignal.addEventListener("abort", () => r(), {
+                    once: true,
+                  });
+                });
+              },
+            }),
+          };
+          break;
+        case "custom":
+          config = {
+            type: "custom",
+            factory: ({ abortSignal }: AgentFactoryContext) => ({
+              [Symbol.asyncIterator]: async function* () {
+                yield {
+                  type: EventType.TEXT_MESSAGE_CHUNK,
+                  role: "assistant",
+                  delta: "tick",
+                } as BaseEvent;
+                emittedFirstChunk();
+                await new Promise<void>((r) => {
+                  if (abortSignal.aborted) return r();
+                  abortSignal.addEventListener("abort", () => r(), {
+                    once: true,
+                  });
+                });
+              },
+            }),
+          };
+          break;
+      }
+
+      const agent = new BuiltInAgent(config);
+      const input = createDefaultInput();
+
+      const completed = await new Promise<boolean>((resolve) => {
+        agent.run(input).subscribe({
+          next: () => {},
+          error: () => resolve(false),
+          complete: () => resolve(true),
+        });
+
+        // Wait for the first chunk to be emitted, then abort
+        firstChunkEmitted.then(() => agent.abortRun());
+      });
+
+      expect(completed).toBe(true);
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // Factory Context
+  // -------------------------------------------------------------------------
+  describe("factory context", () => {
+    it("receives correct input with threadId, runId, and forwardedProps", async () => {
+      let capturedCtx: AgentFactoryContext | null = null;
+
+      let config: BuiltInAgentFactoryConfig;
+      switch (type) {
+        case "aisdk":
+          config = {
+            type: "aisdk",
+            factory: (ctx: AgentFactoryContext) => {
+              capturedCtx = ctx;
+              return {
+                fullStream: (async function* () {
+                  yield { type: "finish", finishReason: "stop" };
+                })(),
+              };
+            },
+          };
+          break;
+        case "tanstack":
+          config = {
+            type: "tanstack",
+            factory: (ctx: AgentFactoryContext) => {
+              capturedCtx = ctx;
+              return (async function* () {
+                // empty stream
+              })();
+            },
+          };
+          break;
+        case "custom":
+          config = {
+            type: "custom",
+            factory: (ctx: AgentFactoryContext) => {
+              capturedCtx = ctx;
+              return (async function* () {
+                // empty stream
+              })();
+            },
+          };
+          break;
+      }
+
+      const agent = new BuiltInAgent(config);
+      const input = createDefaultInput({
+        threadId: "ctx-thread",
+        runId: "ctx-run",
+        forwardedProps: { model: "gpt-4" },
+      });
+
+      await collectEvents(agent.run(input));
+
+      expect(capturedCtx).not.toBeNull();
+      expect(capturedCtx!.input.threadId).toBe("ctx-thread");
+      expect(capturedCtx!.input.runId).toBe("ctx-run");
+      expect(capturedCtx!.input.forwardedProps).toEqual({ model: "gpt-4" });
+    });
+
+    it("receives abortController and abortSignal", async () => {
+      let capturedCtx: AgentFactoryContext | null = null;
+
+      let config: BuiltInAgentFactoryConfig;
+      switch (type) {
+        case "aisdk":
+          config = {
+            type: "aisdk",
+            factory: (ctx: AgentFactoryContext) => {
+              capturedCtx = ctx;
+              return {
+                fullStream: (async function* () {
+                  yield { type: "finish", finishReason: "stop" };
+                })(),
+              };
+            },
+          };
+          break;
+        case "tanstack":
+          config = {
+            type: "tanstack",
+            factory: (ctx: AgentFactoryContext) => {
+              capturedCtx = ctx;
+              return (async function* () {
+                // empty
+              })();
+            },
+          };
+          break;
+        case "custom":
+          config = {
+            type: "custom",
+            factory: (ctx: AgentFactoryContext) => {
+              capturedCtx = ctx;
+              return (async function* () {
+                // empty
+              })();
+            },
+          };
+          break;
+      }
+
+      const agent = new BuiltInAgent(config);
+      const input = createDefaultInput();
+      await collectEvents(agent.run(input));
+
+      expect(capturedCtx!.abortController).toBeInstanceOf(AbortController);
+      expect(capturedCtx!.abortSignal).toBe(
+        capturedCtx!.abortController.signal,
+      );
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // clone()
+  // -------------------------------------------------------------------------
+  describe("clone()", () => {
+    it("returns a new Agent instance (not the same reference)", () => {
+      const agent = createAgent(type, minimalStreamData(type));
+      const cloned = agent.clone();
+
+      expect(cloned).toBeInstanceOf(BuiltInAgent);
+      expect(cloned).not.toBe(agent);
+    });
+
+    it("produces correct lifecycle events from a cloned agent", async () => {
+      const agent = createAgent(type, minimalStreamData(type));
+      const cloned = agent.clone();
+      const input = createDefaultInput({
+        threadId: "clone-t",
+        runId: "clone-r",
+      });
+
+      const events = await collectEvents(cloned.run(input));
+
+      expectLifecycleWrapped(events, "clone-t", "clone-r");
+    });
+  });
+});
+
+// ---------------------------------------------------------------------------
+// Type Discrimination (NOT parameterized)
+// ---------------------------------------------------------------------------
+
+describe("Agent type discrimination", () => {
+  it('"aisdk" routes to AI SDK converter and produces text content', async () => {
+    const agent = createAgent("aisdk", [
+      textDelta("hello from aisdk"),
+      finish(),
+    ]);
+    const input = createDefaultInput();
+    const events = await collectEvents(agent.run(input));
+
+    const textEvents = events.filter(
+      (e) => e.type === EventType.TEXT_MESSAGE_CHUNK,
+    );
+    expect(textEvents.length).toBe(1);
+    expect(eventField<string>(textEvents[0], "delta")).toBe("hello from aisdk");
+  });
+
+  it('"tanstack" routes to TanStack converter and produces text content', async () => {
+    const agent = createAgent("tanstack", [
+      tanstackTextChunk("hello from tanstack"),
+    ]);
+    const input = createDefaultInput();
+    const events = await collectEvents(agent.run(input));
+
+    const textEvents = events.filter(
+      (e) => e.type === EventType.TEXT_MESSAGE_CHUNK,
+    );
+    expect(textEvents.length).toBe(1);
+    expect(eventField<string>(textEvents[0], "delta")).toBe(
+      "hello from tanstack",
+    );
+  });
+
+  it('"custom" forwards events directly without conversion', async () => {
+    const customEvent: BaseEvent = {
+      type: EventType.TEXT_MESSAGE_CHUNK,
+      role: "assistant",
+      delta: "hello from custom",
+    } as BaseEvent;
+
+    const agent = createAgent("custom", [customEvent]);
+    const input = createDefaultInput();
+    const events = await collectEvents(agent.run(input));
+
+    const textEvents = events.filter(
+      (e) => e.type === EventType.TEXT_MESSAGE_CHUNK,
+    );
+    expect(textEvents.length).toBe(1);
+    expect(eventField<string>(textEvents[0], "delta")).toBe(
+      "hello from custom",
+    );
+  });
+});
+
+// ---------------------------------------------------------------------------
+// Async Factory (Promise-returning)
+// ---------------------------------------------------------------------------
+
+describe("Async factory (Promise-returning)", () => {
+  it("aisdk: async factory resolves and streams correctly", async () => {
+    const agent = new BuiltInAgent({
+      type: "aisdk",
+      factory: async () => {
+        // Simulate async setup (e.g., fetching API key)
+        await new Promise((r) => setTimeout(r, 5));
+        return {
+          fullStream: (async function* () {
+            yield { type: "text-delta", text: "async-aisdk" };
+            yield { type: "finish", finishReason: "stop" };
+          })(),
+        };
+      },
+    });
+    const input = createDefaultInput();
+    const events = await collectEvents(agent.run(input));
+
+    expectLifecycleWrapped(events);
+    const textEvents = events.filter(
+      (e) => e.type === EventType.TEXT_MESSAGE_CHUNK,
+    );
+    expect(textEvents).toHaveLength(1);
+    expect(eventField<string>(textEvents[0], "delta")).toBe("async-aisdk");
+  });
+
+  it("tanstack: async factory resolves and streams correctly", async () => {
+    const agent = new BuiltInAgent({
+      type: "tanstack",
+      factory: async () => {
+        await new Promise((r) => setTimeout(r, 5));
+        return (async function* () {
+          yield { type: "TEXT_MESSAGE_CONTENT", delta: "async-tanstack" };
+        })();
+      },
+    });
+    const input = createDefaultInput();
+    const events = await collectEvents(agent.run(input));
+
+    expectLifecycleWrapped(events);
+    const textEvents = events.filter(
+      (e) => e.type === EventType.TEXT_MESSAGE_CHUNK,
+    );
+    expect(textEvents).toHaveLength(1);
+    expect(eventField<string>(textEvents[0], "delta")).toBe("async-tanstack");
+  });
+
+  it("custom: async factory resolves and streams correctly", async () => {
+    const agent = new BuiltInAgent({
+      type: "custom",
+      factory: async () => {
+        await new Promise((r) => setTimeout(r, 5));
+        return (async function* () {
+          yield {
+            type: EventType.TEXT_MESSAGE_CHUNK,
+            role: "assistant",
+            delta: "async-custom",
+          } as BaseEvent;
+        })();
+      },
+    });
+    const input = createDefaultInput();
+    const events = await collectEvents(agent.run(input));
+
+    expectLifecycleWrapped(events);
+    const textEvents = events.filter(
+      (e) => e.type === EventType.TEXT_MESSAGE_CHUNK,
+    );
+    expect(textEvents).toHaveLength(1);
+    expect(eventField<string>(textEvents[0], "delta")).toBe("async-custom");
+  });
+});
+
+// ---------------------------------------------------------------------------
+// RUN_ERROR includes threadId and runId
+// ---------------------------------------------------------------------------
+
+describe("RUN_ERROR correlation fields", () => {
+  it("RUN_ERROR includes threadId and runId for run correlation", async () => {
+    const agent = new BuiltInAgent({
+      type: "aisdk",
+      factory: () => {
+        throw new Error("test-error");
+      },
+    });
+    const input = createDefaultInput({
+      threadId: "err-thread",
+      runId: "err-run",
+    });
+    const { events, errored } = await collectEventsIncludingErrors(
+      agent.run(input),
+    );
+
+    expect(errored).toBe(true);
+    const errorEvents = events.filter((e) => e.type === EventType.RUN_ERROR);
+    expect(errorEvents).toHaveLength(1);
+    expect(eventField<string>(errorEvents[0], "threadId")).toBe("err-thread");
+    expect(eventField<string>(errorEvents[0], "runId")).toBe("err-run");
+  });
+});
+
+// ---------------------------------------------------------------------------
+// Concurrent run guard
+// ---------------------------------------------------------------------------
+
+describe("Concurrent run guard", () => {
+  it("throws when run() is called while another run is in progress", async () => {
+    let resolveFactory: () => void;
+    const factoryBlocked = new Promise<void>((r) => (resolveFactory = r));
+
+    const agent = new BuiltInAgent({
+      type: "custom",
+      factory: async function* ({ abortSignal }) {
+        // Block until resolved externally
+        await new Promise<void>((r) => {
+          if (abortSignal.aborted) return r();
+          abortSignal.addEventListener("abort", () => r(), { once: true });
+          factoryBlocked.then(() => r());
+        });
+      },
+    });
+    const input = createDefaultInput();
+
+    // Start first run — abortController is now set synchronously in run()
+    const sub = agent.run(input).subscribe({ next: () => {} });
+
+    // Second run should throw immediately (no timing dependency)
+    expect(() => agent.run(input)).toThrow("Agent is already running");
+
+    // Cleanup
+    resolveFactory!();
+    sub.unsubscribe();
+  });
+});