test assistant AI SDK v7 transport

2026-08-13 19:57:02 +02:00
1 parent 876f312a2e
commit 8f0efd508e
@@ -0,0 +1,99 @@
+// @vitest-environment node
+
+import { AssistantChatTransport } from "@assistant-ui/react-ai-sdk";
+import { jsonSchema, readUIMessageStream, simulateReadableStream, stepCountIs, streamText, tool } from "ai";
+import { MockLanguageModelV4 } from "ai/test";
+import { describe, expect, test, vi } from "vitest";
+
+const usage = {
+  inputTokens: { total: 8, noCache: 8, cacheRead: 0, cacheWrite: 0 },
+  outputTokens: { total: 4, text: 4, reasoning: 0 },
+};
+
+const finish = (unified: "stop" | "tool-calls") => ({
+  type: "finish" as const,
+  finishReason: { unified, raw: unified },
+  usage,
+});
+
+describe("assistant-ui AI SDK v7 transport", () => {
+  test("round-trips a multi-step server tool response into one UI message", async () => {
+    const model = new MockLanguageModelV4({
+      doStream: [
+        {
+          stream: simulateReadableStream({
+            chunks: [
+              { type: "stream-start" as const, warnings: [] },
+              {
+                type: "tool-call" as const,
+                toolCallId: "lookup-call",
+                toolName: "test_lookup",
+                input: '{"query":"homarr"}',
+              },
+              finish("tool-calls"),
+            ],
+            chunkDelayInMs: null,
+          }),
+        },
+        {
+          stream: simulateReadableStream({
+            chunks: [
+              { type: "stream-start" as const, warnings: [] },
+              { type: "text-start" as const, id: "answer" },
+              { type: "text-delta" as const, id: "answer", delta: "Migration stream complete." },
+              { type: "text-end" as const, id: "answer" },
+              finish("stop"),
+            ],
+            chunkDelayInMs: null,
+          }),
+        },
+      ],
+    });
+    const result = streamText({
+      model,
+      prompt: "Verify the migrated assistant stream.",
+      tools: {
+        test_lookup: tool({
+          inputSchema: jsonSchema({
+            type: "object",
+            properties: { query: { type: "string" } },
+            required: ["query"],
+            additionalProperties: false,
+          }),
+          execute: ({ query }) => ({ query, value: "ok" }),
+        }),
+      },
+      stopWhen: stepCountIs(2),
+    });
+    const fetchMock = vi.fn(async () =>
+      result.toUIMessageStreamResponse({
+        sendReasoning: true,
+        messageMetadata: ({ part }) => (part.type === "finish" ? { runtime: "ai-sdk-v7" } : undefined),
+      }),
+    );
+    const transport = new AssistantChatTransport({
+      api: "https://homarr.test/api/assistant/chat",
+      fetch: fetchMock,
+    });
+
+    const stream = await transport.sendMessages({
+      trigger: "submit-message",
+      chatId: "thread-1",
+      messageId: undefined,
+      messages: [{ id: "user-1", role: "user", parts: [{ type: "text", text: "Run migration test" }] }],
+      abortSignal: undefined,
+    });
+    let finalMessage: unknown;
+    for await (const message of readUIMessageStream({ stream })) finalMessage = message;
+
+    expect(fetchMock).toHaveBeenCalledOnce();
+    expect(finalMessage).toMatchObject({
+      role: "assistant",
+      metadata: { runtime: "ai-sdk-v7" },
+    });
+    expect(JSON.stringify(finalMessage)).toContain('"type":"tool-test_lookup"');
+    expect(JSON.stringify(finalMessage)).toContain('"state":"output-available"');
+    expect(JSON.stringify(finalMessage)).toContain('"value":"ok"');
+    expect(JSON.stringify(finalMessage)).toContain("Migration stream complete.");
+  });
+});