@kognitivedev/vercel-ai-provider 0.1.4 → 0.1.6

package/dist/index.js

@@ -1,4 +1,15 @@
 "use strict";
+var __rest = (this && this.__rest) || function (s, e) {
+    var t = {};
+    for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p) && e.indexOf(p) < 0)
+        t[p] = s[p];
+    if (s != null && typeof Object.getOwnPropertySymbols === "function")
+        for (var i = 0, p = Object.getOwnPropertySymbols(s); i < p.length; i++) {
+            if (e.indexOf(p[i]) < 0 && Object.prototype.propertyIsEnumerable.call(s, p[i]))
+                t[p[i]] = s[p[i]];
+        }
+    return t;
+};
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.createCognitiveLayer = createCognitiveLayer;
 const ai_1 = require("ai");
@@ -37,10 +48,22 @@ function createLogger(logLevel) {
         },
     };
 }
+const PROMPT_CACHE_TTL_MS = 60000; // 1 minute
+/**
+ * Interpolate {{variable}} placeholders in a template string.
+ * Unmatched variables are left as-is.
+ */
+function interpolateTemplate(content, variables) {
+    return content.replace(/\{\{(\w+)\}\}/g, (_, key) => { var _a; return (_a = variables[key]) !== null && _a !== void 0 ? _a : `{{${key}}}`; });
+}
 // Session-scoped snapshot cache: sessionKey → formatted memory block
 const sessionSnapshots = new Map();
 // Regex to detect if memory has already been injected
 const MEMORY_TAG_REGEX = /<MemoryContext>/i;
+// Symbol-keyed property to track session settings on model objects
+const SESSION_KEY = Symbol.for("cl:session");
+// Session key → prompt metadata (populated by cl.streamText/cl.generateText, read by middleware)
+const sessionPromptMetadata = new Map();
 /**
  * Check if any system message already contains a <MemoryContext> block.
  */
@@ -65,11 +88,43 @@ function createCognitiveLayer(config) {
     // Default to 'info' log level
     const logLevel = clConfig.logLevel || 'info';
     const logger = createLogger(logLevel);
+    const authHeaders = {
+        "Content-Type": "application/json",
+        "Authorization": `Bearer ${clConfig.apiKey}`,
+    };
+    // Prompt cache: slug → CachedPrompt
+    const promptCache = new Map();
+    const resolvePrompt = async (slug) => {
+        const cached = promptCache.get(slug);
+        if (cached && Date.now() - cached.fetchedAt < PROMPT_CACHE_TTL_MS) {
+            logger.debug("Using cached prompt", { slug, version: cached.version });
+            return cached;
+        }
+        const res = await fetch(`${baseUrl}/api/cognitive/prompt?slug=${encodeURIComponent(slug)}`, {
+            headers: { "Authorization": `Bearer ${clConfig.apiKey}` },
+        });
+        if (!res.ok) {
+            const body = await res.text();
+            throw new Error(`Failed to resolve prompt "${slug}": ${res.status} ${body}`);
+        }
+        const data = await res.json();
+        const entry = {
+            promptId: data.promptId,
+            slug: data.slug,
+            version: data.version,
+            content: data.content,
+            fetchedAt: Date.now(),
+            gatewaySlug: data.gatewaySlug,
+        };
+        promptCache.set(slug, entry);
+        logger.info("Prompt resolved", { slug, version: entry.version });
+        return entry;
+    };
     const logConversation = async (payload) => {
         try {
             await fetch(`${baseUrl}/api/cognitive/log`, {
                 method: "POST",
-                headers: { "Content-Type": "application/json" },
+                headers: authHeaders,
                 body: JSON.stringify(Object.assign(Object.assign({}, payload), { type: "conversation", timestamp: new Date().toISOString() })),
             });
         }
@@ -77,12 +132,12 @@ function createCognitiveLayer(config) {
             logger.error("Log failed", e);
         }
     };
-    const triggerProcessing = (userId, agentId, sessionId) => {
+    const triggerProcessing = (userId, projectId, sessionId) => {
         const run = () => {
             fetch(`${baseUrl}/api/cognitive/process`, {
                 method: "POST",
-                headers: { "Content-Type": "application/json" },
-                body: JSON.stringify({ userId, agentId, sessionId }),
+                headers: authHeaders,
+                body: JSON.stringify({ userId, sessionId }),
             }).catch(e => logger.error("Process trigger failed", e));
         };
         if (processDelay > 0) {
@@ -95,154 +150,298 @@ function createCognitiveLayer(config) {
     const withMemorySystemPrompt = (params, incomingMessages, memoryPrompt) => {
         var _a;
         const nextParams = Object.assign({}, params);
-        // 1) If caller provided a top-level system prompt, overwrite it.
-        if (nextParams.system) {
-            nextParams.system = memoryPrompt;
-            return { nextParams, messages: incomingMessages, mode: "overwrite-param" };
-        }
-        // 2) If first message is system, replace its content.
+        // 1) If first message is system, append memory to its content (without mutating original).
         if (incomingMessages.length > 0 && ((_a = incomingMessages[0]) === null || _a === void 0 ? void 0 : _a.role) === "system") {
-            const updated = [...incomingMessages];
-            updated[0] = Object.assign(Object.assign({}, updated[0]), { content: memoryPrompt });
+            const original = incomingMessages[0];
+            const updatedContent = typeof original.content === "string"
+                ? original.content + "\n\n" + memoryPrompt
+                : memoryPrompt;
+            const updated = [Object.assign(Object.assign({}, original), { content: updatedContent }), ...incomingMessages.slice(1)];
             return { nextParams, messages: updated, mode: "overwrite-first-system" };
         }
-        // 3) Otherwise prepend a system message.
+        // 2) Otherwise prepend a system message.
         const updated = [{ role: "system", content: memoryPrompt }, ...incomingMessages];
         return { nextParams, messages: updated, mode: "prepend-system" };
     };
-    return (modelId, settings, providerOptions) => {
-        // Pass provider options through to the underlying provider
-        const model = (providerOptions
-            ? provider(modelId, providerOptions)
-            : provider(modelId));
-        const userId = settings === null || settings === void 0 ? void 0 : settings.userId;
-        const agentId = (settings === null || settings === void 0 ? void 0 : settings.agentId) || clConfig.defaultAgentId || "default";
-        const sessionId = settings === null || settings === void 0 ? void 0 : settings.sessionId;
-        const sessionMissing = !!userId && !sessionId;
-        if (sessionMissing) {
-            logger.warn("sessionId is required to log and process memories; skipping logging until provided.");
-        }
-        return (0, ai_1.wrapLanguageModel)({
-            model,
-            middleware: {
-                async transformParams({ params }) {
-                    if (!userId)
-                        return params;
-                    const incomingMessages = Array.isArray(params.prompt)
-                        ? params.prompt
-                        : [];
-                    // 1) Check if memory is already injected in messages
-                    if (hasExistingMemoryInjection(incomingMessages)) {
-                        logger.debug("Memory already injected, skipping");
-                        return params;
-                    }
-                    // 2) Check session cache
-                    const sessionKey = `${userId}:${agentId}:${sessionId || "default"}`;
-                    let systemPromptToAdd = sessionSnapshots.get(sessionKey);
-                    // 3) Fetch snapshot only if not cached
-                    if (systemPromptToAdd === undefined) {
-                        try {
-                            const url = `${baseUrl}/api/cognitive/snapshot?userId=${userId}&agentId=${agentId}&appId=${clConfig.appId}`;
-                            const res = await fetch(url);
-                            if (res.ok) {
-                                const data = await res.json();
-                                const systemBlock = data.systemBlock || "";
-                                const userContextBlock = data.userContextBlock || "";
-                                systemPromptToAdd =
-                                    systemBlock !== "" || userContextBlock !== ""
-                                        ? `
+    const buildMiddleware = (userId, projectId, sessionId, modelId) => ({
+        specificationVersion: 'v3',
+        async transformParams({ params }) {
+            if (!userId)
+                return params;
+            const incomingMessages = Array.isArray(params.prompt)
+                ? params.prompt
+                : [];
+            // 1) Check if memory is already injected in messages
+            if (hasExistingMemoryInjection(incomingMessages)) {
+                logger.debug("Memory already injected, skipping");
+                return params;
+            }
+            // 2) Check session cache
+            const sessionKey = `${userId}:${projectId}:${sessionId || "default"}`;
+            let systemPromptToAdd = sessionSnapshots.get(sessionKey);
+            // 3) Fetch snapshot only if not cached
+            if (systemPromptToAdd === undefined) {
+                try {
+                    const url = `${baseUrl}/api/cognitive/snapshot?userId=${userId}`;
+                    const res = await fetch(url, {
+                        headers: { "Authorization": `Bearer ${clConfig.apiKey}` },
+                    });
+                    if (res.ok) {
+                        const data = await res.json();
+                        const systemBlock = data.systemBlock || "";
+                        const userContextBlock = data.userContextBlock || "";
+                        systemPromptToAdd =
+                            systemBlock !== "" || userContextBlock !== ""
+                                ? `
 <MemoryContext>
 Use the following memory to stay consistent. Prefer UserContext facts for answers; AgentHeuristics guide style, safety, and priorities.
 ${systemBlock || "None"}
 ${userContextBlock || "None"}
 </MemoryContext>
-                                `.trim()
-                                        : "";
-                                // Cache the snapshot for this session
-                                sessionSnapshots.set(sessionKey, systemPromptToAdd);
-                                logger.info("Snapshot fetched and cached", {
-                                    userId,
-                                    agentId,
-                                    sessionId,
-                                    sessionKey,
-                                    systemLen: systemBlock.length,
-                                    userLen: userContextBlock.length,
-                                });
-                                // At debug level, log the full snapshot data
-                                logger.debug("Full snapshot data", {
-                                    systemBlock,
-                                    userContextBlock,
-                                    rawData: data,
-                                });
-                            }
-                            else {
-                                logger.warn("Snapshot fetch failed", { status: res.status });
-                                systemPromptToAdd = "";
-                                sessionSnapshots.set(sessionKey, systemPromptToAdd);
-                            }
-                        }
-                        catch (e) {
-                            logger.warn("Failed to fetch snapshot", e);
-                            systemPromptToAdd = "";
-                            sessionSnapshots.set(sessionKey, systemPromptToAdd);
-                        }
-                    }
-                    else {
-                        logger.debug("Using cached snapshot for session", { sessionKey });
-                    }
-                    if (!systemPromptToAdd) {
-                        return Object.assign(Object.assign({}, params), { messages: incomingMessages });
-                    }
-                    const { nextParams, messages: messagesWithMemory } = withMemorySystemPrompt(params, incomingMessages, systemPromptToAdd);
-                    logger.info("Injecting memory system prompt", {
-                        sessionKey,
-                        promptLength: systemPromptToAdd.length,
-                    });
-                    logger.debug("Injected prompt content", { systemPromptToAdd });
-                    return Object.assign(Object.assign({}, nextParams), { prompt: messagesWithMemory });
-                },
-                async wrapGenerate({ doGenerate, params }) {
-                    var _a;
-                    const result = await doGenerate();
-                    if (userId && sessionId) {
-                        const messagesInput = params.messages || params.prompt || [];
-                        const resultMessages = (_a = result === null || result === void 0 ? void 0 : result.response) === null || _a === void 0 ? void 0 : _a.messages;
-                        const assistantMessage = (result === null || result === void 0 ? void 0 : result.text)
-                            ? [{ role: "assistant", content: [{ type: "text", text: result.text }] }]
-                            : [];
-                        const finalMessages = Array.isArray(resultMessages) && resultMessages.length > 0
-                            ? resultMessages
-                            : [...messagesInput, ...assistantMessage];
-                        logConversation({
+                            `.trim()
+                                : "";
+                        // Cache the snapshot for this session
+                        sessionSnapshots.set(sessionKey, systemPromptToAdd);
+                        logger.info("Snapshot fetched and cached", {
                             userId,
-                            agentId,
+                            projectId,
                             sessionId,
-                            messages: finalMessages,
-                            modelId,
-                        }).then(() => triggerProcessing(userId, agentId, sessionId));
+                            sessionKey,
+                            systemLen: systemBlock.length,
+                            userLen: userContextBlock.length,
+                        });
+                        // At debug level, log the full snapshot data
+                        logger.debug("Full snapshot data", {
+                            systemBlock,
+                            userContextBlock,
+                            rawData: data,
+                        });
                     }
-                    return result;
-                },
-                async wrapStream({ doStream, params }) {
-                    var _a;
-                    const result = await doStream();
-                    if (userId && sessionId) {
-                        const messagesInput = params.messages || params.prompt || [];
-                        const resultMessages = (_a = result === null || result === void 0 ? void 0 : result.response) === null || _a === void 0 ? void 0 : _a.messages;
-                        const finalMessages = Array.isArray(resultMessages) && resultMessages.length > 0
-                            ? resultMessages
-                            : messagesInput;
-                        logConversation({
-                            userId,
-                            agentId,
-                            sessionId,
-                            messages: finalMessages,
-                            modelId,
-                        }).then(() => triggerProcessing(userId, agentId, sessionId));
+                    else {
+                        logger.warn("Snapshot fetch failed", { status: res.status });
+                        systemPromptToAdd = "";
+                        sessionSnapshots.set(sessionKey, systemPromptToAdd);
                     }
-                    return result;
                 }
+                catch (e) {
+                    logger.warn("Failed to fetch snapshot", e);
+                    systemPromptToAdd = "";
+                    sessionSnapshots.set(sessionKey, systemPromptToAdd);
+                }
+            }
+            else {
+                logger.debug("Using cached snapshot for session", { sessionKey });
+            }
+            if (!systemPromptToAdd) {
+                return params;
+            }
+            const { nextParams, messages: messagesWithMemory } = withMemorySystemPrompt(params, incomingMessages, systemPromptToAdd);
+            logger.info("Injecting memory system prompt", {
+                sessionKey,
+                promptLength: systemPromptToAdd.length,
+            });
+            logger.debug("Injected prompt content", { systemPromptToAdd });
+            return Object.assign(Object.assign({}, nextParams), { prompt: messagesWithMemory });
+        },
+        async wrapGenerate({ doGenerate, params }) {
+            var _a, _b;
+            let result;
+            try {
+                result = await doGenerate();
             }
+            catch (err) {
+                logger.error("doGenerate failed", err);
+                logger.error("doGenerate params.prompt", JSON.stringify((_a = params.prompt) === null || _a === void 0 ? void 0 : _a.map((m) => ({ role: m.role, contentType: typeof m.content, contentLength: Array.isArray(m.content) ? m.content.length : undefined })), null, 2));
+                throw err;
+            }
+            if (userId && sessionId) {
+                const sessionKey = `${userId}:${projectId}:${sessionId}`;
+                const promptMeta = sessionPromptMetadata.get(sessionKey);
+                const messagesInput = params.messages || params.prompt || [];
+                const resultMessages = (_b = result === null || result === void 0 ? void 0 : result.response) === null || _b === void 0 ? void 0 : _b.messages;
+                const assistantMessage = (result === null || result === void 0 ? void 0 : result.text)
+                    ? [{ role: "assistant", content: [{ type: "text", text: result.text }] }]
+                    : [];
+                const finalMessages = Array.isArray(resultMessages) && resultMessages.length > 0
+                    ? resultMessages
+                    : [...messagesInput, ...assistantMessage];
+                logConversation(Object.assign({ userId,
+                    projectId,
+                    sessionId, messages: finalMessages, modelId, usage: result.usage }, (promptMeta && {
+                    promptSlug: promptMeta.promptSlug,
+                    promptVersion: promptMeta.promptVersion,
+                    promptId: promptMeta.promptId,
+                }))).then(() => triggerProcessing(userId, projectId, sessionId));
+            }
+            return result;
+        },
+        async wrapStream({ doStream, params }) {
+            var _a;
+            let result;
+            try {
+                logger.debug("Starting doStream with params", JSON.stringify(params, null, 2));
+                result = await doStream();
+            }
+            catch (err) {
+                console.log(err.cause);
+                console.log(err.stack);
+                logger.error("doStream failed", err);
+                throw err;
+            }
+            if (userId && sessionId) {
+                const sessionKey = `${userId}:${projectId}:${sessionId}`;
+                const promptMeta = sessionPromptMetadata.get(sessionKey);
+                const messagesInput = params.messages || params.prompt || [];
+                const resultMessages = (_a = result === null || result === void 0 ? void 0 : result.response) === null || _a === void 0 ? void 0 : _a.messages;
+                const finalMessages = Array.isArray(resultMessages) && resultMessages.length > 0
+                    ? resultMessages
+                    : messagesInput;
+                let streamUsage;
+                let accumulatedText = '';
+                const originalStream = result.stream;
+                const transformStream = new TransformStream({
+                    transform(chunk, controller) {
+                        if (chunk.type === 'text-delta') {
+                            accumulatedText += chunk.delta;
+                        }
+                        if (chunk.type === 'finish' && chunk.usage) {
+                            streamUsage = chunk.usage;
+                        }
+                        controller.enqueue(chunk);
+                    },
+                    flush() {
+                        const allMessages = accumulatedText
+                            ? [...finalMessages, { role: "assistant", content: [{ type: "text", text: accumulatedText }] }]
+                            : finalMessages;
+                        logConversation(Object.assign({ userId,
+                            projectId,
+                            sessionId, messages: allMessages, modelId, usage: streamUsage }, (promptMeta && {
+                            promptSlug: promptMeta.promptSlug,
+                            promptVersion: promptMeta.promptVersion,
+                            promptId: promptMeta.promptId,
+                        }))).then(() => triggerProcessing(userId, projectId, sessionId));
+                    }
+                });
+                result.stream = originalStream.pipeThrough(transformStream);
+            }
+            return result;
+        }
+    });
+    const resolveModel = (originalModel, gatewaySlug) => {
+        if (!gatewaySlug || !clConfig.providerFactory) {
+            if (gatewaySlug && !clConfig.providerFactory) {
+                logger.warn("Gateway config found but no providerFactory provided");
+            }
+            return originalModel;
+        }
+        try {
+            const gatewayURL = `${baseUrl}/api/cognitive/gateway/${gatewaySlug}`;
+            const modelId = originalModel.modelId || 'default';
+            const rawModel = clConfig.providerFactory(gatewayURL)(modelId);
+            const session = originalModel[SESSION_KEY];
+            if (!session)
+                return rawModel;
+            const wrapped = (0, ai_1.wrapLanguageModel)({
+                model: rawModel,
+                middleware: buildMiddleware(session.userId, session.projectId, session.sessionId, modelId),
+            });
+            wrapped[SESSION_KEY] = session;
+            return wrapped;
+        }
+        catch (err) {
+            logger.error("Failed to create gateway model, falling back to original", err);
+            return originalModel;
+        }
+    };
+    const clWrapper = (modelId, settings, providerOptions) => {
+        // Pass provider options through to the underlying provider
+        const model = (providerOptions
+            ? provider(modelId, providerOptions)
+            : provider(modelId));
+        const userId = settings === null || settings === void 0 ? void 0 : settings.userId;
+        const projectId = (settings === null || settings === void 0 ? void 0 : settings.projectId) || clConfig.projectId || "default";
+        const sessionId = settings === null || settings === void 0 ? void 0 : settings.sessionId;
+        const sessionMissing = !!userId && !sessionId;
+        if (sessionMissing) {
+            logger.warn("sessionId is required to log and process memories; skipping logging until provided.");
+        }
+        const wrappedModel = (0, ai_1.wrapLanguageModel)({
+            model: model,
+            middleware: buildMiddleware(userId, projectId, sessionId, modelId),
+        });
+        // Track session settings on the model for use in cl.streamText/cl.generateText
+        if (userId && sessionId) {
+            wrappedModel[SESSION_KEY] = { userId, projectId, sessionId };
+        }
+        return wrappedModel;
+    };
+    const clStreamText = async (options) => {
+        const { prompt: promptConfig } = options, rest = __rest(options, ["prompt"]);
+        // Resolve and interpolate prompt
+        const resolved = await resolvePrompt(promptConfig.slug);
+        const system = promptConfig.variables
+            ? interpolateTemplate(resolved.content, promptConfig.variables)
+            : resolved.content;
+        // Store prompt metadata for the session (read by middleware during logging)
+        const session = options.model[SESSION_KEY];
+        if (session) {
+            const sessionKey = `${session.userId}:${session.projectId}:${session.sessionId}`;
+            sessionPromptMetadata.set(sessionKey, {
+                promptSlug: resolved.slug,
+                promptVersion: resolved.version,
+                promptId: resolved.promptId,
+            });
+        }
+        logger.info("cl.streamText called", {
+            slug: promptConfig.slug,
+            version: resolved.version,
+            systemLength: system.length,
+        });
+        const model = resolveModel(options.model, resolved.gatewaySlug);
+        return (0, ai_1.streamText)(Object.assign(Object.assign({}, rest), { model, system }));
+    };
+    const clGenerateText = async (options) => {
+        const { prompt: promptConfig } = options, rest = __rest(options, ["prompt"]);
+        // Resolve and interpolate prompt
+        const resolved = await resolvePrompt(promptConfig.slug);
+        const system = promptConfig.variables
+            ? interpolateTemplate(resolved.content, promptConfig.variables)
+            : resolved.content;
+        // Store prompt metadata for the session (read by middleware during logging)
+        const session = options.model[SESSION_KEY];
+        if (session) {
+            const sessionKey = `${session.userId}:${session.projectId}:${session.sessionId}`;
+            sessionPromptMetadata.set(sessionKey, {
+                promptSlug: resolved.slug,
+                promptVersion: resolved.version,
+                promptId: resolved.promptId,
+            });
+        }
+        logger.info("cl.generateText called", {
+            slug: promptConfig.slug,
+            version: resolved.version,
+            systemLength: system.length,
         });
+        const model = resolveModel(options.model, resolved.gatewaySlug);
+        return (0, ai_1.generateText)(Object.assign(Object.assign({}, rest), { model, system }));
     };
+    // Return the model wrapper function with streamText/generateText attached
+    return Object.assign(clWrapper, {
+        streamText: clStreamText,
+        generateText: clGenerateText,
+        resolvePrompt,
+        logConversation,
+        triggerProcessing,
+        clearPromptCache: () => promptCache.clear(),
+        clearSessionCache: (sessionKey) => {
+            if (sessionKey) {
+                sessionSnapshots.delete(sessionKey);
+                sessionPromptMetadata.delete(sessionKey);
+            }
+            else {
+                sessionSnapshots.clear();
+                sessionPromptMetadata.clear();
+            }
+        },
+    });
 }

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@kognitivedev/vercel-ai-provider",
-    "version": "0.1.4",
+    "version": "0.1.6",
     "main": "dist/index.js",
     "types": "dist/index.d.ts",
     "publishConfig": {
@@ -9,14 +9,16 @@
     "scripts": {
         "build": "tsc",
         "dev": "tsc -w",
+        "test": "vitest run",
         "prepublishOnly": "npm run build"
     },
     "peerDependencies": {
-        "ai": "^4.0.0 || ^5.0.0"
+        "ai": "^5.0.0 || ^6.0.0"
     },
     "devDependencies": {
         "typescript": "^5.0.0",
-        "ai": "latest",
-        "@types/node": "^20.0.0"
+        "ai": "^6.0.0",
+        "@types/node": "^20.0.0",
+        "vitest": "^3.0.0"
     }
 }

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" },
+        ]);
+    });
+});