@kognitivedev/vercel-ai-provider 0.1.7 → 0.1.9

package/src/index.ts

@@ -4,6 +4,7 @@ import {
     generateText as aiGenerateText,
     type LanguageModel,
 } from "ai";
+import { randomUUID } from "crypto";
 
 /**
  * Log levels for controlling verbosity of CognitiveLayer logging.
@@ -21,6 +22,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,
@@ -111,12 +123,33 @@ 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 & {
     streamText: (options: CLStreamTextOptions) => Promise<ReturnType<typeof aiStreamText>>;
     generateText: (options: CLGenerateTextOptions) => ReturnType<typeof aiGenerateText>;
-    resolvePrompt: (slug: string) => Promise<CachedPrompt>;
+    resolvePrompt: (slug: string, userId?: string) => Promise<CachedPrompt>;
     logConversation: (payload: LogConversationPayload) => Promise<void>;
     triggerProcessing: (userId: string, projectId: string, sessionId: string) => void;
     clearPromptCache: () => void;
@@ -136,6 +169,111 @@ 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(" ");
+}
+
+/**
+ * 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 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;
+}
+
 /**
  * Interpolate {{variable}} placeholders in a template string.
  * Unmatched variables are left as-is.
@@ -194,18 +332,44 @@ export function createCognitiveLayer(config: {
     // Prompt cache: slug → CachedPrompt
     const promptCache = new Map<string, CachedPrompt>();
 
-    const resolvePrompt = async (slug: string): Promise<CachedPrompt> => {
-        const cached = promptCache.get(slug);
+    const resolvePrompt = async (slug: string, userId?: string): Promise<CachedPrompt> => {
+        const cacheKey = userId ? `${slug}:${userId}` : slug;
+        const cached = promptCache.get(cacheKey);
         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)}`, {
+        const url = new URL(`${baseUrl}/api/cognitive/prompt`);
+        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}`);
         }
 
@@ -218,7 +382,15 @@ export function createCognitiveLayer(config: {
             fetchedAt: Date.now(),
             gatewaySlug: data.gatewaySlug,
         };
-        promptCache.set(slug, entry);
+        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;
     };
@@ -301,9 +473,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 || "";
@@ -337,7 +525,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);
                     }
@@ -370,6 +566,7 @@ ${userContextBlock || "None"}
         },
 
         async wrapGenerate({ doGenerate, params }: { doGenerate: any; params: any }) {
+            const startedAt = new Date();
             let result;
             try {
                 result = await doGenerate();
@@ -380,17 +577,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,
@@ -404,12 +624,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));
@@ -426,7 +659,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
@@ -434,6 +667,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({
@@ -444,14 +680,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,
@@ -463,7 +749,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);
                     }
                 });
 
@@ -528,8 +826,10 @@ ${userContextBlock || "None"}
         }) as LanguageModel;
 
         // Track session settings on the model for use in cl.streamText/cl.generateText
-        if (isValidId(userId) && isValidId(sessionId)) {
-            (wrappedModel as any)[SESSION_KEY] = { userId, projectId, sessionId };
+        // Always store if userId is valid — sessionId may be missing but userId is still
+        // needed for prompt resolution (e.g. A/B test assignment)
+        if (isValidId(userId)) {
+            (wrappedModel as any)[SESSION_KEY] = { userId, projectId, sessionId: isValidId(sessionId) ? sessionId : undefined };
         }
 
         return wrappedModel;
@@ -538,61 +838,89 @@ ${userContextBlock || "None"}
     const clStreamText = async (options: CLStreamTextOptions) => {
         const { prompt: promptConfig, ...rest } = options;
 
-        // 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 as any)[SESSION_KEY] as { userId: string; projectId: string; sessionId: string } | undefined;
-        if (session) {
-            const sessionKey = `${session.userId}:${session.projectId}:${session.sessionId}`;
-            sessionPromptMetadata.set(sessionKey, {
-                promptSlug: resolved.slug,
-                promptVersion: resolved.version,
-                promptId: resolved.promptId,
-            });
+        const session = (options.model as any)[SESSION_KEY] as { userId: string; projectId: string; sessionId?: string } | undefined;
+
+        // Resolve and interpolate prompt (graceful fallback on failure)
+        let resolved: CachedPrompt | null = null;
+        try {
+            resolved = await resolvePrompt(promptConfig.slug, session?.userId);
+        } catch (err) {
+            logger.warn(`Failed to resolve prompt "${promptConfig.slug}", streaming without system prompt.`, err);
         }
 
-        logger.info("cl.streamText called", {
-            slug: promptConfig.slug,
-            version: resolved.version,
-            systemLength: system.length,
-        });
+        let system: string | undefined;
+        if (resolved) {
+            system = promptConfig.variables
+                ? interpolateTemplate(resolved.content, promptConfig.variables)
+                : resolved.content;
+
+            // Store prompt metadata for the session (read by middleware during logging)
+            if (session?.sessionId) {
+                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,
+            });
+        } else {
+            logger.info("cl.streamText called without resolved prompt", {
+                slug: promptConfig.slug,
+            });
+        }
 
-        const model = resolveModel(options.model, resolved.gatewaySlug);
-        return aiStreamText({ ...rest, model, system } as any);
+        const model = resolveModel(options.model, resolved?.gatewaySlug);
+        return aiStreamText({ ...rest, model, ...(system && { system }) } as any);
     };
 
     const clGenerateText = async (options: CLGenerateTextOptions) => {
         const { prompt: promptConfig, ...rest } = options;
 
-        // 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 as any)[SESSION_KEY] as { userId: string; projectId: string; sessionId: string } | undefined;
-        if (session) {
-            const sessionKey = `${session.userId}:${session.projectId}:${session.sessionId}`;
-            sessionPromptMetadata.set(sessionKey, {
-                promptSlug: resolved.slug,
-                promptVersion: resolved.version,
-                promptId: resolved.promptId,
-            });
+        const session = (options.model as any)[SESSION_KEY] as { userId: string; projectId: string; sessionId?: string } | undefined;
+
+        // Resolve and interpolate prompt (graceful fallback on failure)
+        let resolved: CachedPrompt | null = null;
+        try {
+            resolved = await resolvePrompt(promptConfig.slug, session?.userId);
+        } catch (err) {
+            logger.warn(`Failed to resolve prompt "${promptConfig.slug}", generating without system prompt.`, err);
         }
 
-        logger.info("cl.generateText called", {
-            slug: promptConfig.slug,
-            version: resolved.version,
-            systemLength: system.length,
-        });
+        let system: string | undefined;
+        if (resolved) {
+            system = promptConfig.variables
+                ? interpolateTemplate(resolved.content, promptConfig.variables)
+                : resolved.content;
+
+            // Store prompt metadata for the session (read by middleware during logging)
+            if (session?.sessionId) {
+                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,
+            });
+        } else {
+            logger.info("cl.generateText called without resolved prompt", {
+                slug: promptConfig.slug,
+            });
+        }
 
-        const model = resolveModel(options.model, resolved.gatewaySlug);
-        return aiGenerateText({ ...rest, model, system } as any);
+        const model = resolveModel(options.model, resolved?.gatewaySlug);
+        return aiGenerateText({ ...rest, model, ...(system && { system }) } as any);
     };
 
     // Return the model wrapper function with streamText/generateText attached