@plumbus/core 0.2.8 → 0.3.0

package/dist/ai/ai-service.js

@@ -3,8 +3,10 @@
 // Integrates: provider adapter, prompt registry, validation, cost tracking, security, RAG, explainability
 import { z } from 'zod';
 import { calculateModelCost } from './model-pricing.js';
+import { AIIncompleteOutputError, AIRefusalError } from './refusal.js';
 import { checkPromptSecurity } from './security.js';
-import { generateWithValidation } from './validation.js';
+import { AIValidationError, generateWithValidation, } from './validation.js';
+import { zodToProviderJsonSchema } from './zod-to-provider-schema.js';
 /** Convenience: create config from a single provider (backward compat) */
 export function singleProviderConfig(provider, rest) {
     return {
@@ -14,7 +16,8 @@ export function singleProviderConfig(provider, rest) {
     };
 }
 export function createAIService(config) {
-    const { providers, defaultProvider, promptRegistry, costTracker, ragPipeline, explainability, security, } = config;
+    const { providers, defaultProvider, promptRegistry, costTracker, ragPipeline, explainability, security, onAICostRecorded, } = config;
+    const responseSchemaCache = new WeakMap();
     function resolveProvider(providerName) {
         const name = providerName ?? defaultProvider;
         const adapter = providers[name];
@@ -24,6 +27,96 @@ export function createAIService(config) {
         }
         return adapter;
     }
+    function getSingleTextFieldName(schema) {
+        const shape = schema instanceof z.ZodObject ? schema.shape : undefined;
+        if (!shape)
+            return undefined;
+        const keys = Object.keys(shape);
+        const firstKey = keys[0];
+        if (keys.length === 1 && firstKey && shape[firstKey] instanceof z.ZodString) {
+            return firstKey;
+        }
+        return undefined;
+    }
+    function getStrictResponseSchema(promptName, promptDef) {
+        if (!promptDef)
+            return undefined;
+        if (!config.enableStrictStructuredOutputs) {
+            if (promptDef.requireStrictStructuredOutputs) {
+                throw new Error(`Prompt "${promptName}" requires provider-side structured outputs, but enableStrictStructuredOutputs is disabled`);
+            }
+            return undefined;
+        }
+        if (promptDef.disableStrictStructuredOutputs) {
+            if (promptDef.requireStrictStructuredOutputs) {
+                throw new Error(`Prompt "${promptName}" cannot require strict structured outputs while disableStrictStructuredOutputs is set`);
+            }
+            return undefined;
+        }
+        if (getSingleTextFieldName(promptDef.output)) {
+            if (promptDef.requireStrictStructuredOutputs) {
+                throw new Error(`Prompt "${promptName}" requires provider-side structured outputs, but single-string output schemas use text mode`);
+            }
+            return undefined;
+        }
+        const cached = responseSchemaCache.get(promptDef.output);
+        if (cached)
+            return cached;
+        const converted = zodToProviderJsonSchema(promptDef.output, { promptName });
+        responseSchemaCache.set(promptDef.output, converted.schema);
+        return converted.schema;
+    }
+    function getTypedAIErrorStatus(err) {
+        if (err instanceof AIRefusalError)
+            return 'refused';
+        if (err instanceof AIIncompleteOutputError)
+            return 'incomplete';
+        return 'failed';
+    }
+    function getTypedAIErrorUsage(err) {
+        if (err instanceof AIRefusalError || err instanceof AIIncompleteOutputError) {
+            return err.usage;
+        }
+        return undefined;
+    }
+    /**
+     * Centralized cost recorder: updates the in-memory tracker and fires the
+     * optional `onAICostRecorded` hook. Hook errors are swallowed so a
+     * misbehaving hook can never break the AI call. The `status` default and
+     * id/timestamp generation mirror what `createCostTracker` does internally
+     * so the hook gets a fully-formed `AICostRecord`.
+     */
+    async function recordCost(entry, costContext) {
+        if (costTracker) {
+            costTracker.record(entry);
+        }
+        if (onAICostRecorded) {
+            const finalEntry = {
+                ...entry,
+                status: entry.status ?? 'success',
+                id: crypto.randomUUID(),
+                timestamp: new Date(),
+            };
+            try {
+                await onAICostRecorded(finalEntry, costContext);
+            }
+            catch (hookErr) {
+                const msg = hookErr instanceof Error ? hookErr.message : String(hookErr);
+                console.error(`[plumbus:ai] onAICostRecorded hook threw: ${msg}`);
+            }
+        }
+    }
+    function sumUsage(a, b) {
+        const left = a ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 };
+        const right = b ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 };
+        return {
+            inputTokens: left.inputTokens + right.inputTokens,
+            outputTokens: left.outputTokens + right.outputTokens,
+            totalTokens: left.totalTokens + right.totalTokens,
+            cachedInputTokens: (left.cachedInputTokens ?? 0) + (right.cachedInputTokens ?? 0),
+            cacheWriteTokens: (left.cacheWriteTokens ?? 0) + (right.cacheWriteTokens ?? 0),
+        };
+    }
     function checkBudget(estimatedTokens) {
         if (!costTracker)
             return;
@@ -38,32 +131,64 @@ export function createAIService(config) {
     function buildPromptText(promptName, input) {
         if (!promptRegistry) {
             // No registry — treat promptName as raw prompt text
-            return { text: promptName };
+            return { text: promptName, substitutedKeys: new Set() };
         }
         const def = promptRegistry.get(promptName);
         // Validate input against prompt input schema
         def.input.parse(input);
-        // Resolve template: replace {{key}} with input values
-        let text = promptName;
-        if (def.description) {
-            text = def.description;
-        }
-        // Simple template resolution: replace {{key}} placeholders
-        for (const [key, value] of Object.entries(input)) {
-            text = text.replaceAll(`{{${key}}}`, String(value));
-        }
+        // Simple template resolution: replace top-level {{key}} placeholders.
+        const substitutedKeys = new Set();
+        const renderTemplate = (template) => {
+            let text = template;
+            for (const [key, value] of Object.entries(input)) {
+                const placeholder = `{{${key}}}`;
+                if (text.includes(placeholder)) {
+                    substitutedKeys.add(key);
+                    text = text.replaceAll(placeholder, String(value));
+                }
+            }
+            return text;
+        };
+        const system = def.system ? renderTemplate(def.system) : undefined;
+        const text = renderTemplate(def.description ?? promptName);
         // Resolution chain: config/env override → prompt-level → defaults
         // Look up override by prompt name (dots replaced with underscores for env var matching)
         const overrideKey = promptName.toLowerCase().replaceAll('.', '_');
         const override = config.promptOverrides?.[overrideKey] ?? config.promptOverrides?.[promptName];
         return {
+            system,
             text,
             model: override?.model ?? def.model?.name ?? config.defaultModel,
             provider: override?.provider ?? def.model?.provider,
             temperature: override?.temperature ?? def.model?.temperature,
             maxTokens: override?.maxTokens ?? def.model?.maxTokens,
+            appendUnsubstitutedInput: def.appendUnsubstitutedInput,
+            substitutedKeys,
         };
     }
+    /**
+     * Builds the final prompt string the provider sees, given a rendered
+     * `promptInfo` (with placeholders already substituted) and the original
+     * input object. Any keys that were inlined as `{{key}}` in the description
+     * are stripped from the appended `Input: ...` JSON blob to prevent the
+     * same value from being sent twice in one HTTP body.
+     */
+    function buildBasePrompt(promptInfo, inputForAI) {
+        if (promptInfo.appendUnsubstitutedInput === false) {
+            return promptInfo.text;
+        }
+        const substituted = promptInfo.substitutedKeys ?? new Set();
+        const remaining = {};
+        for (const [key, value] of Object.entries(inputForAI)) {
+            if (!substituted.has(key)) {
+                remaining[key] = value;
+            }
+        }
+        if (Object.keys(remaining).length === 0) {
+            return promptInfo.text;
+        }
+        return `${promptInfo.text}\n\nInput: ${JSON.stringify(remaining)}`;
+    }
     // ── Shared generate implementation (returns data + usage) ──
     async function _generateCore(params) {
         const start = performance.now();
@@ -76,51 +201,100 @@ export function createAIService(config) {
         const hasPromptDef = promptRegistry?.has(params.prompt);
         const promptInfo = hasPromptDef
             ? buildPromptText(params.prompt, inputForAI)
-            : { text: params.prompt };
+            : { text: params.prompt, substitutedKeys: new Set() };
         // Resolve provider: prompt-level > default
         const activeProvider = resolveProvider('provider' in promptInfo ? promptInfo.provider : undefined);
         // Check if we have a schema to validate against
         const promptDef = hasPromptDef ? promptRegistry?.get(params.prompt) : undefined;
+        const singleTextField = promptDef ? getSingleTextFieldName(promptDef.output) : undefined;
+        const responseSchema = getStrictResponseSchema(params.prompt, promptDef);
+        const basePrompt = buildBasePrompt(promptInfo, inputForAI);
+        const useMultiTurn = Boolean(params.messages && params.messages.length > 0);
+        const mergedSystem = useMultiTurn
+            ? [promptInfo.system, basePrompt].filter((s) => s && String(s).trim().length > 0).join('\n\n')
+            : promptInfo.system;
         const request = {
-            prompt: Object.keys(inputForAI).length > 0
-                ? `${promptInfo.text}\n\nInput: ${JSON.stringify(inputForAI)}`
-                : promptInfo.text,
+            system: mergedSystem || undefined,
+            prompt: useMultiTurn ? '' : basePrompt,
+            messages: params.messages,
             model: promptInfo.model ?? config.defaultModel,
             temperature: promptInfo.temperature,
             maxTokens: promptInfo.maxTokens,
+            responseFormat: promptDef && !singleTextField ? 'json' : undefined,
+            responseSchema,
+            structuredOutputTransport: promptDef?.structuredOutputTransport,
+            signal: params.signal,
+            seed: params.seed,
         };
+        const resolvedModel = promptInfo.model ?? config.defaultModel ?? activeProvider.name;
         let result;
         let totalUsage = { inputTokens: 0, outputTokens: 0, totalTokens: 0 };
         let validationAttempts = 1;
         const validationPassed = true;
-        if (promptDef) {
-            // Use validated generation
-            const validated = await generateWithValidation(activeProvider, request, promptDef.output, config.validation);
-            result = validated.data;
-            totalUsage = validated.usage;
-            validationAttempts = validated.attempts;
-        }
-        else {
-            // Raw generation
-            const response = await activeProvider.complete(request);
-            result = response.content;
-            totalUsage = response.usage;
+        try {
+            if (promptDef) {
+                // Use validated generation
+                const validated = await generateWithValidation(activeProvider, request, promptDef.output, {
+                    ...config.validation,
+                    ...params.validation,
+                    model: resolvedModel,
+                    provider: activeProvider.name,
+                });
+                result = validated.data;
+                totalUsage = validated.usage;
+                validationAttempts = validated.attempts;
+            }
+            else {
+                // Raw generation
+                const response = await activeProvider.complete(request);
+                result = response.content;
+                totalUsage = response.usage;
+            }
         }
-        const latencyMs = performance.now() - start;
-        // Track cost
-        if (costTracker) {
-            costTracker.record({
-                model: promptInfo.model ?? config.defaultModel ?? activeProvider.name,
+        catch (err) {
+            // Failure-path cost recording: capture whatever usage we know about
+            // (from AIValidationError or zero for opaque fetch errors) so the
+            // sunk provider-side spend is still visible to budgets and ledgers.
+            const latencyMs = performance.now() - start;
+            const failureUsage = err instanceof AIValidationError
+                ? err.usage
+                : (getTypedAIErrorUsage(err) ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 });
+            const errorMessage = err instanceof Error ? err.message : String(err);
+            const failureCost = failureUsage.totalTokens > 0
+                ? calculateModelCost(failureUsage.inputTokens, failureUsage.outputTokens, resolvedModel, {
+                    cachedInputTokens: failureUsage.cachedInputTokens,
+                    cacheWriteTokens: failureUsage.cacheWriteTokens,
+                })
+                : null;
+            await recordCost({
+                model: resolvedModel,
                 provider: activeProvider.name,
                 promptName: hasPromptDef ? params.prompt : undefined,
                 operation: 'generate',
-                usage: totalUsage,
-                cost: null,
+                usage: failureUsage,
+                cost: failureCost,
                 latencyMs,
                 tenantId: config.budget?.tenantId,
                 actor: config.budget?.actor,
-            });
+                status: getTypedAIErrorStatus(err),
+                errorMessage,
+            }, params.costContext);
+            throw err;
         }
+        const latencyMs = performance.now() - start;
+        // Track cost (success path)
+        await recordCost({
+            model: resolvedModel,
+            provider: activeProvider.name,
+            promptName: hasPromptDef ? params.prompt : undefined,
+            operation: 'generate',
+            usage: totalUsage,
+            cost: null,
+            latencyMs,
+            tenantId: config.budget?.tenantId,
+            actor: config.budget?.actor,
+            status: 'success',
+        }, params.costContext);
         // Explainability
         if (explainability) {
             explainability.record({
@@ -141,7 +315,6 @@ export function createAIService(config) {
                 latencyMs,
             });
         }
-        const resolvedModel = promptInfo.model ?? config.defaultModel ?? activeProvider.name;
         const cost = calculateModelCost(totalUsage.inputTokens, totalUsage.outputTokens, resolvedModel, {
             cachedInputTokens: totalUsage.cachedInputTokens,
             cacheWriteTokens: totalUsage.cacheWriteTokens,
@@ -172,69 +345,311 @@ export function createAIService(config) {
             const hasPromptDef = promptRegistry?.has(params.prompt);
             const promptInfo = hasPromptDef
                 ? buildPromptText(params.prompt, inputForAI)
-                : { text: params.prompt };
+                : { text: params.prompt, substitutedKeys: new Set() };
             // Resolve provider
             const activeProvider = resolveProvider('provider' in promptInfo ? promptInfo.provider : undefined);
             // Detect if the output schema is a simple single-string-field object
             // (e.g. z.object({ question: z.string() })). If so, stream plain text
             // so the user sees readable words instead of JSON tokens.
             const promptDef = hasPromptDef ? promptRegistry?.get(params.prompt) : undefined;
-            let singleTextField;
-            if (promptDef) {
-                const shape = promptDef.output instanceof z.ZodObject ? promptDef.output.shape : undefined;
-                if (shape) {
-                    const keys = Object.keys(shape);
-                    const firstKey = keys[0];
-                    if (keys.length === 1 && firstKey && shape[firstKey] instanceof z.ZodString) {
-                        singleTextField = firstKey;
-                    }
-                }
-            }
-            const basePrompt = Object.keys(inputForAI).length > 0
-                ? `${promptInfo.text}\n\nInput: ${JSON.stringify(inputForAI)}`
-                : promptInfo.text;
-            const request = {
-                prompt: singleTextField
-                    ? `${basePrompt}\n\nRespond with ONLY the plain text content. Do NOT wrap your response in JSON or any other format.`
-                    : !basePrompt.toLowerCase().includes('json')
+            const singleTextField = promptDef ? getSingleTextFieldName(promptDef.output) : undefined;
+            const responseSchema = getStrictResponseSchema(params.prompt, promptDef);
+            const basePrompt = buildBasePrompt(promptInfo, inputForAI);
+            // The text-mode brevity hint is the right default for Q&A-style prompts
+            // (it stops the model from accidentally wrapping a one-line answer in
+            // JSON), but some models read "ONLY the plain text content" as a
+            // brevity instruction and collapse long-form outputs into bullet
+            // points. Prompts can opt out via `disableTextModeBrevityHint: true`
+            // when they own their own length contract in the prompt body.
+            const skipTextBrevityHint = promptDef?.disableTextModeBrevityHint === true;
+            const useMultiTurn = Boolean(params.messages && params.messages.length > 0);
+            const promptForProvider = useMultiTurn
+                ? ''
+                : singleTextField
+                    ? skipTextBrevityHint
+                        ? basePrompt
+                        : `${basePrompt}\n\nRespond with ONLY the plain text content. Do NOT wrap your response in JSON or any other format.`
+                    : !responseSchema && !basePrompt.toLowerCase().includes('json')
                         ? `${basePrompt}\n\nRespond with a valid JSON object.`
-                        : basePrompt,
+                        : basePrompt;
+            const mergedSystem = useMultiTurn
+                ? [promptInfo.system, basePrompt]
+                    .filter((s) => s && String(s).trim().length > 0)
+                    .join('\n\n')
+                : promptInfo.system;
+            const request = {
+                system: mergedSystem || undefined,
+                prompt: promptForProvider,
+                // Native multi-turn: when caller supplies a messages array, providers
+                // use it verbatim and ignore `prompt`. Otherwise legacy single-user
+                // mode applies.
+                messages: params.messages,
                 model: promptInfo.model ?? config.defaultModel,
                 temperature: promptInfo.temperature,
                 maxTokens: promptInfo.maxTokens,
                 responseFormat: singleTextField ? 'text' : 'json',
+                responseSchema,
+                structuredOutputTransport: promptDef?.structuredOutputTransport,
+                signal: params.signal,
+                seed: params.seed,
             };
-            // Stream from provider, collecting full text
+            const resolvedModel = promptInfo.model ?? config.defaultModel ?? activeProvider.name;
+            // Stream from provider, collecting full text and usage
             let fullText = '';
-            for await (const event of activeProvider.stream(request)) {
-                if (event.type === 'content_delta' && event.delta) {
-                    fullText += event.delta;
-                    yield { type: 'delta', text: event.delta };
-                }
-                else if (event.type === 'error') {
-                    yield { type: 'error', error: event.error };
-                    return;
+            let streamUsage;
+            let providerFinishReason;
+            try {
+                for await (const event of activeProvider.stream(request)) {
+                    if (event.type === 'content_delta' && event.delta) {
+                        fullText += event.delta;
+                        yield { type: 'delta', text: event.delta };
+                    }
+                    else if (event.type === 'usage' && event.usage) {
+                        streamUsage = {
+                            inputTokens: event.usage.inputTokens,
+                            outputTokens: event.usage.outputTokens,
+                            totalTokens: event.usage.totalTokens,
+                            cachedInputTokens: event.usage.cachedInputTokens,
+                            cacheWriteTokens: event.usage.cacheWriteTokens,
+                        };
+                    }
+                    else if (event.type === 'done') {
+                        // Capture the provider's finish_reason so downstream callers can
+                        // distinguish natural completion ('stop') from truncation
+                        // ('length') — essential when diagnosing short-output compression
+                        // vs a hit max_tokens ceiling.
+                        if (event.finishReason) {
+                            providerFinishReason = event.finishReason;
+                        }
+                    }
+                    else if (event.type === 'error') {
+                        // Provider-reported mid-stream error — record whatever usage we
+                        // have so the sunk spend is visible, then yield error to caller.
+                        const errCost = streamUsage
+                            ? calculateModelCost(streamUsage.inputTokens, streamUsage.outputTokens, resolvedModel, {
+                                cachedInputTokens: streamUsage.cachedInputTokens,
+                                cacheWriteTokens: streamUsage.cacheWriteTokens,
+                            })
+                            : null;
+                        await recordCost({
+                            model: resolvedModel,
+                            provider: activeProvider.name,
+                            promptName: hasPromptDef ? params.prompt : undefined,
+                            operation: 'generate',
+                            usage: streamUsage ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
+                            cost: errCost,
+                            latencyMs: 0,
+                            tenantId: config.budget?.tenantId,
+                            actor: config.budget?.actor,
+                            status: 'failed',
+                            errorMessage: event.error,
+                        }, params.costContext);
+                        yield { type: 'error', error: event.error };
+                        return;
+                    }
                 }
             }
+            catch (err) {
+                // Transport-level failure mid-stream (fetch error, abort, etc.).
+                const errUsage = getTypedAIErrorUsage(err) ??
+                    streamUsage ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 };
+                const errCost = errUsage.totalTokens > 0
+                    ? calculateModelCost(errUsage.inputTokens, errUsage.outputTokens, resolvedModel, {
+                        cachedInputTokens: errUsage.cachedInputTokens,
+                        cacheWriteTokens: errUsage.cacheWriteTokens,
+                    })
+                    : null;
+                await recordCost({
+                    model: resolvedModel,
+                    provider: activeProvider.name,
+                    promptName: hasPromptDef ? params.prompt : undefined,
+                    operation: 'generate',
+                    usage: errUsage,
+                    cost: errCost,
+                    latencyMs: 0,
+                    tenantId: config.budget?.tenantId,
+                    actor: config.budget?.actor,
+                    status: getTypedAIErrorStatus(err),
+                    errorMessage: err instanceof Error ? err.message : String(err),
+                }, params.costContext);
+                throw err;
+            }
+            // Compute cost and model info for the done event
+            const streamCost = streamUsage
+                ? calculateModelCost(streamUsage.inputTokens, streamUsage.outputTokens, resolvedModel, {
+                    cachedInputTokens: streamUsage.cachedInputTokens,
+                    cacheWriteTokens: streamUsage.cacheWriteTokens,
+                })
+                : 0;
+            const doneBase = {
+                usage: streamUsage,
+                model: resolvedModel,
+                provider: activeProvider.name,
+                cost: streamCost,
+                finishReason: providerFinishReason,
+            };
+            if (request.responseSchema &&
+                (providerFinishReason === 'length' || providerFinishReason === 'max_tokens')) {
+                const failureUsage = streamUsage ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 };
+                await recordCost({
+                    model: resolvedModel,
+                    provider: activeProvider.name,
+                    promptName: hasPromptDef ? params.prompt : undefined,
+                    operation: 'generate',
+                    usage: failureUsage,
+                    cost: null,
+                    latencyMs: 0,
+                    tenantId: config.budget?.tenantId,
+                    actor: config.budget?.actor,
+                    status: 'incomplete',
+                    errorMessage: `stream stopped before completing structured output: ${providerFinishReason}`,
+                }, params.costContext);
+                throw new AIIncompleteOutputError({
+                    provider: activeProvider.name,
+                    model: resolvedModel,
+                    partialText: fullText,
+                    usage: failureUsage,
+                    finishReason: providerFinishReason,
+                });
+            }
             // Build the final validated result
             if (singleTextField) {
                 // Plain text mode — wrap the accumulated text in the schema field
-                yield { type: 'done', data: { [singleTextField]: fullText.trim() } };
+                await recordCost({
+                    model: resolvedModel,
+                    provider: activeProvider.name,
+                    promptName: hasPromptDef ? params.prompt : undefined,
+                    operation: 'generate',
+                    usage: streamUsage ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
+                    cost: null,
+                    latencyMs: 0,
+                    tenantId: config.budget?.tenantId,
+                    actor: config.budget?.actor,
+                    status: 'success',
+                }, params.costContext);
+                yield { type: 'done', data: { [singleTextField]: fullText.trim() }, ...doneBase };
             }
             else if (promptDef) {
+                let validatedData;
+                let validationFellBackToNonStreaming = false;
+                let fallbackUsage;
                 try {
                     const parsed = JSON.parse(fullText);
-                    const validated = promptDef.output.parse(parsed);
-                    yield { type: 'done', data: validated };
+                    validatedData = promptDef.output.parse(parsed);
                 }
-                catch {
-                    // Validation failed — fall back to non-streaming retry with JSON mode
-                    const validated = await generateWithValidation(activeProvider, { ...request, responseFormat: 'json' }, promptDef.output, config.validation);
-                    yield { type: 'done', data: validated.data };
+                catch (parseErr) {
+                    // Streamed output failed local validation. The default behaviour
+                    // is to fall back to a non-streaming JSON-mode retry — but that
+                    // silently doubles the input-token cost of large prompts. Prompts
+                    // that opt out via `skipStreamValidationFallback` get a hard
+                    // failure instead, so callers can implement their own recovery.
+                    if (promptDef.skipStreamValidationFallback) {
+                        const errMsg = parseErr instanceof Error ? parseErr.message : String(parseErr);
+                        console.error(`[stream:${params.prompt}] validation-fallback SKIPPED by prompt config — stream produced invalid JSON (${errMsg}); throwing. streamedChars=${fullText.length} finishReason=${providerFinishReason ?? 'unknown'}`);
+                        const failureUsage = streamUsage ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 };
+                        const failureCost = failureUsage.totalTokens > 0
+                            ? calculateModelCost(failureUsage.inputTokens, failureUsage.outputTokens, resolvedModel, {
+                                cachedInputTokens: failureUsage.cachedInputTokens,
+                                cacheWriteTokens: failureUsage.cacheWriteTokens,
+                            })
+                            : null;
+                        await recordCost({
+                            model: resolvedModel,
+                            provider: activeProvider.name,
+                            promptName: hasPromptDef ? params.prompt : undefined,
+                            operation: 'generate',
+                            usage: failureUsage,
+                            cost: failureCost,
+                            latencyMs: 0,
+                            tenantId: config.budget?.tenantId,
+                            actor: config.budget?.actor,
+                            status: 'failed',
+                            errorMessage: `stream-validation failed and fallback disabled: ${errMsg}`,
+                        }, params.costContext);
+                        throw new Error(`${params.prompt}: streaming JSON validation failed and fallback is disabled — ${errMsg}`);
+                    }
+                    // Default path: fall back to a non-streaming JSON-mode retry with
+                    // full validation+feedback. Hidden cost doubler — log loudly and
+                    // tag the cost-ledger row so billing can surface how often it
+                    // fires.
+                    validationFellBackToNonStreaming = true;
+                    console.warn(`[stream:${params.prompt}] validation-fallback fired — stream produced invalid JSON (${parseErr instanceof Error ? parseErr.message : String(parseErr)}); replaying non-streaming. streamedChars=${fullText.length} finishReason=${providerFinishReason ?? 'unknown'}`);
+                    try {
+                        const validated = await generateWithValidation(activeProvider, { ...request, responseFormat: 'json' }, promptDef.output, { ...config.validation, model: resolvedModel, provider: activeProvider.name });
+                        validatedData = validated.data;
+                        fallbackUsage = validated.usage;
+                    }
+                    catch (fallbackErr) {
+                        // Bill the caller for the combined streamUsage + whatever the
+                        // fallback provider ate before giving up (AIValidationError
+                        // carries the accumulated retry usage since A1).
+                        const fallbackErrorUsage = fallbackErr instanceof AIValidationError
+                            ? fallbackErr.usage
+                            : (getTypedAIErrorUsage(fallbackErr) ?? {
+                                inputTokens: 0,
+                                outputTokens: 0,
+                                totalTokens: 0,
+                            });
+                        const combinedUsage = sumUsage(streamUsage, fallbackErrorUsage);
+                        const combinedCost = combinedUsage.totalTokens > 0
+                            ? calculateModelCost(combinedUsage.inputTokens, combinedUsage.outputTokens, resolvedModel, {
+                                cachedInputTokens: combinedUsage.cachedInputTokens,
+                                cacheWriteTokens: combinedUsage.cacheWriteTokens,
+                            })
+                            : null;
+                        await recordCost({
+                            model: resolvedModel,
+                            provider: activeProvider.name,
+                            promptName: hasPromptDef ? params.prompt : undefined,
+                            operation: 'generate',
+                            usage: combinedUsage,
+                            cost: combinedCost,
+                            latencyMs: 0,
+                            tenantId: config.budget?.tenantId,
+                            actor: config.budget?.actor,
+                            status: getTypedAIErrorStatus(fallbackErr),
+                            errorMessage: fallbackErr instanceof Error ? fallbackErr.message : String(fallbackErr),
+                        }, params.costContext);
+                        throw fallbackErr;
+                    }
                 }
+                const successUsage = validationFellBackToNonStreaming
+                    ? sumUsage(streamUsage, fallbackUsage)
+                    : (streamUsage ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 });
+                await recordCost({
+                    model: resolvedModel,
+                    provider: activeProvider.name,
+                    promptName: hasPromptDef ? params.prompt : undefined,
+                    operation: 'generate',
+                    fallbackUsed: validationFellBackToNonStreaming || undefined,
+                    usage: successUsage,
+                    cost: null,
+                    latencyMs: 0,
+                    tenantId: config.budget?.tenantId,
+                    actor: config.budget?.actor,
+                    status: 'success',
+                }, params.costContext);
+                yield {
+                    type: 'done',
+                    data: validatedData,
+                    ...doneBase,
+                    validationFallbackFired: validationFellBackToNonStreaming,
+                };
             }
             else {
-                yield { type: 'done', data: { content: fullText } };
+                await recordCost({
+                    model: resolvedModel,
+                    provider: activeProvider.name,
+                    promptName: hasPromptDef ? params.prompt : undefined,
+                    operation: 'generate',
+                    usage: streamUsage ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
+                    cost: null,
+                    latencyMs: 0,
+                    tenantId: config.budget?.tenantId,
+                    actor: config.budget?.actor,
+                    status: 'success',
+                }, params.costContext);
+                yield { type: 'done', data: { content: fullText }, ...doneBase };
             }
         },
         async extract(params) {
@@ -242,27 +657,57 @@ export function createAIService(config) {
             checkBudget();
             // extract uses default provider (no prompt-level routing)
             const activeProvider = resolveProvider();
+            const resolvedModel = config.defaultModel ?? activeProvider.name;
             const systemPrompt = 'Extract structured data from the following text. Return valid JSON matching the required schema.';
             const request = {
                 system: systemPrompt,
                 prompt: params.text,
                 model: config.defaultModel,
                 responseFormat: 'json',
+                responseSchema: config.enableStrictStructuredOutputs
+                    ? zodToProviderJsonSchema(params.schema, { promptName: 'extract' }).schema
+                    : undefined,
+                signal: params.signal,
             };
-            const validated = await generateWithValidation(activeProvider, request, params.schema, config.validation);
-            const latencyMs = performance.now() - start;
-            if (costTracker) {
-                costTracker.record({
-                    model: config.defaultModel ?? activeProvider.name,
+            let validated;
+            try {
+                validated = (await generateWithValidation(activeProvider, request, params.schema, {
+                    ...config.validation,
+                    model: resolvedModel,
+                    provider: activeProvider.name,
+                }));
+            }
+            catch (err) {
+                const latencyMs = performance.now() - start;
+                const failureUsage = err instanceof AIValidationError
+                    ? err.usage
+                    : (getTypedAIErrorUsage(err) ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 });
+                await recordCost({
+                    model: resolvedModel,
                     provider: activeProvider.name,
                     operation: 'extract',
-                    usage: validated.usage,
+                    usage: failureUsage,
                     cost: null,
                     latencyMs,
                     tenantId: config.budget?.tenantId,
                     actor: config.budget?.actor,
-                });
+                    status: getTypedAIErrorStatus(err),
+                    errorMessage: err instanceof Error ? err.message : String(err),
+                }, params.costContext);
+                throw err;
             }
+            const latencyMs = performance.now() - start;
+            await recordCost({
+                model: resolvedModel,
+                provider: activeProvider.name,
+                operation: 'extract',
+                usage: validated.usage,
+                cost: null,
+                latencyMs,
+                tenantId: config.budget?.tenantId,
+                actor: config.budget?.actor,
+                status: 'success',
+            }, params.costContext);
             if (explainability) {
                 explainability.record({
                     operation: 'extract',
@@ -284,30 +729,60 @@ export function createAIService(config) {
             checkBudget();
             // classify uses default provider (no prompt-level routing)
             const activeProvider = resolveProvider();
+            const resolvedModel = config.defaultModel ?? activeProvider.name;
             const systemPrompt = 'Classify the following text into one or more of the provided labels. Return a JSON array of matching label strings.';
             const request = {
                 system: systemPrompt,
                 prompt: `Labels: ${JSON.stringify(params.labels)}\n\nText: ${params.text}`,
                 model: config.defaultModel,
                 responseFormat: 'json',
+                responseSchema: config.enableStrictStructuredOutputs
+                    ? zodToProviderJsonSchema(z.array(z.string()), { promptName: 'classify' }).schema
+                    : undefined,
+                signal: params.signal,
             };
             const schema = z.array(z.string());
-            const validated = await generateWithValidation(activeProvider, request, schema, config.validation);
-            const latencyMs = performance.now() - start;
-            // Filter to only valid labels
-            const result = validated.data.filter((l) => params.labels.includes(l));
-            if (costTracker) {
-                costTracker.record({
-                    model: config.defaultModel ?? activeProvider.name,
+            let validated;
+            try {
+                validated = await generateWithValidation(activeProvider, request, schema, {
+                    ...config.validation,
+                    model: resolvedModel,
+                    provider: activeProvider.name,
+                });
+            }
+            catch (err) {
+                const latencyMs = performance.now() - start;
+                const failureUsage = err instanceof AIValidationError
+                    ? err.usage
+                    : (getTypedAIErrorUsage(err) ?? { inputTokens: 0, outputTokens: 0, totalTokens: 0 });
+                await recordCost({
+                    model: resolvedModel,
                     provider: activeProvider.name,
                     operation: 'classify',
-                    usage: validated.usage,
+                    usage: failureUsage,
                     cost: null,
                     latencyMs,
                     tenantId: config.budget?.tenantId,
                     actor: config.budget?.actor,
-                });
+                    status: getTypedAIErrorStatus(err),
+                    errorMessage: err instanceof Error ? err.message : String(err),
+                }, params.costContext);
+                throw err;
             }
+            const latencyMs = performance.now() - start;
+            // Filter to only valid labels
+            const result = validated.data.filter((l) => params.labels.includes(l));
+            await recordCost({
+                model: resolvedModel,
+                provider: activeProvider.name,
+                operation: 'classify',
+                usage: validated.usage,
+                cost: null,
+                latencyMs,
+                tenantId: config.budget?.tenantId,
+                actor: config.budget?.actor,
+                status: 'success',
+            }, params.costContext);
             if (explainability) {
                 explainability.record({
                     operation: 'classify',