@kognitivedev/vercel-ai-provider 0.1.5 → 0.1.7

package/src/__tests__/wrap-stream-logging.test.ts

@@ -0,0 +1,104 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { createCognitiveLayer } from "../index";
+import { streamText } from "ai";
+import { MockLanguageModelV3, convertArrayToReadableStream } from "ai/test";
+
+describe("wrapStream logging", () => {
+    let fetchCalls: { url: string; body: any }[];
+
+    beforeEach(() => {
+        fetchCalls = [];
+
+        vi.stubGlobal(
+            "fetch",
+            vi.fn(async (url: string | URL | Request, init?: RequestInit) => {
+                const urlStr = typeof url === "string" ? url : url.toString();
+
+                if (urlStr.includes("/api/cognitive/snapshot")) {
+                    return new Response(
+                        JSON.stringify({ systemBlock: "", userContextBlock: "" }),
+                        { status: 200, headers: { "Content-Type": "application/json" } }
+                    );
+                }
+
+                if (urlStr.includes("/api/cognitive/log")) {
+                    const body = JSON.parse(init?.body as string);
+                    fetchCalls.push({ url: urlStr, body });
+                    return new Response(JSON.stringify({ ok: true }), { status: 200 });
+                }
+
+                if (urlStr.includes("/api/cognitive/process")) {
+                    return new Response(JSON.stringify({ ok: true }), { status: 200 });
+                }
+
+                return new Response("not found", { status: 404 });
+            })
+        );
+    });
+
+    it("should include assistant message in logged conversation after streaming", async () => {
+        const mockModel = new MockLanguageModelV3({
+            doStream: async () => ({
+                stream: convertArrayToReadableStream([
+                    { type: "text-start" as const, id: "t1" },
+                    { type: "text-delta" as const, id: "t1", delta: "Hello" },
+                    { type: "text-delta" as const, id: "t1", delta: " world" },
+                    { type: "text-end" as const, id: "t1" },
+                    {
+                        type: "finish" as const,
+                        finishReason: {
+                            unified: "stop" as const,
+                            raw: undefined,
+                        },
+                        usage: {
+                            inputTokens: { total: 10, noCache: undefined, cacheRead: undefined, cacheWrite: undefined },
+                            outputTokens: { total: 5, text: undefined, reasoning: undefined },
+                        },
+                    },
+                ] satisfies import("@ai-sdk/provider").LanguageModelV3StreamPart[]),
+            }),
+        });
+
+        const mockProvider = () => mockModel;
+
+        const cl = createCognitiveLayer({
+            provider: mockProvider,
+            clConfig: {
+                apiKey: "test-api-key",
+                appId: "test-app",
+                projectId: "test-project",
+                processDelayMs: 0,
+                logLevel: "none",
+            },
+        });
+
+        const model = cl("mock-model", {
+            userId: "user-1",
+            projectId: "project-1",
+            sessionId: "session-1",
+        });
+
+        const result = streamText({
+            model,
+            messages: [{ role: "user", content: "Hi" }],
+        });
+
+        // Fully consume the stream
+        const text = await result.text;
+        expect(text).toBe("Hello world");
+
+        // Wait for async logConversation to complete
+        await new Promise((r) => setTimeout(r, 100));
+
+        // Find the log call
+        const logCall = fetchCalls.find((c) => c.url.includes("/api/cognitive/log"));
+        expect(logCall).toBeDefined();
+
+        const messages = logCall!.body.messages;
+        const assistantMsg = messages.find((m: any) => m.role === "assistant");
+        expect(assistantMsg).toBeDefined();
+        expect(assistantMsg.content).toEqual([
+            { type: "text", text: "Hello world" },
+        ]);
+    });
+});