@kognitivedev/vercel-ai-provider 0.1.8 → 0.2.0

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

@@ -36,6 +36,110 @@ describe("wrapStream logging", () => {
         );
     });
 
+    it("should capture tool-call chunks and include them in logged conversation", async () => {
+        const mockModel = new MockLanguageModelV3({
+            doStream: async () => ({
+                stream: convertArrayToReadableStream([
+                    { type: "text-start" as const, id: "t1" },
+                    { type: "text-delta" as const, id: "t1", delta: "Let me check" },
+                    { type: "text-end" as const, id: "t1" },
+                    {
+                        type: "tool-call" as const,
+                        toolCallId: "call-1",
+                        toolName: "get_weather",
+                        input: '{"city":"London"}',
+                    },
+                    {
+                        type: "tool-result" as const,
+                        toolCallId: "call-1",
+                        toolName: "get_weather",
+                        result: { temperature: 15, unit: "celsius" },
+                    },
+                    {
+                        type: "finish" as const,
+                        finishReason: {
+                            unified: "tool-calls" as const,
+                            raw: undefined,
+                        },
+                        usage: {
+                            inputTokens: { total: 20, noCache: undefined, cacheRead: undefined, cacheWrite: undefined },
+                            outputTokens: { total: 15, 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: "What's the weather in London?" }],
+        });
+
+        // Fully consume the stream
+        await result.text;
+
+        // 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;
+
+        // Assistant message should contain text + tool-call parts
+        const assistantMsg = messages.find((m: any) => m.role === "assistant");
+        expect(assistantMsg).toBeDefined();
+        expect(assistantMsg.content).toEqual([
+            { type: "text", text: "Let me check" },
+            {
+                type: "tool-call",
+                toolCallId: "call-1",
+                toolName: "get_weather",
+                input: '{"city":"London"}',
+            },
+        ]);
+
+        // Tool results should be in a separate tool message
+        const toolMsg = messages.find((m: any) => m.role === "tool");
+        expect(toolMsg).toBeDefined();
+        expect(toolMsg.content).toEqual([
+            {
+                type: "tool-result",
+                toolCallId: "call-1",
+                toolName: "get_weather",
+                result: { temperature: 15, unit: "celsius" },
+            },
+        ]);
+
+        // Spans should include the tool call with populated previews
+        const spans = logCall!.body.spans;
+        const toolSpan = spans?.find((s: any) => s.spanType === "tool");
+        expect(toolSpan).toBeDefined();
+        expect(toolSpan.toolName).toBe("get_weather");
+        expect(toolSpan.inputPreview).toContain("London");
+        expect(toolSpan.outputPreview).toContain("15");
+    });
+
     it("should include assistant message in logged conversation after streaming", async () => {
         const mockModel = new MockLanguageModelV3({
             doStream: async () => ({

package/src/handlebars.d.ts

@@ -0,0 +1,4 @@
+declare module "handlebars/dist/cjs/handlebars" {
+    import Handlebars from "handlebars";
+    export default Handlebars;
+}

package/src/index.ts

@@ -4,6 +4,9 @@ import {
     generateText as aiGenerateText,
     type LanguageModel,
 } from "ai";
+import { randomUUID } from "crypto";
+export { renderTemplate, type TemplateVariables } from "./template";
+import { renderTemplate } from "./template";
 
 /**
  * Log levels for controlling verbosity of CognitiveLayer logging.
@@ -21,6 +24,17 @@ function isValidId(value: string | undefined | null): value is string {
     return trimmed !== "" && trimmed !== "null" && trimmed !== "undefined";
 }
 
+function maskSecret(secret: string | undefined | null): string {
+    if (!secret) return "missing";
+    if (secret.length <= 8) return `${secret.slice(0, 2)}***`;
+    return `${secret.slice(0, 4)}...${secret.slice(-4)}`;
+}
+
+function previewText(value: string, maxLength = 240): string {
+    if (value.length <= maxLength) return value;
+    return `${value.slice(0, maxLength)}...`;
+}
+
 const LOG_LEVEL_PRIORITY: Record<LogLevel, number> = {
     none: 0,
     error: 1,
@@ -89,7 +103,7 @@ export type CLModelWrapper = (
 
 export interface PromptConfig {
     slug: string;
-    variables?: Record<string, string>;
+    variables?: Record<string, string | boolean>;
 }
 
 export type CLStreamTextOptions = Omit<Parameters<typeof aiStreamText>[0], 'system' | 'prompt'> & {
@@ -111,6 +125,27 @@ export interface LogConversationPayload {
     promptSlug?: string;
     promptVersion?: number;
     promptId?: string;
+    traceId?: string;
+    parentSpanId?: string;
+    requestPreview?: string;
+    responsePreview?: string;
+    state?: "active" | "completed" | "error";
+    startedAt?: string;
+    endedAt?: string;
+    durationMs?: number;
+    metadata?: Record<string, unknown>;
+    spans?: Array<{
+        spanKey: string;
+        parentSpanKey?: string;
+        name: string;
+        spanType: string;
+        status?: "active" | "completed" | "error";
+        inputPreview?: string;
+        outputPreview?: string;
+        toolName?: string;
+        errorMessage?: string;
+        metadata?: Record<string, unknown>;
+    }>;
 }
 
 export type CognitiveLayer = CLModelWrapper & {
@@ -136,12 +171,109 @@ export interface CachedPrompt {
 
 const PROMPT_CACHE_TTL_MS = 60_000; // 1 minute
 
+function getContentText(content: any): string {
+    if (typeof content === "string") return content;
+    if (!Array.isArray(content)) return "";
+
+    return content.map((part) => {
+        if (!part || typeof part !== "object") return "";
+        if (typeof part.text === "string") return part.text;
+        if (part.type === "tool-call" && typeof part.toolName === "string") return `Called ${part.toolName}`;
+        if (part.type === "tool-result") return "Received tool result";
+        return "";
+    }).filter(Boolean).join(" ");
+}
+
 /**
- * Interpolate {{variable}} placeholders in a template string.
- * Unmatched variables are left as-is.
+ * Unwraps V2/V3 ToolResultOutput discriminated union to a displayable value.
+ * Stream ToolResult uses plain `result` (passthrough), while prompt ToolResultPart
+ * uses `output` with a discriminated union: text, json, error-text, error-json, content, execution-denied.
  */
-function interpolateTemplate(content: string, variables: Record<string, string>): string {
-    return content.replace(/\{\{(\w+)\}\}/g, (_, key) => variables[key] ?? `{{${key}}}`);
+function extractOutputValue(raw: unknown): unknown {
+    if (raw == null) return raw;
+    if (typeof raw !== 'object') return raw;
+    const obj = raw as Record<string, unknown>;
+    if (typeof obj.type !== 'string') return raw;
+    switch (obj.type) {
+        case 'text':
+        case 'json':
+        case 'error-text':
+        case 'error-json':
+        case 'content':
+            return obj.value;
+        case 'execution-denied':
+            return `Execution denied: ${obj.reason ?? 'unknown'}`;
+        default:
+            return raw;
+    }
+}
+
+function buildTracePreviews(messages: any[]): { requestPreview: string; responsePreview: string } {
+    const request = [...messages].reverse().find((message) => message?.role === "user");
+    const response = [...messages].reverse().find((message) => message?.role === "assistant");
+
+    return {
+        requestPreview: request ? getContentText(request.content).slice(0, 220) : "No request captured",
+        responsePreview: response ? getContentText(response.content).slice(0, 240) : "No response captured",
+    };
+}
+
+function buildTraceSpansFromMessages(messages: any[]): Array<{
+    spanKey: string;
+    parentSpanKey?: string;
+    name: string;
+    spanType: string;
+    status?: "active" | "completed" | "error";
+    inputPreview?: string;
+    outputPreview?: string;
+    toolName?: string;
+    errorMessage?: string;
+    metadata?: Record<string, unknown>;
+}> {
+    const resultMap = new Map<string, unknown>();
+
+    for (const message of messages) {
+        if (!Array.isArray(message?.content)) continue;
+        for (const part of message.content) {
+            if (part?.type === "tool-result" && typeof part.toolCallId === "string") {
+                resultMap.set(part.toolCallId, part.result ?? part.output);
+            }
+        }
+    }
+
+    const spans: Array<{
+        spanKey: string;
+        parentSpanKey?: string;
+        name: string;
+        spanType: string;
+        status?: "active" | "completed" | "error";
+        inputPreview?: string;
+        outputPreview?: string;
+        toolName?: string;
+        errorMessage?: string;
+        metadata?: Record<string, unknown>;
+    }> = [];
+
+    for (const message of messages) {
+        if (!Array.isArray(message?.content)) continue;
+        for (const part of message.content) {
+            if (part?.type === "tool-call" && typeof part.toolCallId === "string") {
+                const result = resultMap.get(part.toolCallId);
+                spans.push({
+                    spanKey: part.toolCallId,
+                    parentSpanKey: "root",
+                    name: typeof part.toolName === "string" ? part.toolName : "tool",
+                    spanType: "tool",
+                    status: "completed",
+                    inputPreview: JSON.stringify(part.input ?? {}).slice(0, 220),
+                    outputPreview: result != null ? JSON.stringify(extractOutputValue(result)).slice(0, 220) : "No tool result captured",
+                    toolName: typeof part.toolName === "string" ? part.toolName : undefined,
+                });
+            }
+        }
+    }
+
+    return spans;
 }
 
 // Session-scoped snapshot cache: sessionKey → formatted memory block
@@ -206,11 +338,32 @@ export function createCognitiveLayer(config: {
         url.searchParams.set("slug", slug);
         if (userId) url.searchParams.set("userId", userId);
 
+        logger.debug("Resolving prompt from backend", {
+            slug,
+            userId,
+            url: url.toString(),
+            baseUrl,
+            apiKeyHint: maskSecret(clConfig.apiKey),
+        });
+
         const res = await fetch(url.toString(), {
             headers: { "Authorization": `Bearer ${clConfig.apiKey}` },
         });
+        logger.debug("Prompt resolve response received", {
+            slug,
+            userId,
+            status: res.status,
+            ok: res.ok,
+            contentType: res.headers.get("content-type"),
+        });
         if (!res.ok) {
             const body = await res.text();
+            logger.debug("Prompt resolve response body preview", {
+                slug,
+                userId,
+                status: res.status,
+                bodyPreview: previewText(body),
+            });
             throw new Error(`Failed to resolve prompt "${slug}": ${res.status} ${body}`);
         }
 
@@ -224,6 +377,14 @@ export function createCognitiveLayer(config: {
             gatewaySlug: data.gatewaySlug,
         };
         promptCache.set(cacheKey, entry);
+        logger.debug("Prompt resolved payload", {
+            slug,
+            resolvedSlug: entry.slug,
+            version: entry.version,
+            promptId: entry.promptId,
+            contentLength: entry.content.length,
+            gatewaySlug: entry.gatewaySlug ?? null,
+        });
         logger.info("Prompt resolved", { slug, version: entry.version });
         return entry;
     };
@@ -306,9 +467,25 @@ export function createCognitiveLayer(config: {
             if (systemPromptToAdd === undefined) {
                 try {
                     const url = `${baseUrl}/api/cognitive/snapshot?userId=${userId}`;
+                    logger.debug("Fetching snapshot from backend", {
+                        userId,
+                        projectId,
+                        sessionId,
+                        url,
+                        baseUrl,
+                        apiKeyHint: maskSecret(clConfig.apiKey),
+                    });
                     const res = await fetch(url, {
                         headers: { "Authorization": `Bearer ${clConfig.apiKey}` },
                     });
+                    logger.debug("Snapshot response received", {
+                        userId,
+                        projectId,
+                        sessionId,
+                        status: res.status,
+                        ok: res.ok,
+                        contentType: res.headers.get("content-type"),
+                    });
                     if (res.ok) {
                         const data = await res.json();
                         const systemBlock = data.systemBlock || "";
@@ -342,7 +519,15 @@ ${userContextBlock || "None"}
                             rawData: data,
                         });
                     } else {
+                        const body = await res.text();
                         logger.warn("Snapshot fetch failed", { status: res.status });
+                        logger.debug("Snapshot response body preview", {
+                            userId,
+                            projectId,
+                            sessionId,
+                            status: res.status,
+                            bodyPreview: previewText(body),
+                        });
                         systemPromptToAdd = "";
                         sessionSnapshots.set(sessionKey, systemPromptToAdd);
                     }
@@ -375,6 +560,7 @@ ${userContextBlock || "None"}
         },
 
         async wrapGenerate({ doGenerate, params }: { doGenerate: any; params: any }) {
+            const startedAt = new Date();
             let result;
             try {
                 result = await doGenerate();
@@ -385,17 +571,40 @@ ${userContextBlock || "None"}
             }
 
             if (isValidId(userId) && isValidId(sessionId)) {
+                const endedAt = new Date();
                 const sessionKey = `${userId}:${projectId}:${sessionId}`;
                 const promptMeta = sessionPromptMetadata.get(sessionKey);
 
-                const messagesInput = (params as any).messages || (params as any).prompt || [];
-                const resultMessages = (result as any)?.response?.messages;
-                const assistantMessage = (result as any)?.text
-                    ? [{ role: "assistant", content: [{ type: "text", text: (result as any).text }] }]
+                const messagesInput = (params as any).prompt || (params as any).messages || [];
+
+                // Build assistant message from result.content (V2/V3 GenerateResult)
+                const resultContent = Array.isArray(result?.content) ? result.content : [];
+                const assistantParts: any[] = [];
+                for (const part of resultContent) {
+                    if (part?.type === 'text') {
+                        assistantParts.push({ type: 'text', text: part.text });
+                    } else if (part?.type === 'tool-call') {
+                        assistantParts.push({
+                            type: 'tool-call',
+                            toolCallId: part.toolCallId,
+                            toolName: part.toolName,
+                            input: part.input,
+                        });
+                    } else if (part?.type === 'tool-result') {
+                        assistantParts.push({
+                            type: 'tool-result',
+                            toolCallId: part.toolCallId,
+                            toolName: part.toolName,
+                            result: part.result,
+                        });
+                    }
+                }
+                const assistantMessage = assistantParts.length > 0
+                    ? [{ role: "assistant", content: assistantParts }]
                     : [];
-                const finalMessages = Array.isArray(resultMessages) && resultMessages.length > 0
-                    ? resultMessages
-                    : [...messagesInput, ...assistantMessage];
+                const finalMessages = [...messagesInput, ...assistantMessage];
+                const { requestPreview, responsePreview } = buildTracePreviews(finalMessages);
+                const spans = buildTraceSpansFromMessages(finalMessages);
 
                 logConversation({
                     userId,
@@ -409,12 +618,25 @@ ${userContextBlock || "None"}
                         promptVersion: promptMeta.promptVersion,
                         promptId: promptMeta.promptId,
                     }),
+                    traceId: randomUUID(),
+                    requestPreview,
+                    responsePreview,
+                    state: "completed",
+                    startedAt: startedAt.toISOString(),
+                    endedAt: endedAt.toISOString(),
+                    durationMs: endedAt.getTime() - startedAt.getTime(),
+                    metadata: {
+                        appId: clConfig.appId,
+                    },
+                    spans,
                 }).then(() => triggerProcessing(userId, projectId, sessionId));
             }
 
             return result;
         },
         async wrapStream({ doStream, params }: { doStream: any; params: any }) {
+            const startedAt = new Date();
+            const traceId = randomUUID();
             let result;
             try {
                 logger.debug("Starting doStream with params", JSON.stringify(params, null, 2));
@@ -431,7 +653,7 @@ ${userContextBlock || "None"}
                 const sessionKey = `${userId}:${projectId}:${sessionId}`;
                 const promptMeta = sessionPromptMetadata.get(sessionKey);
 
-                const messagesInput = (params as any).messages || (params as any).prompt || [];
+                const messagesInput = (params as any).prompt || (params as any).messages || [];
                 const resultMessages = (result as any)?.response?.messages;
                 const finalMessages = Array.isArray(resultMessages) && resultMessages.length > 0
                     ? resultMessages
@@ -439,6 +661,9 @@ ${userContextBlock || "None"}
 
                 let streamUsage: Record<string, unknown> | undefined;
                 let accumulatedText = '';
+                const toolCallInputs = new Map<string, { toolName: string; chunks: string[] }>();
+                const completedToolCalls: any[] = [];
+                const completedToolResults: any[] = [];
 
                 const originalStream = result.stream;
                 const transformStream = new TransformStream({
@@ -449,14 +674,64 @@ ${userContextBlock || "None"}
                         if (chunk.type === 'finish' && chunk.usage) {
                             streamUsage = chunk.usage;
                         }
+                        // Capture tool-call stream chunks (V2/V3 shared types)
+                        if (chunk.type === 'tool-input-start') {
+                            toolCallInputs.set(chunk.id, { toolName: chunk.toolName, chunks: [] });
+                        }
+                        if (chunk.type === 'tool-input-delta') {
+                            const entry = toolCallInputs.get(chunk.id);
+                            if (entry) entry.chunks.push(chunk.delta);
+                        }
+                        if (chunk.type === 'tool-call') {
+                            completedToolCalls.push({
+                                type: 'tool-call',
+                                toolCallId: chunk.toolCallId,
+                                toolName: chunk.toolName,
+                                input: chunk.input,
+                            });
+                        }
+                        if (chunk.type === 'tool-result') {
+                            completedToolResults.push({
+                                type: 'tool-result',
+                                toolCallId: chunk.toolCallId,
+                                toolName: chunk.toolName,
+                                result: chunk.result,
+                            });
+                        }
                         controller.enqueue(chunk);
                     },
-                    flush() {
-                        const allMessages = accumulatedText
-                            ? [...finalMessages, { role: "assistant", content: [{ type: "text", text: accumulatedText }] }]
+                    async flush() {
+                        const endedAt = new Date();
+
+                        // Finalize any tool calls from incremental input chunks
+                        for (const [id, entry] of toolCallInputs) {
+                            // Only add if not already captured via a tool-call chunk
+                            if (!completedToolCalls.some((tc: any) => tc.toolCallId === id)) {
+                                completedToolCalls.push({
+                                    type: 'tool-call',
+                                    toolCallId: id,
+                                    toolName: entry.toolName,
+                                    input: entry.chunks.join(''),
+                                });
+                            }
+                        }
+
+                        const assistantParts: any[] = [];
+                        if (accumulatedText) assistantParts.push({ type: "text", text: accumulatedText });
+                        for (const tc of completedToolCalls) assistantParts.push(tc);
+
+                        const allMessages = assistantParts.length > 0
+                            ? [...finalMessages, { role: "assistant", content: assistantParts }]
                             : finalMessages;
 
-                        logConversation({
+                        if (completedToolResults.length > 0) {
+                            allMessages.push({ role: "tool", content: completedToolResults });
+                        }
+
+                        const { requestPreview, responsePreview } = buildTracePreviews(allMessages);
+                        const spans = buildTraceSpansFromMessages(allMessages);
+
+                        await logConversation({
                             userId,
                             projectId,
                             sessionId,
@@ -468,7 +743,19 @@ ${userContextBlock || "None"}
                                 promptVersion: promptMeta.promptVersion,
                                 promptId: promptMeta.promptId,
                             }),
-                        }).then(() => triggerProcessing(userId, projectId, sessionId));
+                            traceId,
+                            requestPreview,
+                            responsePreview,
+                            state: "completed",
+                            startedAt: startedAt.toISOString(),
+                            endedAt: endedAt.toISOString(),
+                            durationMs: endedAt.getTime() - startedAt.getTime(),
+                            metadata: {
+                                appId: clConfig.appId,
+                            },
+                            spans,
+                        });
+                        triggerProcessing(userId, projectId, sessionId);
                     }
                 });
 
@@ -558,7 +845,7 @@ ${userContextBlock || "None"}
         let system: string | undefined;
         if (resolved) {
             system = promptConfig.variables
-                ? interpolateTemplate(resolved.content, promptConfig.variables)
+                ? renderTemplate(resolved.content, promptConfig.variables)
                 : resolved.content;
 
             // Store prompt metadata for the session (read by middleware during logging)
@@ -602,7 +889,7 @@ ${userContextBlock || "None"}
         let system: string | undefined;
         if (resolved) {
             system = promptConfig.variables
-                ? interpolateTemplate(resolved.content, promptConfig.variables)
+                ? renderTemplate(resolved.content, promptConfig.variables)
                 : resolved.content;
 
             // Store prompt metadata for the session (read by middleware during logging)

package/src/template.ts

@@ -0,0 +1,10 @@
+// Use the pre-built dist to avoid `require.extensions` warning in webpack/Next.js
+// eslint-disable-next-line @typescript-eslint/no-require-imports
+import Handlebars from "handlebars/dist/cjs/handlebars";
+
+export type TemplateVariables = Record<string, string | boolean>;
+
+export function renderTemplate(template: string, variables: TemplateVariables): string {
+    const compiled = Handlebars.compile(template, { noEscape: true });
+    return compiled(variables);
+}