@assistant-ui/react-ag-ui 0.1.0 → 0.1.2

package/src/runtime/adapter/replay-content.test.ts

@@ -0,0 +1,235 @@
+import { describe, expect, it } from "vitest";
+import type {
+  ThreadAssistantMessage,
+  ToolCallMessagePart,
+} from "@assistant-ui/core";
+import {
+  createReplayContentMerger,
+  preserveToolResult,
+} from "./replay-content";
+
+const call = (
+  id: string,
+  extra: Partial<ToolCallMessagePart> = {},
+): ToolCallMessagePart => ({
+  type: "tool-call",
+  toolCallId: id,
+  toolName: "lookup",
+  args: {},
+  argsText: "{}",
+  ...extra,
+});
+
+const assistant = (
+  id: string,
+  content: ThreadAssistantMessage["content"],
+): ThreadAssistantMessage => ({
+  id,
+  role: "assistant",
+  createdAt: new Date(0),
+  status: { type: "requires-action", reason: "tool-calls" },
+  content,
+  metadata: {
+    unstable_state: null,
+    unstable_annotations: [],
+    unstable_data: [],
+    steps: [],
+    custom: {},
+  },
+});
+
+describe("live replay content", () => {
+  it("keeps a locally supplied falsy result and its details until a fresh result arrives", () => {
+    const prior = call("tool", {
+      result: 0,
+      artifact: null,
+      isError: false,
+      modelContent: [],
+    });
+    expect(preserveToolResult(prior, call("tool"))).toMatchObject({
+      result: 0,
+      artifact: null,
+      isError: false,
+      modelContent: [],
+    });
+    expect(
+      preserveToolResult(prior, call("tool", { result: "fresh" })),
+    ).toMatchObject({
+      result: "fresh",
+    });
+  });
+
+  it("updates a flattened child in its nested owner while retaining unreplayed children", () => {
+    const mergeReplayContent = createReplayContentMerger();
+    const previous = [
+      call("root", {
+        messages: [
+          assistant("sub", [
+            call("child", { args: { old: true }, argsText: '{"old":true}' }),
+            call("pending"),
+          ]),
+        ],
+      }),
+    ];
+    const next = [
+      call("root", {
+        messages: [
+          {
+            ...assistant("sub", [{ type: "text", text: "fresh output" }]),
+            status: { type: "complete", reason: "stop" } as const,
+          },
+        ],
+      }),
+      call("child", {
+        args: { fresh: true },
+        argsText: '{"fresh":true}',
+        result: "ok",
+      }),
+    ];
+
+    const merged = mergeReplayContent(previous, next);
+    expect(merged).toHaveLength(1);
+    const root = merged[0] as ToolCallMessagePart;
+    const sub = root.messages?.[0] as ThreadAssistantMessage;
+    expect(sub.status).toEqual({ type: "complete", reason: "stop" });
+    expect(sub.content).toEqual([
+      expect.objectContaining({
+        toolCallId: "child",
+        args: { fresh: true },
+        result: "ok",
+      }),
+      expect.objectContaining({ toolCallId: "pending" }),
+      { type: "text", text: "fresh output" },
+    ]);
+  });
+
+  it("uses incoming part and nested message order while retaining an unreplayed child", () => {
+    const mergeReplayContent = createReplayContentMerger();
+    const previous = [
+      call("root", {
+        messages: [
+          assistant("first", [call("pending")]),
+          assistant("second", [{ type: "text", text: "old" }]),
+        ],
+      }),
+    ];
+    const next = [
+      { type: "text" as const, text: "fresh introduction" },
+      call("root", {
+        messages: [
+          assistant("second", [{ type: "text", text: "fresh second" }]),
+          assistant("first", [{ type: "text", text: "fresh first" }]),
+        ],
+      }),
+    ];
+
+    const merged = mergeReplayContent(previous, next);
+    expect(merged[0]).toEqual({ type: "text", text: "fresh introduction" });
+    const root = merged[1] as ToolCallMessagePart;
+    expect(root.messages?.map((message) => message.id)).toEqual([
+      "second",
+      "first",
+    ]);
+    expect((root.messages?.[1] as ThreadAssistantMessage).content).toEqual([
+      expect.objectContaining({ toolCallId: "pending" }),
+      { type: "text", text: "fresh first" },
+    ]);
+  });
+  it("transfers a completed flat call into its incoming nested owner without retaining a duplicate", () => {
+    const mergeReplayContent = createReplayContentMerger();
+    const interactions = {
+      entries: [
+        { type: "action" as const, occurredAt: 1, payload: { approved: true } },
+      ],
+    };
+    const merged = mergeReplayContent(
+      [call("child", { result: 0, unstable_interactions: interactions })],
+      [
+        call("parent", {
+          messages: [
+            assistant("sub", [
+              call("child", {
+                args: { fresh: true },
+                argsText: '{"fresh":true}',
+              }),
+            ]),
+          ],
+        }),
+      ],
+    );
+    expect(merged).toHaveLength(1);
+    const parent = merged[0] as ToolCallMessagePart;
+    expect(parent.toolCallId).toBe("parent");
+    expect((parent.messages?.[0] as ThreadAssistantMessage).content).toEqual([
+      expect.objectContaining({
+        toolCallId: "child",
+        result: 0,
+        args: { fresh: true },
+        unstable_interactions: interactions,
+      }),
+    ]);
+  });
+
+  it("removes a replay-owned tool after an unrelated content update", () => {
+    const merge = createReplayContentMerger();
+    const historical = call("historical", { result: "saved" });
+    const surface = call("a2ui:surface");
+    const first = merge([historical], [surface]);
+    expect(
+      first.map((part) => part.type === "tool-call" && part.toolCallId),
+    ).toEqual(["historical", "a2ui:surface"]);
+
+    const second = merge(first, [{ type: "text", text: "continued" }]);
+    expect(second).toEqual([historical, { type: "text", text: "continued" }]);
+  });
+
+  it("retains a local tool result until the replay explicitly removes its call", () => {
+    const merge = createReplayContentMerger();
+    const updated = merge([call("child", { result: 0 })], [call("child")]);
+    expect(updated).toEqual([expect.objectContaining({ result: 0 })]);
+    expect(merge(updated, [])).toEqual([]);
+  });
+
+  it("keeps a removed parent only while its incoming flattened child needs it", () => {
+    const merge = createReplayContentMerger();
+    const parent = call("parent", {
+      messages: [assistant("sub", [call("child")])],
+    });
+    const first = merge([], [parent]);
+    const second = merge(first, [call("child", { result: "done" })]);
+    expect(second).toHaveLength(1);
+    expect((second[0] as ToolCallMessagePart).messages?.[0]).toMatchObject({
+      content: [
+        expect.objectContaining({ toolCallId: "child", result: "done" }),
+      ],
+    });
+    expect(merge(second, [])).toEqual([]);
+  });
+
+  it("scopes matching nested message identities to their tool container", () => {
+    const merge = createReplayContentMerger();
+    const first = merge(
+      [],
+      [
+        call("left", {
+          messages: [assistant("same", [{ type: "text", text: "left" }])],
+        }),
+      ],
+    );
+    const second = merge(first, [
+      call("right", {
+        messages: [assistant("same", [{ type: "text", text: "right" }])],
+      }),
+    ]);
+    expect(second).toEqual([
+      expect.objectContaining({
+        toolCallId: "right",
+        messages: [
+          expect.objectContaining({
+            content: [{ type: "text", text: "right" }],
+          }),
+        ],
+      }),
+    ]);
+  });
+});

package/src/runtime/adapter/replay-content.ts

@@ -0,0 +1,306 @@
+import type {
+  ThreadAssistantMessage,
+  ThreadMessage,
+  ToolCallMessagePart,
+} from "@assistant-ui/core";
+import {
+  iterateToolCallParts,
+  mapToolCallPartsDeep,
+} from "@assistant-ui/core/internal";
+
+import { preserveToolApprovalDecision } from "./tool-approval";
+
+type Content = ThreadAssistantMessage["content"];
+type Part = Content[number];
+
+export function indexReplayToolOwners(
+  owners: readonly ThreadAssistantMessage[],
+): ReadonlyMap<string, string> {
+  const ownerByToolCallId = new Map<string, string>();
+  for (const owner of owners) {
+    for (const part of iterateToolCallParts(owner.content)) {
+      if (!ownerByToolCallId.has(part.toolCallId))
+        ownerByToolCallId.set(part.toolCallId, owner.id);
+    }
+  }
+  return ownerByToolCallId;
+}
+
+export function groupReplayToolContent(
+  owners: readonly ThreadAssistantMessage[],
+  content: Content,
+  fallbackMessageId: string,
+): Map<string, Part[]> {
+  const ownerByToolCallId = indexReplayToolOwners(owners);
+  const grouped = new Map<string, Part[]>([[fallbackMessageId, []]]);
+  for (const part of content) {
+    const ownerId =
+      part.type === "tool-call"
+        ? (ownerByToolCallId.get(part.toolCallId) ?? fallbackMessageId)
+        : fallbackMessageId;
+    const parts = grouped.get(ownerId) ?? [];
+    parts.push(part);
+    grouped.set(ownerId, parts);
+  }
+  return grouped;
+}
+
+export function preserveToolResult(
+  previous: ToolCallMessagePart,
+  next: ToolCallMessagePart,
+): ToolCallMessagePart {
+  if (previous.result === undefined || next.result !== undefined) return next;
+  return {
+    ...next,
+    result: previous.result,
+    ...(previous.artifact !== undefined ? { artifact: previous.artifact } : {}),
+    ...(previous.isError !== undefined ? { isError: previous.isError } : {}),
+    ...(previous.modelContent !== undefined
+      ? { modelContent: previous.modelContent }
+      : {}),
+  };
+}
+
+function mergeOrdered<T>(
+  previous: readonly (readonly [T, string])[],
+  next: readonly (readonly [T, string])[],
+  merge: (old: T, fresh: T) => T,
+): T[] {
+  const old = new Map(previous.map(([value, key]) => [key, value] as const));
+  const indexes = new Map(next.map(([, key], index) => [key, index] as const));
+  const insertions: T[][] = Array.from({ length: next.length + 1 }, () => []);
+  let predecessor = -1;
+  for (const [value, key] of previous) {
+    const index = indexes.get(key);
+    if (index === undefined) insertions[predecessor + 1]!.push(value);
+    else predecessor = index;
+  }
+  const result: T[] = [];
+  for (const [index, [value, key]] of next.entries()) {
+    const prior = old.get(key);
+    result.push(
+      ...insertions[index]!,
+      prior === undefined ? value : merge(prior, value),
+    );
+  }
+  result.push(...insertions[next.length]!);
+  return result;
+}
+
+function mergeMessages(
+  previous: readonly ThreadMessage[],
+  next: readonly ThreadMessage[],
+): readonly ThreadMessage[] {
+  return mergeOrdered(
+    previous.map((message) => [message, message.id] as const),
+    next.map((message) => [message, message.id] as const),
+    (old, fresh) =>
+      old.role === "assistant" && fresh.role === "assistant"
+        ? { ...fresh, content: mergeReplayContent(old.content, fresh.content) }
+        : fresh,
+  );
+}
+
+function mergeTool(
+  previous: ToolCallMessagePart,
+  next: ToolCallMessagePart,
+): ToolCallMessagePart {
+  const result = preserveToolResult(previous, next);
+  const approval = preserveToolApprovalDecision(previous, next);
+  return {
+    ...result,
+    ...(approval ? { approval } : {}),
+    ...(result.unstable_interactions === undefined &&
+    previous.unstable_interactions !== undefined
+      ? { unstable_interactions: previous.unstable_interactions }
+      : {}),
+    ...(previous.messages !== undefined || next.messages !== undefined
+      ? {
+          messages: mergeMessages(previous.messages ?? [], next.messages ?? []),
+        }
+      : {}),
+  };
+}
+
+function partKey(part: Part): string {
+  if (part.type === "tool-call") return `tool:${part.toolCallId}`;
+  if (part.type === "text" || part.type === "reasoning") return part.type;
+  if ("id" in part && typeof part.id === "string" && part.id)
+    return `${part.type}:${part.id}`;
+  if (part.type === "data") return `data:${part.name}`;
+  return part.type;
+}
+
+function keyed(content: Content): [Part, string][] {
+  const ordinals = new Map<string, number>();
+  return content.map((part) => {
+    const base = partKey(part);
+    const ordinal = ordinals.get(base) ?? 0;
+    ordinals.set(base, ordinal + 1);
+    return [part, `${base}:${ordinal}`] as [Part, string];
+  });
+}
+
+function identity(part: Part, key: string, scope: readonly string[]): string {
+  return JSON.stringify(
+    part.type === "tool-call" ? ["tool", part.toolCallId] : [...scope, key],
+  );
+}
+
+function visitMessages(
+  messages: readonly ThreadMessage[],
+  scope: readonly string[],
+  visit: (key: string) => void,
+): void {
+  for (const message of messages) {
+    const path = [...scope, "message", message.id];
+    visit(JSON.stringify(path));
+    if (message.role === "assistant")
+      visitContent(message.content, path, visit);
+  }
+}
+
+function visitContent(
+  content: Content,
+  scope: readonly string[],
+  visit: (key: string) => void,
+): void {
+  for (const [part, key] of keyed(content)) {
+    visit(identity(part, key, scope));
+    if (part.type === "tool-call" && part.messages)
+      visitMessages(part.messages, ["tool", part.toolCallId], visit);
+  }
+}
+
+function pruneMessages(
+  messages: readonly ThreadMessage[],
+  scope: readonly string[],
+  present: ReadonlySet<string>,
+  owned: ReadonlySet<string>,
+): { messages: readonly ThreadMessage[]; hasCurrent: boolean } {
+  let hasCurrent = false;
+  const kept = messages.flatMap((message) => {
+    const path = [...scope, "message", message.id];
+    const key = JSON.stringify(path);
+    const nested =
+      message.role === "assistant"
+        ? pruneContent(message.content, path, present, owned)
+        : undefined;
+    const current = present.has(key) || (nested?.hasCurrent ?? false);
+    if (owned.has(key) && !current) return [];
+    hasCurrent ||= current;
+    return [
+      nested && message.role === "assistant"
+        ? { ...message, content: nested.content }
+        : message,
+    ];
+  });
+  return { messages: kept, hasCurrent };
+}
+
+function pruneContent(
+  content: Content,
+  scope: readonly string[],
+  present: ReadonlySet<string>,
+  owned: ReadonlySet<string>,
+): { content: Content; hasCurrent: boolean } {
+  let hasCurrent = false;
+  const kept = keyed(content).flatMap(([part, localKey]) => {
+    const key = identity(part, localKey, scope);
+    const nested =
+      part.type === "tool-call" && part.messages
+        ? pruneMessages(
+            part.messages,
+            ["tool", part.toolCallId],
+            present,
+            owned,
+          )
+        : undefined;
+    const current = present.has(key) || (nested?.hasCurrent ?? false);
+    if (owned.has(key) && !current) return [];
+    hasCurrent ||= current;
+    return [nested ? { ...part, messages: nested.messages } : part];
+  });
+  return { content: kept, hasCurrent };
+}
+
+export function createReplayContentMerger(): (
+  previous: Content,
+  next: Content,
+) => Content {
+  // Remember stream-owned identities, never a second copy of tool payloads.
+  const owned = new Set<string>();
+  return (previous, next) => {
+    const present = new Set<string>();
+    visitContent(next, [], (key) => present.add(key));
+    for (const key of present) owned.add(key);
+    return pruneContent(mergeReplayContent(previous, next), [], present, owned)
+      .content;
+  };
+}
+
+function nestedToolIds(content: Content): Set<string> {
+  const nestedIds = new Set<string>();
+  for (const part of content) {
+    if (part.type !== "tool-call") continue;
+    for (const message of part.messages ?? []) {
+      if (message.role !== "assistant") continue;
+      for (const child of iterateToolCallParts(message.content)) {
+        nestedIds.add(child.toolCallId);
+      }
+    }
+  }
+  return nestedIds;
+}
+
+function mergeReplayContent(previous: Content, next: Content): Content {
+  const nestedIds = nestedToolIds(previous);
+  const incomingNestedIds = nestedToolIds(next);
+  const movingRoots = new Map(
+    previous.flatMap((part) =>
+      part.type === "tool-call" && incomingNestedIds.has(part.toolCallId)
+        ? [[part.toolCallId, part] as const]
+        : [],
+    ),
+  );
+  const flattened = new Map<string, ToolCallMessagePart>();
+  for (const part of next) {
+    if (part.type === "tool-call" && nestedIds.has(part.toolCallId)) {
+      flattened.set(part.toolCallId, part);
+    }
+  }
+  const prior = (
+    flattened.size
+      ? mapToolCallPartsDeep(previous, (part) => {
+          const fresh = flattened.get(part.toolCallId);
+          return fresh ? mergeTool(part, fresh) : part;
+        }).content
+      : previous
+  ).filter(
+    (part) => part.type !== "tool-call" || !movingRoots.has(part.toolCallId),
+  );
+  let incoming = next.filter(
+    (part) => part.type !== "tool-call" || !flattened.has(part.toolCallId),
+  );
+  if (movingRoots.size) {
+    incoming = [
+      ...mapToolCallPartsDeep(incoming, (part) => {
+        const prior = movingRoots.get(part.toolCallId);
+        return prior ? mergeTool(prior, part) : part;
+      }).content,
+    ];
+  }
+  return mergeOrdered(keyed(prior), keyed(incoming), (old, fresh) => {
+    if (old.type === "tool-call" && fresh.type === "tool-call")
+      return mergeTool(old, fresh);
+    if (
+      (old.type === "text" || old.type === "reasoning") &&
+      (fresh.type === "text" || fresh.type === "reasoning") &&
+      old.text.startsWith(fresh.text) &&
+      old.text.length > fresh.text.length
+    ) {
+      return { ...fresh, text: old.text };
+    }
+    return fresh;
+  });
+}

package/src/runtime/adapter/run-aggregator.metadata.test.ts

@@ -0,0 +1,151 @@
+import { describe, expect, it } from "vitest";
+import type { ChatModelRunResult } from "@assistant-ui/core";
+import { RunAggregator } from "./run-aggregator";
+import { fromAgUiMessages, toAgUiMessages } from "./conversions";
+import { parseAgUiEvent } from "../event-parser";
+
+describe("AG-UI tool result metadata", () => {
+  it("keeps interleaved tool result metadata through finish and message export/import", () => {
+    const expectedSecondMetadata = {
+      trace: "second-final",
+      attempts: 0,
+      optional: null,
+      nested: { value: "after" },
+    };
+    const updates: ChatModelRunResult[] = [];
+    const aggregator = new RunAggregator({
+      showThinking: false,
+      logger: { debug: () => {}, error: () => {} },
+      emit: (update) => updates.push(update),
+    });
+    aggregator.handle({ type: "RUN_STARTED", runId: "run" });
+    for (const toolCallId of ["first", "second", "empty"]) {
+      aggregator.handle({
+        type: "TOOL_CALL_START",
+        toolCallId,
+        toolCallName: "lookup",
+      });
+    }
+
+    for (const event of [
+      {
+        type: "TOOL_CALL_RESULT",
+        toolCallId: "second",
+        messageId: "result-second",
+        content: "0",
+        metadata: {
+          trace: "second",
+          attempts: 0,
+          optional: null,
+          nested: { old: true, value: "before" },
+        },
+      },
+      {
+        type: "TOOL_CALL_RESULT",
+        toolCallId: "first",
+        messageId: "result-first",
+        content: "false",
+        metadata: { trace: "first", labels: [] },
+      },
+      {
+        type: "TOOL_CALL_RESULT",
+        toolCallId: "empty",
+        messageId: "result-empty",
+        content: "",
+        metadata: { trace: "empty" },
+      },
+      {
+        type: "TOOL_CALL_RESULT",
+        toolCallId: "second",
+        messageId: "result-second",
+        content: "0",
+        metadata: { trace: "second-final", nested: { value: "after" } },
+      },
+    ]) {
+      const parsed = parseAgUiEvent(event);
+      if (!parsed) throw new Error("Expected a parsed result event");
+      aggregator.handle(parsed);
+    }
+
+    const liveSecond = updates
+      .at(-1)
+      ?.content?.find(
+        (part) => part.type === "tool-call" && part.toolCallId === "second",
+      );
+    if (liveSecond?.type !== "tool-call") {
+      throw new Error("Expected the streamed second tool call");
+    }
+    expect(liveSecond).toMatchObject({ result: 0 });
+    expect(liveSecond.providerMetadata).toEqual({
+      agui: { messageMetadata: expectedSecondMetadata },
+    });
+    aggregator.handle({ type: "RUN_FINISHED", runId: "run" });
+    const content = updates.at(-1)?.content;
+    expect(content).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          toolCallId: "first",
+          result: false,
+          providerMetadata: {
+            agui: { messageMetadata: { trace: "first", labels: [] } },
+          },
+        }),
+        expect.objectContaining({
+          toolCallId: "second",
+          result: 0,
+          providerMetadata: {
+            agui: {
+              messageMetadata: expectedSecondMetadata,
+            },
+          },
+        }),
+        expect.objectContaining({
+          toolCallId: "empty",
+          result: "",
+          providerMetadata: {
+            agui: { messageMetadata: { trace: "empty" } },
+          },
+        }),
+      ]),
+    );
+    const exported = toAgUiMessages([
+      { id: "assistant", role: "assistant", content: content ?? [] },
+    ]);
+    expect(exported).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          role: "tool",
+          toolCallId: "first",
+          metadata: { trace: "first", labels: [] },
+        }),
+        expect.objectContaining({
+          role: "tool",
+          toolCallId: "second",
+          metadata: expectedSecondMetadata,
+        }),
+        expect.objectContaining({
+          role: "tool",
+          toolCallId: "empty",
+          metadata: { trace: "empty" },
+        }),
+      ]),
+    );
+    expect(fromAgUiMessages(exported)).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          content: expect.arrayContaining([
+            expect.objectContaining({
+              toolCallId: "second",
+              result: 0,
+              providerMetadata: {
+                agui: {
+                  messageMetadata: expectedSecondMetadata,
+                },
+              },
+            }),
+          ]),
+        }),
+      ]),
+    );
+  });
+});