@copilotkit/runtime 1.56.2 → 1.56.4-canary.1777529757

package/src/agent/__tests__/converter-tanstack.test.ts

@@ -11,6 +11,12 @@ import {
   tanstackToolCallStart,
   tanstackToolCallArgs,
   tanstackToolCallEnd,
+  tanstackToolCallResult,
+  tanstackReasoningStart,
+  tanstackReasoningMessageStart,
+  tanstackReasoningMessageContent,
+  tanstackReasoningMessageEnd,
+  tanstackReasoningEnd,
 } from "./agent-test-helpers";
 
 describe("TanStack AI converter (via Agent)", () => {
@@ -312,3 +318,277 @@ describe("TanStack AI converter (via Agent)", () => {
     });
   });
 });
+
+describe("TanStack AI converter — state tools", () => {
+  it("emits STATE_SNAPSHOT before TOOL_CALL_RESULT for AGUISendStateSnapshot", async () => {
+    const snapshot = { counter: 5, items: ["x", "y"] };
+    const agent = createAgent("tanstack", [
+      tanstackToolCallStart("call1", "AGUISendStateSnapshot"),
+      tanstackToolCallEnd("call1"),
+      tanstackToolCallResult("call1", { success: true, snapshot }),
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+
+    expectLifecycleWrapped(events);
+
+    const snapshotIdx = events.findIndex(
+      (e) => e.type === EventType.STATE_SNAPSHOT,
+    );
+    const resultIdx = events.findIndex(
+      (e) => e.type === EventType.TOOL_CALL_RESULT,
+    );
+
+    expect(snapshotIdx).toBeGreaterThanOrEqual(0);
+    expect(resultIdx).toBeGreaterThanOrEqual(0);
+    expect(snapshotIdx).toBeLessThan(resultIdx);
+    expect(eventField<unknown>(events[snapshotIdx], "snapshot")).toEqual(
+      snapshot,
+    );
+  });
+
+  it("emits STATE_DELTA before TOOL_CALL_RESULT for AGUISendStateDelta", async () => {
+    const delta = [{ op: "replace", path: "/counter", value: 7 }];
+    const agent = createAgent("tanstack", [
+      tanstackToolCallStart("call1", "AGUISendStateDelta"),
+      tanstackToolCallEnd("call1"),
+      tanstackToolCallResult("call1", { success: true, delta }),
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+
+    expectLifecycleWrapped(events);
+
+    const deltaIdx = events.findIndex((e) => e.type === EventType.STATE_DELTA);
+    const resultIdx = events.findIndex(
+      (e) => e.type === EventType.TOOL_CALL_RESULT,
+    );
+
+    expect(deltaIdx).toBeGreaterThanOrEqual(0);
+    expect(deltaIdx).toBeLessThan(resultIdx);
+    expect(eventField<unknown>(events[deltaIdx], "delta")).toEqual(delta);
+  });
+
+  it("emits STATE_SNAPSHOT when payload arrives in raw.result instead of raw.content", async () => {
+    // Regression: serialization fell back to raw.result (`?? raw.result ?? null`)
+    // but state-tool detection only inspected raw.content, so STATE_* events
+    // were silently dropped if upstream used `result` for the state-tool body.
+    // See the "TOOL_CALL_RESULT with object content serializes to JSON" test
+    // above which already exercises the `result` field on a non-state tool.
+    const snapshot = { counter: 99 };
+    const agent = createAgent("tanstack", [
+      tanstackToolCallStart("call1", "AGUISendStateSnapshot"),
+      tanstackToolCallEnd("call1"),
+      {
+        type: "TOOL_CALL_RESULT",
+        toolCallId: "call1",
+        result: { success: true, snapshot },
+      },
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+
+    expectLifecycleWrapped(events);
+
+    const snapshotIdx = events.findIndex(
+      (e) => e.type === EventType.STATE_SNAPSHOT,
+    );
+    expect(snapshotIdx).toBeGreaterThanOrEqual(0);
+    expect(eventField<unknown>(events[snapshotIdx], "snapshot")).toEqual(
+      snapshot,
+    );
+  });
+
+  it("does NOT emit STATE_* for non-state tool results", async () => {
+    const agent = createAgent("tanstack", [
+      tanstackToolCallStart("call1", "getWeather"),
+      tanstackToolCallEnd("call1"),
+      tanstackToolCallResult("call1", {
+        snapshot: { spoofed: true },
+        delta: [{ op: "x" }],
+      }),
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+
+    expectLifecycleWrapped(events);
+
+    expect(
+      events.find(
+        (e) =>
+          e.type === EventType.STATE_SNAPSHOT ||
+          e.type === EventType.STATE_DELTA,
+      ),
+    ).toBeUndefined();
+    expect(
+      events.find((e) => e.type === EventType.TOOL_CALL_RESULT),
+    ).toBeDefined();
+  });
+});
+
+describe("TanStack AI converter — reasoning", () => {
+  it("emits the full REASONING lifecycle for reasoning chunks", async () => {
+    const agent = createAgent("tanstack", [
+      tanstackReasoningStart("r1"),
+      tanstackReasoningMessageStart("r1"),
+      tanstackReasoningMessageContent("r1", "thinking"),
+      tanstackReasoningMessageEnd("r1"),
+      tanstackReasoningEnd("r1"),
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+
+    expectLifecycleWrapped(events);
+
+    // Strip the lifecycle wrap and inspect the inner sequence.
+    const inner = events.slice(1, -1).map((e) => e.type);
+    expect(inner).toEqual([
+      EventType.REASONING_START,
+      EventType.REASONING_MESSAGE_START,
+      EventType.REASONING_MESSAGE_CONTENT,
+      EventType.REASONING_MESSAGE_END,
+      EventType.REASONING_END,
+    ]);
+  });
+
+  it("auto-closes an open reasoning lifecycle when text starts", async () => {
+    const agent = createAgent("tanstack", [
+      tanstackReasoningStart("r1"),
+      tanstackReasoningMessageStart("r1"),
+      tanstackReasoningMessageContent("r1", "thinking"),
+      // No REASONING_MESSAGE_END / REASONING_END before text
+      tanstackTextChunk("Hi"),
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+    const types = events.map((e) => e.type);
+
+    const reasonEndIdx = types.indexOf(EventType.REASONING_END);
+    const reasonMsgEndIdx = types.indexOf(EventType.REASONING_MESSAGE_END);
+    const textIdx = types.indexOf(EventType.TEXT_MESSAGE_CHUNK);
+
+    expect(reasonMsgEndIdx).toBeGreaterThan(-1);
+    expect(reasonEndIdx).toBeGreaterThan(-1);
+    expect(reasonEndIdx).toBeLessThan(textIdx);
+  });
+
+  it("auto-closes an open reasoning lifecycle when a tool call starts", async () => {
+    const agent = createAgent("tanstack", [
+      tanstackReasoningStart("r1"),
+      tanstackReasoningMessageStart("r1"),
+      tanstackReasoningMessageContent("r1", "..."),
+      tanstackToolCallStart("t1", "x"),
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+    const types = events.map((e) => e.type);
+
+    expect(types).toContain(EventType.REASONING_MESSAGE_END);
+    expect(types).toContain(EventType.REASONING_END);
+    expect(types.indexOf(EventType.REASONING_END)).toBeLessThan(
+      types.indexOf(EventType.TOOL_CALL_START),
+    );
+  });
+
+  it("auto-closes when the stream ends without explicit reasoning end", async () => {
+    const agent = createAgent("tanstack", [
+      tanstackReasoningStart("r1"),
+      tanstackReasoningMessageStart("r1"),
+      tanstackReasoningMessageContent("r1", "x"),
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+    const types = events.map((e) => e.type);
+
+    expect(types).toContain(EventType.REASONING_MESSAGE_END);
+    expect(types).toContain(EventType.REASONING_END);
+  });
+
+  it("emits REASONING_MESSAGE_END before REASONING_END when upstream sends END with message still open", async () => {
+    // Regression: if the converter received REASONING_END while a message
+    // was still open, it cleared run-open and emitted END only — leaving
+    // message-open true. The next non-reasoning chunk then triggered
+    // closeReasoningIfOpen, which emitted MSG_END AFTER END (wrong order).
+    // Fix: REASONING_END handler closes any open message first.
+    const agent = createAgent("tanstack", [
+      tanstackReasoningStart("r1"),
+      tanstackReasoningMessageStart("r1"),
+      tanstackReasoningMessageContent("r1", "thinking"),
+      // Upstream skips MSG_END and goes straight to END
+      tanstackReasoningEnd("r1"),
+      tanstackTextChunk("Hi"),
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+    const types = events.map((e) => e.type);
+
+    const msgEndIdx = types.indexOf(EventType.REASONING_MESSAGE_END);
+    const endIdx = types.indexOf(EventType.REASONING_END);
+
+    expect(msgEndIdx).toBeGreaterThan(-1);
+    expect(endIdx).toBeGreaterThan(-1);
+    expect(msgEndIdx).toBeLessThan(endIdx);
+    // No duplicate MSG_END or END from auto-close on the text chunk
+    expect(
+      types.filter((t) => t === EventType.REASONING_MESSAGE_END),
+    ).toHaveLength(1);
+    expect(types.filter((t) => t === EventType.REASONING_END)).toHaveLength(1);
+  });
+
+  it("auto-closes prior reasoning run when a new REASONING_START arrives without END", async () => {
+    // Regression: REASONING_START used to overwrite reasoningMessageId
+    // unconditionally, orphaning the prior run's MSG_END / END.
+    // Fix: REASONING_START handler calls closeReasoningIfOpen() first.
+    const agent = createAgent("tanstack", [
+      tanstackReasoningStart("r1"),
+      tanstackReasoningMessageStart("r1"),
+      tanstackReasoningMessageContent("r1", "first"),
+      // Second START with no intervening END
+      tanstackReasoningStart("r2"),
+      tanstackReasoningMessageStart("r2"),
+      tanstackReasoningMessageContent("r2", "second"),
+      tanstackReasoningMessageEnd("r2"),
+      tanstackReasoningEnd("r2"),
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+    const types = events.map((e) => e.type);
+
+    // Two complete START → MSG_START → ... → MSG_END → END sequences
+    expect(types.filter((t) => t === EventType.REASONING_START)).toHaveLength(
+      2,
+    );
+    expect(types.filter((t) => t === EventType.REASONING_END)).toHaveLength(2);
+    expect(
+      types.filter((t) => t === EventType.REASONING_MESSAGE_END),
+    ).toHaveLength(2);
+
+    // First START's prior message gets closed BEFORE the second START
+    const firstEndIdx = types.indexOf(EventType.REASONING_END);
+    const secondStartIdx = types.indexOf(
+      EventType.REASONING_START,
+      firstEndIdx + 1,
+    );
+    expect(firstEndIdx).toBeLessThan(secondStartIdx);
+  });
+
+  it("does NOT duplicate REASONING_MESSAGE_END when upstream emits it explicitly before text", async () => {
+    // Regression: a single isInReasoning flag conflated message-open with
+    // run-open, so closeReasoningIfOpen on TEXT_MESSAGE_CONTENT emitted a
+    // second REASONING_MESSAGE_END after upstream's own. Track message-open
+    // and run-open separately so closeReasoningIfOpen owes only what's still
+    // open.
+    const agent = createAgent("tanstack", [
+      tanstackReasoningStart("r1"),
+      tanstackReasoningMessageStart("r1"),
+      tanstackReasoningMessageContent("r1", "thinking"),
+      tanstackReasoningMessageEnd("r1"),
+      // No explicit REASONING_END before text — closeReasoningIfOpen should
+      // emit REASONING_END but NOT a second REASONING_MESSAGE_END.
+      tanstackTextChunk("Hi"),
+    ]);
+    const events = await collectEvents(agent.run(createDefaultInput()));
+    const types = events.map((e) => e.type);
+
+    const msgEndCount = types.filter(
+      (t) => t === EventType.REASONING_MESSAGE_END,
+    ).length;
+    const endCount = types.filter((t) => t === EventType.REASONING_END).length;
+
+    expect(msgEndCount).toBe(1);
+    expect(endCount).toBe(1);
+    expect(types.indexOf(EventType.REASONING_END)).toBeLessThan(
+      types.indexOf(EventType.TEXT_MESSAGE_CHUNK),
+    );
+  });
+});

package/src/agent/converters/tanstack.ts

@@ -8,6 +8,13 @@ import {
   ToolCallEndEvent,
   ToolCallStartEvent,
   ToolCallResultEvent,
+  StateSnapshotEvent,
+  StateDeltaEvent,
+  ReasoningStartEvent,
+  ReasoningMessageStartEvent,
+  ReasoningMessageContentEvent,
+  ReasoningMessageEndEvent,
+  ReasoningEndEvent,
 } from "@ag-ui/client";
 import { randomUUID } from "@copilotkit/shared";
 
@@ -229,6 +236,44 @@ export async function* convertTanStackStream(
   abortSignal: AbortSignal,
 ): AsyncGenerator<BaseEvent> {
   const messageId = randomUUID();
+  const toolNamesById = new Map<string, string>();
+  // Track the reasoning lifecycle at two granularities so closeReasoningIfOpen
+  // emits exactly the events still owed. A single boolean conflates the run
+  // (REASONING_START → REASONING_END) with the message
+  // (REASONING_MESSAGE_START → REASONING_MESSAGE_END) and produces a duplicate
+  // REASONING_MESSAGE_END when upstream emits MSG_END but not END before
+  // text/tools resume.
+  let reasoningRunOpen = false;
+  let reasoningMessageOpen = false;
+  let reasoningMessageId = randomUUID();
+
+  function* closeReasoningIfOpen(): Generator<BaseEvent> {
+    if (reasoningMessageOpen) {
+      reasoningMessageOpen = false;
+      const msgEnd: ReasoningMessageEndEvent = {
+        type: EventType.REASONING_MESSAGE_END,
+        messageId: reasoningMessageId,
+      };
+      yield msgEnd;
+    }
+    if (reasoningRunOpen) {
+      reasoningRunOpen = false;
+      const end: ReasoningEndEvent = {
+        type: EventType.REASONING_END,
+        messageId: reasoningMessageId,
+      };
+      yield end;
+    }
+  }
+
+  // TanStack's chat() engine runs a multi-turn agent loop: after the model
+  // returns tool calls, the engine tries to execute them and re-prompt. This
+  // produces a second round of TOOL_CALL_START / TOOL_CALL_END events that
+  // duplicate the ones from the first streaming pass. The CopilotKit runtime
+  // handles tool execution externally (via the frontend SDK), so we must stop
+  // converting events once the TanStack adapter signals the first turn is
+  // complete with RUN_FINISHED.
+  let runFinished = false;
 
   for await (const chunk of stream) {
     if (abortSignal.aborted) break;
@@ -236,7 +281,17 @@ export async function* convertTanStackStream(
     const raw = chunk as Record<string, unknown>;
     const type = raw.type as string;
 
-    if (type === "TEXT_MESSAGE_CONTENT" && raw.delta) {
+    // Stop converting after the first RUN_FINISHED — any subsequent events
+    // come from TanStack's internal tool-execution loop and would produce
+    // duplicate TOOL_CALL_END events that violate the ag-ui verify middleware.
+    if (type === "RUN_FINISHED") {
+      runFinished = true;
+      continue;
+    }
+    if (runFinished) continue;
+
+    if (type === "TEXT_MESSAGE_CONTENT" && raw.delta != null) {
+      yield* closeReasoningIfOpen();
       const textEvent: TextMessageChunkEvent = {
         type: EventType.TEXT_MESSAGE_CHUNK,
         role: "assistant",
@@ -245,6 +300,8 @@ export async function* convertTanStackStream(
       };
       yield textEvent;
     } else if (type === "TOOL_CALL_START") {
+      yield* closeReasoningIfOpen();
+      toolNamesById.set(raw.toolCallId as string, raw.toolCallName as string);
       const startEvent: ToolCallStartEvent = {
         type: EventType.TOOL_CALL_START,
         parentMessageId: messageId,
@@ -253,6 +310,7 @@ export async function* convertTanStackStream(
       };
       yield startEvent;
     } else if (type === "TOOL_CALL_ARGS") {
+      yield* closeReasoningIfOpen();
       const argsEvent: ToolCallArgsEvent = {
         type: EventType.TOOL_CALL_ARGS,
         toolCallId: raw.toolCallId as string,
@@ -260,35 +318,134 @@ export async function* convertTanStackStream(
       };
       yield argsEvent;
     } else if (type === "TOOL_CALL_END") {
+      yield* closeReasoningIfOpen();
       const endEvent: ToolCallEndEvent = {
         type: EventType.TOOL_CALL_END,
         toolCallId: raw.toolCallId as string,
       };
       yield endEvent;
     } else if (type === "TOOL_CALL_RESULT") {
+      yield* closeReasoningIfOpen();
+      const toolCallId = raw.toolCallId as string;
+      const toolName = toolNamesById.get(toolCallId);
+      // Accept the payload from either `content` (canonical TanStack shape)
+      // or `result` (alternate shape used by some adapters / tests). Both
+      // state-tool detection and the final TOOL_CALL_RESULT serialization
+      // must read the same field, otherwise STATE_SNAPSHOT/STATE_DELTA can
+      // be silently dropped when upstream uses `result`.
+      const rawPayload = raw.content ?? raw.result;
+
+      const parsedContent =
+        typeof rawPayload === "string" ? safeParse(rawPayload) : rawPayload;
+
+      if (
+        toolName === "AGUISendStateSnapshot" &&
+        parsedContent &&
+        typeof parsedContent === "object" &&
+        "snapshot" in parsedContent
+      ) {
+        const stateSnapshotEvent: StateSnapshotEvent = {
+          type: EventType.STATE_SNAPSHOT,
+          snapshot: (parsedContent as Record<string, unknown>).snapshot,
+        };
+        yield stateSnapshotEvent;
+      }
+
+      if (
+        toolName === "AGUISendStateDelta" &&
+        parsedContent &&
+        typeof parsedContent === "object" &&
+        "delta" in parsedContent
+      ) {
+        const stateDeltaEvent: StateDeltaEvent = {
+          type: EventType.STATE_DELTA,
+          delta: (parsedContent as Record<string, unknown>).delta as never,
+        };
+        yield stateDeltaEvent;
+      }
+
       let serializedContent: string;
-      if (typeof raw.content === "string") {
-        serializedContent = raw.content;
+      if (typeof rawPayload === "string") {
+        serializedContent = rawPayload;
       } else {
         try {
-          serializedContent = JSON.stringify(raw.content ?? raw.result ?? null);
+          serializedContent = JSON.stringify(rawPayload ?? null);
         } catch {
           serializedContent = "[Unserializable tool result]";
         }
       }
+
       const resultEvent: ToolCallResultEvent = {
         type: EventType.TOOL_CALL_RESULT,
         role: "tool",
         messageId: randomUUID(),
-        toolCallId: raw.toolCallId as string,
+        toolCallId,
         content: serializedContent,
       };
       yield resultEvent;
+      toolNamesById.delete(toolCallId);
+    } else if (type === "REASONING_START") {
+      // If a prior reasoning run is still open (no REASONING_END before this
+      // new START), close it cleanly first so MSG_END / END pair correctly.
+      yield* closeReasoningIfOpen();
+      reasoningRunOpen = true;
+      reasoningMessageId = (raw.messageId as string) ?? randomUUID();
+      const startEvt: ReasoningStartEvent = {
+        type: EventType.REASONING_START,
+        messageId: reasoningMessageId,
+      };
+      yield startEvt;
+    } else if (type === "REASONING_MESSAGE_START") {
+      reasoningMessageOpen = true;
+      const evt: ReasoningMessageStartEvent = {
+        type: EventType.REASONING_MESSAGE_START,
+        messageId: reasoningMessageId,
+        role: "reasoning",
+      };
+      yield evt;
+    } else if (type === "REASONING_MESSAGE_CONTENT") {
+      const evt: ReasoningMessageContentEvent = {
+        type: EventType.REASONING_MESSAGE_CONTENT,
+        messageId: reasoningMessageId,
+        delta: raw.delta as string,
+      };
+      yield evt;
+    } else if (type === "REASONING_MESSAGE_END") {
+      reasoningMessageOpen = false;
+      const evt: ReasoningMessageEndEvent = {
+        type: EventType.REASONING_MESSAGE_END,
+        messageId: reasoningMessageId,
+      };
+      yield evt;
+    } else if (type === "REASONING_END") {
+      // If upstream sends REASONING_END while a message is still open, emit
+      // the missing REASONING_MESSAGE_END FIRST so the closing pair stays in
+      // order (MSG_END before END). Otherwise the next non-reasoning chunk
+      // would trigger closeReasoningIfOpen and emit MSG_END after END.
+      if (reasoningMessageOpen) {
+        reasoningMessageOpen = false;
+        const msgEnd: ReasoningMessageEndEvent = {
+          type: EventType.REASONING_MESSAGE_END,
+          messageId: reasoningMessageId,
+        };
+        yield msgEnd;
+      }
+      reasoningRunOpen = false;
+      const evt: ReasoningEndEvent = {
+        type: EventType.REASONING_END,
+        messageId: reasoningMessageId,
+      };
+      yield evt;
     }
-    // Unhandled chunk types are silently ignored.
-    // Known gaps: STATE_SNAPSHOT, STATE_DELTA, and REASONING events are not
-    // converted from TanStack streams. Shared state and reasoning will not
-    // surface when using the TanStack backend. Use the AI SDK backend if these
-    // features are required.
+  }
+
+  yield* closeReasoningIfOpen();
+}
+
+function safeParse(value: string): unknown {
+  try {
+    return JSON.parse(value);
+  } catch {
+    return value;
   }
 }