@stackfactor/agent-utils 1.0.19 → 1.0.21

package/dist/cjs/const.d.ts

@@ -1,6 +1,7 @@
 declare const _default: {
     HTTP_CODES: {
         BAD_REQUEST: number;
+        PAYMENT_REQUIRED: number;
         UNPROCESSABLE_ENTITY: number;
         INTERNAL_SERVER_ERROR: number;
         BAD_GATEWAY: number;
@@ -9,6 +10,7 @@ declare const _default: {
         UNABLE_TO_GENERATE_CONTENT: string;
         UNEXPECTED_ERROR: string;
         UNSUPPORTED_MODEL: string;
+        QUOTA_EXHAUSTED: string;
     };
 };
 export default _default;

package/dist/cjs/const.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"const.d.ts","sourceRoot":"","sources":["../../src/const.ts"],"names":[],"mappings":";;;;;;;;;;;;;AAAA,wBAaE"}
+{"version":3,"file":"const.d.ts","sourceRoot":"","sources":["../../src/const.ts"],"names":[],"mappings":";;;;;;;;;;;;;;;AAAA,wBAgBE"}

package/dist/cjs/const.js

@@ -3,6 +3,7 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.default = {
     HTTP_CODES: {
         BAD_REQUEST: 400,
+        PAYMENT_REQUIRED: 402,
         UNPROCESSABLE_ENTITY: 422,
         INTERNAL_SERVER_ERROR: 500,
         BAD_GATEWAY: 502,
@@ -11,5 +12,6 @@ exports.default = {
         UNABLE_TO_GENERATE_CONTENT: "Unable to generate content",
         UNEXPECTED_ERROR: "An unexpected error occured. If the issue persists please contact the StackFactor support team at support@stackfactor.ai",
         UNSUPPORTED_MODEL: "The specified model is not supported",
+        QUOTA_EXHAUSTED: "Agent session quota exhausted: no remaining budget for additional LLM calls.",
     },
 };

package/dist/cjs/langChain.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":";wBAgWQ,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBAyBG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,KAC1B,OAAO,CAAC,GAAG,CAAC;wCAmfF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,KACpB,GAAG;oCAxYO,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,KACX,OAAO,CAAC,GAAG,CAAC;sDA6rBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,KACX,OAAO,CAAC,GAAG,CAAC;0CA7vB8B,GAAG,KAAG,MAAM;;AAwzBzD,wBAOE"}
+{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":";wBAogBQ,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBAyBG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,KAC1B,OAAO,CAAC,GAAG,CAAC;wCAkkBF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,KACpB,GAAG;oCA7XO,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,KACX,OAAO,CAAC,GAAG,CAAC;sDAyrBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,KACX,OAAO,CAAC,GAAG,CAAC;0CA30B8B,GAAG,KAAG,MAAM;;AAs4BzD,wBAOE"}

package/dist/cjs/langChain.js

@@ -16,12 +16,157 @@ const zod_to_json_schema_1 = require("zod-to-json-schema");
 const JSON_ESCAPE_INSTRUCTION = `
 CRITICAL - Your response must be valid JSON. Escape ALL special characters in string values:
 - Newlines → \\n
-- Tabs → \\t  
+- Tabs → \\t
 - Carriage returns → \\r
 - Double quotes inside strings → \\"
 - Backslashes → \\\\
 Do NOT include raw newlines, tabs, or unescaped quotes inside JSON string values.
 `.trim();
+/**
+ * Reads `global.quota` (set by the StackFactor host before invoking `main()`) and
+ * throws a `PAYMENT_REQUIRED` error when the session has no remaining budget. No-op
+ * when `global.quota` is absent (e.g. local tests) or `remaining` is not a number,
+ * which keeps the library usable outside the StackFactor runtime.
+ */
+const assertQuotaAvailable = () => {
+    const quota = globalThis.quota;
+    if (!quota || typeof quota.remaining !== "number")
+        return;
+    if (quota.remaining <= 0) {
+        throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.PAYMENT_REQUIRED, const_js_1.default.ERROR.QUOTA_EXHAUSTED);
+    }
+};
+/**
+ * Subtracts the given USD cost from `global.quota.remaining` and adds it to
+ * `global.quota.usedThisSession`. Silently no-ops when the quota globals are not
+ * present so callers do not need to branch. Negative, zero, or non-finite costs
+ * are ignored to keep counters monotonic.
+ */
+const recordCost = (cost) => {
+    if (!cost || cost <= 0 || !Number.isFinite(cost))
+        return;
+    const quota = globalThis.quota;
+    if (!quota)
+        return;
+    if (typeof quota.remaining === "number") {
+        quota.remaining -= cost;
+    }
+    if (typeof quota.usedThisSession === "number") {
+        quota.usedThisSession += cost;
+    }
+    else {
+        quota.usedThisSession = cost;
+    }
+};
+/**
+ * Looks up the per-model pricing entry from `config.modelPricing`. Returns `null`
+ * when pricing is not configured for the model, in which case cost recording
+ * becomes a no-op (the call still succeeds — pricing data is the host's
+ * responsibility, not the agent's).
+ */
+const getModelPrice = (modelName, config) => {
+    const pricing = config?.modelPricing;
+    if (!pricing)
+        return null;
+    return pricing[modelName] || null;
+};
+/**
+ * Computes USD cost for a text LLM call. `config.modelPricing[modelName]` is
+ * expected to provide `input` and `output` rates in dollars per million tokens.
+ */
+const calculateTextCost = (modelName, usage, config) => {
+    if (!usage)
+        return 0;
+    const price = getModelPrice(modelName, config);
+    if (!price)
+        return 0;
+    const inputCost = ((usage.input_tokens || 0) / 1_000_000) * (price.input || 0);
+    const outputCost = ((usage.output_tokens || 0) / 1_000_000) * (price.output || 0);
+    return inputCost + outputCost;
+};
+/**
+ * Computes USD cost for an image generation call. `config.modelPricing[modelName]`
+ * is expected to provide a `perImage` rate in dollars.
+ */
+const calculateImageCost = (modelName, numImages, config) => {
+    const price = getModelPrice(modelName, config);
+    if (!price)
+        return 0;
+    return (numImages || 0) * (price.perImage || 0);
+};
+/**
+ * Extracts a normalized token-usage object from a single LangChain `invoke()`
+ * response. Reads from `usage_metadata` first (standardized in LangChain v1),
+ * then falls back to provider-specific shapes under `response_metadata`.
+ */
+const extractUsageFromInvoke = (response) => {
+    if (!response)
+        return null;
+    const um = response.usage_metadata;
+    if (um) {
+        return {
+            input_tokens: um.input_tokens || 0,
+            output_tokens: um.output_tokens || 0,
+        };
+    }
+    const rm = response.response_metadata;
+    if (rm?.usage) {
+        return {
+            input_tokens: rm.usage.input_tokens || rm.usage.prompt_tokens || 0,
+            output_tokens: rm.usage.output_tokens || rm.usage.completion_tokens || 0,
+        };
+    }
+    if (rm?.tokenUsage) {
+        return {
+            input_tokens: rm.tokenUsage.promptTokens || 0,
+            output_tokens: rm.tokenUsage.completionTokens || 0,
+        };
+    }
+    return null;
+};
+/**
+ * Accumulates token usage from a streaming chunk into a running total. LangChain
+ * typically attaches `usage_metadata` to the final chunk; earlier chunks carry no
+ * usage and are no-ops here.
+ */
+const accumulateChunkUsage = (acc, chunk) => {
+    if (!chunk)
+        return acc;
+    const um = chunk.usage_metadata;
+    if (um) {
+        acc.input_tokens += um.input_tokens || 0;
+        acc.output_tokens += um.output_tokens || 0;
+        return acc;
+    }
+    const rm = chunk.response_metadata;
+    if (rm?.usage) {
+        acc.input_tokens += rm.usage.input_tokens || rm.usage.prompt_tokens || 0;
+        acc.output_tokens +=
+            rm.usage.output_tokens || rm.usage.completion_tokens || 0;
+    }
+    return acc;
+};
+/**
+ * Sums token usage across every message in an agent invocation response. Each
+ * AIMessage in `response.messages` may carry its own `usage_metadata` (one per
+ * LLM round-trip the agent made).
+ */
+const sumAgentResponseUsage = (response) => {
+    const messages = response?.messages;
+    if (!Array.isArray(messages) || messages.length === 0)
+        return null;
+    const total = { input_tokens: 0, output_tokens: 0 };
+    let found = false;
+    for (const msg of messages) {
+        const um = msg?.usage_metadata;
+        if (um) {
+            total.input_tokens += um.input_tokens || 0;
+            total.output_tokens += um.output_tokens || 0;
+            found = true;
+        }
+    }
+    return found ? total : null;
+};
 /**
  * Converts a Zod validation error object into a human-readable multi-line string.
  * Each failing field is described with its dot-notation path and a contextual message
@@ -377,12 +522,15 @@ const runAgent = async (agent, prompt, config, onProgress = null) => {
             },
         ]
         : undefined;
+    assertQuotaAvailable();
     const response = await agent.invoke({
         messages: [{ role: "user", content: prompt }],
     }, {
         recursionLimit: config.recursionLimit || 25,
         ...(callbacks ? { callbacks } : {}),
     });
+    const modelName = agent.options?.model?.modelName || agent.options?.model?.model || "";
+    recordCost(calculateTextCost(modelName, sumAgentResponseUsage(response), config));
     const endTime = Date.now();
     const duration = endTime - startTime;
     logger_js_1.default.log(null, logger_js_1.default.levels.info, `Agent "${agent.options?.name}" completed in ${Math.round(duration / 1000)} seconds.`);
@@ -403,6 +551,75 @@ const throwErrorIfNotSuccessful = (response) => {
         throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.INTERNAL_SERVER_ERROR, const_js_1.default.ERROR.UNABLE_TO_GENERATE_CONTENT);
     }
 };
+const MAX_VALIDATION_RETRIES = 3;
+/**
+ * Parses raw LLM content as JSON and validates against an optional Zod schema,
+ * returning the canonical JSON string. Tries a direct parse of the trimmed text
+ * first, then falls back to `extractJSONFromResponse` for markdown-wrapped output.
+ * Throws `UNPROCESSABLE_ENTITY` when the parsed value fails schema validation, or
+ * `INTERNAL_SERVER_ERROR` when no JSON can be extracted from the response at all.
+ */
+const parseAndValidateJSONResponse = (rawContent, schema) => {
+    if (typeof rawContent === "string") {
+        const trimmed = rawContent.trim();
+        if ((trimmed.startsWith("{") && trimmed.endsWith("}")) ||
+            (trimmed.startsWith("[") && trimmed.endsWith("]"))) {
+            try {
+                const parsed = JSON.parse(trimmed);
+                if (schema) {
+                    validateWithSchema(parsed, schema);
+                }
+                return JSON.stringify(parsed);
+            }
+            catch (parseError) {
+                if (parseError?.code === const_js_1.default.HTTP_CODES.UNPROCESSABLE_ENTITY) {
+                    throw parseError;
+                }
+                // JSON.parse failed (not a validation error) — fall through to extraction.
+            }
+        }
+        const extracted = extractJSONFromResponse(rawContent);
+        if (extracted) {
+            if (schema) {
+                validateWithSchema(extracted, schema);
+            }
+            return JSON.stringify(extracted);
+        }
+    }
+    let preview = "";
+    if (typeof rawContent === "string") {
+        preview = rawContent.substring(0, 100);
+    }
+    else if (typeof rawContent === "object" && rawContent !== null) {
+        preview = JSON.stringify(rawContent).substring(0, 100);
+    }
+    else if (rawContent !== undefined && rawContent !== null) {
+        preview = String(rawContent).substring(0, 100);
+    }
+    else {
+        preview = "[empty response]";
+    }
+    throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.INTERNAL_SERVER_ERROR, `LLM returned non-JSON response when JSON was expected: ${preview}...`);
+};
+/**
+ * Appends the model's failed output as an assistant turn followed by a user turn
+ * containing the Zod validation errors. Including the bad assistant turn (not just
+ * the critique) lets the model see what it actually produced — without it,
+ * self-correction is guesswork. Callers pass the *current* message list (not the
+ * original base) so that across multiple retries the full failure history
+ * accumulates and the model can avoid oscillating between previous broken outputs.
+ */
+const buildValidationRetryMessages = (priorMessages, rawContent, validationError) => {
+    const assistantContent = typeof rawContent === "string" ? rawContent : JSON.stringify(rawContent);
+    return [
+        ...priorMessages,
+        { role: "assistant", content: assistantContent },
+        {
+            role: "user",
+            content: `The previous response failed schema validation with the following errors. Return a corrected JSON response that fixes ALL of these issues:\n\n${validationError.message}`,
+        },
+    ];
+};
 /**
  * Sends a prompt to an LLM and returns the response, with support for streaming,
  * agentic execution, progress reporting, JSON extraction, and Zod schema validation.
@@ -580,80 +797,79 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
         else {
             messagesToSend = messages;
         }
-        // Stream and report server-side progress based on time elapsed
-        let rawContent = "";
-        let chunkCount = 0;
+        // Retry once on schema validation failure. The retry is silent to the progress
+        // callback — same "Generating content..." message and the progress curve
+        // continues from the original startTime so the bar doesn't visibly reset.
+        // Skipped when OpenAI native response_format is in use (a 422 there is a
+        // server-side schema bug, not a model output issue). Each attempt is billed
+        // independently via recordCost so retry cost is still visible in telemetry.
+        const canRetryOnValidation = expectsJsonResponse && !!schema && !useNativeSchema;
         const progressReportInterval = 10; // Report every N chunks
-        const startTime = Date.now();
+        const overallStartTime = Date.now();
         // Use a time-based asymptotic curve: progress approaches maxPercent but never
         // overshoots. This avoids the magic "expected length" constant — longer responses
         // simply slow the curve down rather than exceeding the range.
         const expectedDurationMs = 15_000; // Tune: expected typical response time
-        const stream = await llm.stream(messagesToSend);
-        for await (const chunk of stream) {
-            const content = chunk?.content || chunk;
-            if (typeof content === "string") {
-                rawContent += content;
-                chunkCount++;
-                if (chunkCount % progressReportInterval === 0) {
-                    const elapsed = Date.now() - startTime;
-                    // Asymptotic curve: fast early progress that slows as it approaches max
-                    const progress = Math.round(minPercent +
-                        (maxPercent - minPercent - 5) *
-                            (1 - Math.exp(-elapsed / expectedDurationMs)));
-                    await onProgressReport({
-                        message: "Generating content...",
-                        progress: Math.min(progress, maxPercent - 5),
-                    });
+        let activeMessages = messagesToSend;
+        let attempt = 0;
+        while (true) {
+            let rawContent = "";
+            let chunkCount = 0;
+            const streamUsage = { input_tokens: 0, output_tokens: 0 };
+            assertQuotaAvailable();
+            const stream = await llm.stream(activeMessages);
+            for await (const chunk of stream) {
+                accumulateChunkUsage(streamUsage, chunk);
+                const content = chunk?.content || chunk;
+                if (typeof content === "string") {
+                    rawContent += content;
+                    chunkCount++;
+                    if (chunkCount % progressReportInterval === 0) {
+                        const elapsed = Date.now() - overallStartTime;
+                        // Asymptotic curve: fast early progress that slows as it approaches max
+                        const progress = Math.round(minPercent +
+                            (maxPercent - minPercent - 5) *
+                                (1 - Math.exp(-elapsed / expectedDurationMs)));
+                        await onProgressReport({
+                            message: "Generating content...",
+                            progress: Math.min(progress, maxPercent - 5),
+                        });
+                    }
                 }
             }
-        }
-        await onProgressReport({
-            message: "Processing complete",
-            progress: maxPercent,
-        });
-        if (!rawContent) {
-            throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.INTERNAL_SERVER_ERROR, "LLM returned no content");
-        }
-        // If not expecting JSON, return raw content directly
-        if (!expectsJsonResponse) {
-            return rawContent;
-        }
-        // Parse and validate JSON response
-        const trimmed = rawContent.trim();
-        if ((trimmed.startsWith("{") && trimmed.endsWith("}")) ||
-            (trimmed.startsWith("[") && trimmed.endsWith("]"))) {
+            recordCost(calculateTextCost(modelName, streamUsage, config));
+            if (!rawContent) {
+                throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.INTERNAL_SERVER_ERROR, "LLM returned no content");
+            }
+            // If not expecting JSON, return raw content directly
+            if (!expectsJsonResponse) {
+                await onProgressReport({
+                    message: "Processing complete",
+                    progress: maxPercent,
+                });
+                return rawContent;
+            }
             try {
-                const parsed = JSON.parse(trimmed);
-                if (schema) {
-                    validateWithSchema(parsed, schema);
-                }
-                return JSON.stringify(parsed);
+                const result = parseAndValidateJSONResponse(rawContent, schema);
+                await onProgressReport({
+                    message: "Processing complete",
+                    progress: maxPercent,
+                });
+                return result;
             }
-            catch (parseError) {
-                if (parseError?.code === const_js_1.default.HTTP_CODES.UNPROCESSABLE_ENTITY) {
-                    throw parseError;
+            catch (err) {
+                const isValidationFailure = err?.code === const_js_1.default.HTTP_CODES.UNPROCESSABLE_ENTITY;
+                if (isValidationFailure &&
+                    canRetryOnValidation &&
+                    attempt < MAX_VALIDATION_RETRIES) {
+                    attempt++;
+                    logger_js_1.default.log(null, logger_js_1.default.levels.warn, `Schema validation failed for "${modelName}", retrying (${attempt}/${MAX_VALIDATION_RETRIES}): ${err.message}`);
+                    activeMessages = buildValidationRetryMessages(activeMessages, rawContent, err);
+                    continue;
                 }
+                throw err;
             }
         }
-        const extracted = extractJSONFromResponse(rawContent);
-        if (extracted) {
-            if (schema) {
-                validateWithSchema(extracted, schema);
-            }
-            return JSON.stringify(extracted);
-        }
-        let preview = "";
-        if (typeof rawContent === "string") {
-            preview = rawContent.substring(0, 100);
-        }
-        else if (typeof rawContent === "object" && rawContent !== null) {
-            preview = JSON.stringify(rawContent).substring(0, 100);
-        }
-        else {
-            preview = "[empty response]";
-        }
-        throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.INTERNAL_SERVER_ERROR, `LLM returned non-JSON response when JSON was expected: ${preview}...`);
     }
     else {
         // Non-streaming mode: use native response_format for OpenAI when schema is provided
@@ -693,65 +909,36 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
         else {
             messagesToSend = messages;
         }
-        // Simply invoke without streaming
-        const response = await llm.invoke(messagesToSend);
-        const rawContent = response?.content || response;
-        // If not expecting JSON, return raw content directly
-        if (!expectsJsonResponse) {
-            return rawContent;
-        }
-        // If the response is already a string that looks like JSON, return it
-        // Otherwise, try to extract JSON from potential markdown wrapping
-        if (typeof rawContent === "string") {
-            const trimmed = rawContent.trim();
-            // Check if it's already clean JSON
-            if ((trimmed.startsWith("{") && trimmed.endsWith("}")) ||
-                (trimmed.startsWith("[") && trimmed.endsWith("]"))) {
-                try {
-                    // Parse and re-stringify
-                    const parsed = JSON.parse(trimmed);
-                    // Validate against schema if provided
-                    if (schema) {
-                        validateWithSchema(parsed, schema);
-                    }
-                    return JSON.stringify(parsed);
-                }
-                catch (parseError) {
-                    // If it's a validation error, re-throw it
-                    if (parseError?.code === const_js_1.default.HTTP_CODES.UNPROCESSABLE_ENTITY) {
-                        throw parseError;
-                    }
-                    // Not valid JSON, try extraction
-                }
+        // Retry once on schema validation failure (see streaming branch for rationale).
+        // Each attempt is billed independently via recordCost so retry cost is visible.
+        const canRetryOnValidation = expectsJsonResponse && !!schema && !useNativeSchema;
+        let activeMessages = messagesToSend;
+        let attempt = 0;
+        while (true) {
+            assertQuotaAvailable();
+            const response = await llm.invoke(activeMessages);
+            recordCost(calculateTextCost(modelName, extractUsageFromInvoke(response), config));
+            const rawContent = response?.content || response;
+            // If not expecting JSON, return raw content directly
+            if (!expectsJsonResponse) {
+                return rawContent;
             }
-            // Try to extract JSON from markdown or other wrapping
-            const extracted = extractJSONFromResponse(rawContent);
-            if (extracted) {
-                // Validate against schema if provided
-                if (schema) {
-                    validateWithSchema(extracted, schema);
-                }
-                return JSON.stringify(extracted);
-            }
-        }
-        // Return raw content as last resort - but throw if JSON was expected
-        if (expectsJsonResponse) {
-            let preview = "";
-            if (typeof rawContent === "string") {
-                preview = rawContent.substring(0, 100);
-            }
-            else if (typeof rawContent === "object" && rawContent !== null) {
-                preview = JSON.stringify(rawContent).substring(0, 100);
-            }
-            else if (rawContent !== undefined && rawContent !== null) {
-                preview = String(rawContent).substring(0, 100);
+            try {
+                return parseAndValidateJSONResponse(rawContent, schema);
             }
-            else {
-                preview = "[empty response]";
+            catch (err) {
+                const isValidationFailure = err?.code === const_js_1.default.HTTP_CODES.UNPROCESSABLE_ENTITY;
+                if (isValidationFailure &&
+                    canRetryOnValidation &&
+                    attempt < MAX_VALIDATION_RETRIES) {
+                    attempt++;
+                    logger_js_1.default.log(null, logger_js_1.default.levels.warn, `Schema validation failed for "${modelName}", retrying (${attempt}/${MAX_VALIDATION_RETRIES}): ${err.message}`);
+                    activeMessages = buildValidationRetryMessages(activeMessages, rawContent, err);
+                    continue;
+                }
+                throw err;
             }
-            throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.INTERNAL_SERVER_ERROR, `LLM returned non-JSON response when JSON was expected: ${preview}...`);
         }
-        return rawContent;
     }
 };
 /**
@@ -877,7 +1064,9 @@ const generateImageWithOpenAI = async (modelName, config, prompt, options) => {
             requestParams.style = style;
         }
     }
+    assertQuotaAvailable();
     const response = await openai.images.generate(requestParams);
+    recordCost(calculateImageCost(modelName, response.data?.length || n, config));
     // Format response based on number of images
     if (n === 1) {
         const imageData = response.data[0];
@@ -975,6 +1164,7 @@ const generateImageWithGoogle = async (modelName, config, prompt, options) => {
             safetySettings: safetySettings,
         },
     };
+    assertQuotaAvailable();
     const response = await ai.models.generateContent(req);
     // Extract images from response
     const images = [];
@@ -993,6 +1183,7 @@ const generateImageWithGoogle = async (modelName, config, prompt, options) => {
     if (images.length === 0) {
         throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.INTERNAL_SERVER_ERROR, `No images were generated by ${modelName}`);
     }
+    recordCost(calculateImageCost(modelName, images.length, config));
     if (numberOfImages === 1 || images.length === 1) {
         return images[0];
     }

package/dist/esm/const.d.ts

@@ -1,6 +1,7 @@
 declare const _default: {
     HTTP_CODES: {
         BAD_REQUEST: number;
+        PAYMENT_REQUIRED: number;
         UNPROCESSABLE_ENTITY: number;
         INTERNAL_SERVER_ERROR: number;
         BAD_GATEWAY: number;
@@ -9,6 +10,7 @@ declare const _default: {
         UNABLE_TO_GENERATE_CONTENT: string;
         UNEXPECTED_ERROR: string;
         UNSUPPORTED_MODEL: string;
+        QUOTA_EXHAUSTED: string;
     };
 };
 export default _default;

package/dist/esm/const.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"const.d.ts","sourceRoot":"","sources":["../../src/const.ts"],"names":[],"mappings":";;;;;;;;;;;;;AAAA,wBAaE"}
+{"version":3,"file":"const.d.ts","sourceRoot":"","sources":["../../src/const.ts"],"names":[],"mappings":";;;;;;;;;;;;;;;AAAA,wBAgBE"}

package/dist/esm/const.js

@@ -1,6 +1,7 @@
 export default {
     HTTP_CODES: {
         BAD_REQUEST: 400,
+        PAYMENT_REQUIRED: 402,
         UNPROCESSABLE_ENTITY: 422,
         INTERNAL_SERVER_ERROR: 500,
         BAD_GATEWAY: 502,
@@ -9,5 +10,6 @@ export default {
         UNABLE_TO_GENERATE_CONTENT: "Unable to generate content",
         UNEXPECTED_ERROR: "An unexpected error occured. If the issue persists please contact the StackFactor support team at support@stackfactor.ai",
         UNSUPPORTED_MODEL: "The specified model is not supported",
+        QUOTA_EXHAUSTED: "Agent session quota exhausted: no remaining budget for additional LLM calls.",
     },
 };

package/dist/esm/langChain.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":";wBAgWQ,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBAyBG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,KAC1B,OAAO,CAAC,GAAG,CAAC;wCAmfF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,KACpB,GAAG;oCAxYO,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,KACX,OAAO,CAAC,GAAG,CAAC;sDA6rBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,KACX,OAAO,CAAC,GAAG,CAAC;0CA7vB8B,GAAG,KAAG,MAAM;;AAwzBzD,wBAOE"}
+{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":";wBAogBQ,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBAyBG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,KAC1B,OAAO,CAAC,GAAG,CAAC;wCAkkBF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,KACpB,GAAG;oCA7XO,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,KACX,OAAO,CAAC,GAAG,CAAC;sDAyrBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,KACX,OAAO,CAAC,GAAG,CAAC;0CA30B8B,GAAG,KAAG,MAAM;;AAs4BzD,wBAOE"}

package/dist/esm/langChain.js

@@ -11,12 +11,157 @@ import { zodToJsonSchema } from "zod-to-json-schema";
 const JSON_ESCAPE_INSTRUCTION = `
 CRITICAL - Your response must be valid JSON. Escape ALL special characters in string values:
 - Newlines → \\n
-- Tabs → \\t  
+- Tabs → \\t
 - Carriage returns → \\r
 - Double quotes inside strings → \\"
 - Backslashes → \\\\
 Do NOT include raw newlines, tabs, or unescaped quotes inside JSON string values.
 `.trim();
+/**
+ * Reads `global.quota` (set by the StackFactor host before invoking `main()`) and
+ * throws a `PAYMENT_REQUIRED` error when the session has no remaining budget. No-op
+ * when `global.quota` is absent (e.g. local tests) or `remaining` is not a number,
+ * which keeps the library usable outside the StackFactor runtime.
+ */
+const assertQuotaAvailable = () => {
+    const quota = globalThis.quota;
+    if (!quota || typeof quota.remaining !== "number")
+        return;
+    if (quota.remaining <= 0) {
+        throw errorHandlingHelper.create(constants.HTTP_CODES.PAYMENT_REQUIRED, constants.ERROR.QUOTA_EXHAUSTED);
+    }
+};
+/**
+ * Subtracts the given USD cost from `global.quota.remaining` and adds it to
+ * `global.quota.usedThisSession`. Silently no-ops when the quota globals are not
+ * present so callers do not need to branch. Negative, zero, or non-finite costs
+ * are ignored to keep counters monotonic.
+ */
+const recordCost = (cost) => {
+    if (!cost || cost <= 0 || !Number.isFinite(cost))
+        return;
+    const quota = globalThis.quota;
+    if (!quota)
+        return;
+    if (typeof quota.remaining === "number") {
+        quota.remaining -= cost;
+    }
+    if (typeof quota.usedThisSession === "number") {
+        quota.usedThisSession += cost;
+    }
+    else {
+        quota.usedThisSession = cost;
+    }
+};
+/**
+ * Looks up the per-model pricing entry from `config.modelPricing`. Returns `null`
+ * when pricing is not configured for the model, in which case cost recording
+ * becomes a no-op (the call still succeeds — pricing data is the host's
+ * responsibility, not the agent's).
+ */
+const getModelPrice = (modelName, config) => {
+    const pricing = config?.modelPricing;
+    if (!pricing)
+        return null;
+    return pricing[modelName] || null;
+};
+/**
+ * Computes USD cost for a text LLM call. `config.modelPricing[modelName]` is
+ * expected to provide `input` and `output` rates in dollars per million tokens.
+ */
+const calculateTextCost = (modelName, usage, config) => {
+    if (!usage)
+        return 0;
+    const price = getModelPrice(modelName, config);
+    if (!price)
+        return 0;
+    const inputCost = ((usage.input_tokens || 0) / 1_000_000) * (price.input || 0);
+    const outputCost = ((usage.output_tokens || 0) / 1_000_000) * (price.output || 0);
+    return inputCost + outputCost;
+};
+/**
+ * Computes USD cost for an image generation call. `config.modelPricing[modelName]`
+ * is expected to provide a `perImage` rate in dollars.
+ */
+const calculateImageCost = (modelName, numImages, config) => {
+    const price = getModelPrice(modelName, config);
+    if (!price)
+        return 0;
+    return (numImages || 0) * (price.perImage || 0);
+};
+/**
+ * Extracts a normalized token-usage object from a single LangChain `invoke()`
+ * response. Reads from `usage_metadata` first (standardized in LangChain v1),
+ * then falls back to provider-specific shapes under `response_metadata`.
+ */
+const extractUsageFromInvoke = (response) => {
+    if (!response)
+        return null;
+    const um = response.usage_metadata;
+    if (um) {
+        return {
+            input_tokens: um.input_tokens || 0,
+            output_tokens: um.output_tokens || 0,
+        };
+    }
+    const rm = response.response_metadata;
+    if (rm?.usage) {
+        return {
+            input_tokens: rm.usage.input_tokens || rm.usage.prompt_tokens || 0,
+            output_tokens: rm.usage.output_tokens || rm.usage.completion_tokens || 0,
+        };
+    }
+    if (rm?.tokenUsage) {
+        return {
+            input_tokens: rm.tokenUsage.promptTokens || 0,
+            output_tokens: rm.tokenUsage.completionTokens || 0,
+        };
+    }
+    return null;
+};
+/**
+ * Accumulates token usage from a streaming chunk into a running total. LangChain
+ * typically attaches `usage_metadata` to the final chunk; earlier chunks carry no
+ * usage and are no-ops here.
+ */
+const accumulateChunkUsage = (acc, chunk) => {
+    if (!chunk)
+        return acc;
+    const um = chunk.usage_metadata;
+    if (um) {
+        acc.input_tokens += um.input_tokens || 0;
+        acc.output_tokens += um.output_tokens || 0;
+        return acc;
+    }
+    const rm = chunk.response_metadata;
+    if (rm?.usage) {
+        acc.input_tokens += rm.usage.input_tokens || rm.usage.prompt_tokens || 0;
+        acc.output_tokens +=
+            rm.usage.output_tokens || rm.usage.completion_tokens || 0;
+    }
+    return acc;
+};
+/**
+ * Sums token usage across every message in an agent invocation response. Each
+ * AIMessage in `response.messages` may carry its own `usage_metadata` (one per
+ * LLM round-trip the agent made).
+ */
+const sumAgentResponseUsage = (response) => {
+    const messages = response?.messages;
+    if (!Array.isArray(messages) || messages.length === 0)
+        return null;
+    const total = { input_tokens: 0, output_tokens: 0 };
+    let found = false;
+    for (const msg of messages) {
+        const um = msg?.usage_metadata;
+        if (um) {
+            total.input_tokens += um.input_tokens || 0;
+            total.output_tokens += um.output_tokens || 0;
+            found = true;
+        }
+    }
+    return found ? total : null;
+};
 /**
  * Converts a Zod validation error object into a human-readable multi-line string.
  * Each failing field is described with its dot-notation path and a contextual message
@@ -372,12 +517,15 @@ const runAgent = async (agent, prompt, config, onProgress = null) => {
             },
         ]
         : undefined;
+    assertQuotaAvailable();
     const response = await agent.invoke({
         messages: [{ role: "user", content: prompt }],
     }, {
         recursionLimit: config.recursionLimit || 25,
         ...(callbacks ? { callbacks } : {}),
     });
+    const modelName = agent.options?.model?.modelName || agent.options?.model?.model || "";
+    recordCost(calculateTextCost(modelName, sumAgentResponseUsage(response), config));
     const endTime = Date.now();
     const duration = endTime - startTime;
     logger.log(null, logger.levels.info, `Agent "${agent.options?.name}" completed in ${Math.round(duration / 1000)} seconds.`);
@@ -398,6 +546,75 @@ const throwErrorIfNotSuccessful = (response) => {
         throw errorHandlingHelper.create(constants.HTTP_CODES.INTERNAL_SERVER_ERROR, constants.ERROR.UNABLE_TO_GENERATE_CONTENT);
     }
 };
+const MAX_VALIDATION_RETRIES = 3;
+/**
+ * Parses raw LLM content as JSON and validates against an optional Zod schema,
+ * returning the canonical JSON string. Tries a direct parse of the trimmed text
+ * first, then falls back to `extractJSONFromResponse` for markdown-wrapped output.
+ * Throws `UNPROCESSABLE_ENTITY` when the parsed value fails schema validation, or
+ * `INTERNAL_SERVER_ERROR` when no JSON can be extracted from the response at all.
+ */
+const parseAndValidateJSONResponse = (rawContent, schema) => {
+    if (typeof rawContent === "string") {
+        const trimmed = rawContent.trim();
+        if ((trimmed.startsWith("{") && trimmed.endsWith("}")) ||
+            (trimmed.startsWith("[") && trimmed.endsWith("]"))) {
+            try {
+                const parsed = JSON.parse(trimmed);
+                if (schema) {
+                    validateWithSchema(parsed, schema);
+                }
+                return JSON.stringify(parsed);
+            }
+            catch (parseError) {
+                if (parseError?.code === constants.HTTP_CODES.UNPROCESSABLE_ENTITY) {
+                    throw parseError;
+                }
+                // JSON.parse failed (not a validation error) — fall through to extraction.
+            }
+        }
+        const extracted = extractJSONFromResponse(rawContent);
+        if (extracted) {
+            if (schema) {
+                validateWithSchema(extracted, schema);
+            }
+            return JSON.stringify(extracted);
+        }
+    }
+    let preview = "";
+    if (typeof rawContent === "string") {
+        preview = rawContent.substring(0, 100);
+    }
+    else if (typeof rawContent === "object" && rawContent !== null) {
+        preview = JSON.stringify(rawContent).substring(0, 100);
+    }
+    else if (rawContent !== undefined && rawContent !== null) {
+        preview = String(rawContent).substring(0, 100);
+    }
+    else {
+        preview = "[empty response]";
+    }
+    throw errorHandlingHelper.create(constants.HTTP_CODES.INTERNAL_SERVER_ERROR, `LLM returned non-JSON response when JSON was expected: ${preview}...`);
+};
+/**
+ * Appends the model's failed output as an assistant turn followed by a user turn
+ * containing the Zod validation errors. Including the bad assistant turn (not just
+ * the critique) lets the model see what it actually produced — without it,
+ * self-correction is guesswork. Callers pass the *current* message list (not the
+ * original base) so that across multiple retries the full failure history
+ * accumulates and the model can avoid oscillating between previous broken outputs.
+ */
+const buildValidationRetryMessages = (priorMessages, rawContent, validationError) => {
+    const assistantContent = typeof rawContent === "string" ? rawContent : JSON.stringify(rawContent);
+    return [
+        ...priorMessages,
+        { role: "assistant", content: assistantContent },
+        {
+            role: "user",
+            content: `The previous response failed schema validation with the following errors. Return a corrected JSON response that fixes ALL of these issues:\n\n${validationError.message}`,
+        },
+    ];
+};
 /**
  * Sends a prompt to an LLM and returns the response, with support for streaming,
  * agentic execution, progress reporting, JSON extraction, and Zod schema validation.
@@ -575,80 +792,79 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
         else {
             messagesToSend = messages;
         }
-        // Stream and report server-side progress based on time elapsed
-        let rawContent = "";
-        let chunkCount = 0;
+        // Retry once on schema validation failure. The retry is silent to the progress
+        // callback — same "Generating content..." message and the progress curve
+        // continues from the original startTime so the bar doesn't visibly reset.
+        // Skipped when OpenAI native response_format is in use (a 422 there is a
+        // server-side schema bug, not a model output issue). Each attempt is billed
+        // independently via recordCost so retry cost is still visible in telemetry.
+        const canRetryOnValidation = expectsJsonResponse && !!schema && !useNativeSchema;
         const progressReportInterval = 10; // Report every N chunks
-        const startTime = Date.now();
+        const overallStartTime = Date.now();
         // Use a time-based asymptotic curve: progress approaches maxPercent but never
         // overshoots. This avoids the magic "expected length" constant — longer responses
         // simply slow the curve down rather than exceeding the range.
         const expectedDurationMs = 15_000; // Tune: expected typical response time
-        const stream = await llm.stream(messagesToSend);
-        for await (const chunk of stream) {
-            const content = chunk?.content || chunk;
-            if (typeof content === "string") {
-                rawContent += content;
-                chunkCount++;
-                if (chunkCount % progressReportInterval === 0) {
-                    const elapsed = Date.now() - startTime;
-                    // Asymptotic curve: fast early progress that slows as it approaches max
-                    const progress = Math.round(minPercent +
-                        (maxPercent - minPercent - 5) *
-                            (1 - Math.exp(-elapsed / expectedDurationMs)));
-                    await onProgressReport({
-                        message: "Generating content...",
-                        progress: Math.min(progress, maxPercent - 5),
-                    });
+        let activeMessages = messagesToSend;
+        let attempt = 0;
+        while (true) {
+            let rawContent = "";
+            let chunkCount = 0;
+            const streamUsage = { input_tokens: 0, output_tokens: 0 };
+            assertQuotaAvailable();
+            const stream = await llm.stream(activeMessages);
+            for await (const chunk of stream) {
+                accumulateChunkUsage(streamUsage, chunk);
+                const content = chunk?.content || chunk;
+                if (typeof content === "string") {
+                    rawContent += content;
+                    chunkCount++;
+                    if (chunkCount % progressReportInterval === 0) {
+                        const elapsed = Date.now() - overallStartTime;
+                        // Asymptotic curve: fast early progress that slows as it approaches max
+                        const progress = Math.round(minPercent +
+                            (maxPercent - minPercent - 5) *
+                                (1 - Math.exp(-elapsed / expectedDurationMs)));
+                        await onProgressReport({
+                            message: "Generating content...",
+                            progress: Math.min(progress, maxPercent - 5),
+                        });
+                    }
                 }
             }
-        }
-        await onProgressReport({
-            message: "Processing complete",
-            progress: maxPercent,
-        });
-        if (!rawContent) {
-            throw errorHandlingHelper.create(constants.HTTP_CODES.INTERNAL_SERVER_ERROR, "LLM returned no content");
-        }
-        // If not expecting JSON, return raw content directly
-        if (!expectsJsonResponse) {
-            return rawContent;
-        }
-        // Parse and validate JSON response
-        const trimmed = rawContent.trim();
-        if ((trimmed.startsWith("{") && trimmed.endsWith("}")) ||
-            (trimmed.startsWith("[") && trimmed.endsWith("]"))) {
+            recordCost(calculateTextCost(modelName, streamUsage, config));
+            if (!rawContent) {
+                throw errorHandlingHelper.create(constants.HTTP_CODES.INTERNAL_SERVER_ERROR, "LLM returned no content");
+            }
+            // If not expecting JSON, return raw content directly
+            if (!expectsJsonResponse) {
+                await onProgressReport({
+                    message: "Processing complete",
+                    progress: maxPercent,
+                });
+                return rawContent;
+            }
             try {
-                const parsed = JSON.parse(trimmed);
-                if (schema) {
-                    validateWithSchema(parsed, schema);
-                }
-                return JSON.stringify(parsed);
+                const result = parseAndValidateJSONResponse(rawContent, schema);
+                await onProgressReport({
+                    message: "Processing complete",
+                    progress: maxPercent,
+                });
+                return result;
             }
-            catch (parseError) {
-                if (parseError?.code === constants.HTTP_CODES.UNPROCESSABLE_ENTITY) {
-                    throw parseError;
+            catch (err) {
+                const isValidationFailure = err?.code === constants.HTTP_CODES.UNPROCESSABLE_ENTITY;
+                if (isValidationFailure &&
+                    canRetryOnValidation &&
+                    attempt < MAX_VALIDATION_RETRIES) {
+                    attempt++;
+                    logger.log(null, logger.levels.warn, `Schema validation failed for "${modelName}", retrying (${attempt}/${MAX_VALIDATION_RETRIES}): ${err.message}`);
+                    activeMessages = buildValidationRetryMessages(activeMessages, rawContent, err);
+                    continue;
                 }
+                throw err;
             }
         }
-        const extracted = extractJSONFromResponse(rawContent);
-        if (extracted) {
-            if (schema) {
-                validateWithSchema(extracted, schema);
-            }
-            return JSON.stringify(extracted);
-        }
-        let preview = "";
-        if (typeof rawContent === "string") {
-            preview = rawContent.substring(0, 100);
-        }
-        else if (typeof rawContent === "object" && rawContent !== null) {
-            preview = JSON.stringify(rawContent).substring(0, 100);
-        }
-        else {
-            preview = "[empty response]";
-        }
-        throw errorHandlingHelper.create(constants.HTTP_CODES.INTERNAL_SERVER_ERROR, `LLM returned non-JSON response when JSON was expected: ${preview}...`);
     }
     else {
         // Non-streaming mode: use native response_format for OpenAI when schema is provided
@@ -688,65 +904,36 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
         else {
             messagesToSend = messages;
         }
-        // Simply invoke without streaming
-        const response = await llm.invoke(messagesToSend);
-        const rawContent = response?.content || response;
-        // If not expecting JSON, return raw content directly
-        if (!expectsJsonResponse) {
-            return rawContent;
-        }
-        // If the response is already a string that looks like JSON, return it
-        // Otherwise, try to extract JSON from potential markdown wrapping
-        if (typeof rawContent === "string") {
-            const trimmed = rawContent.trim();
-            // Check if it's already clean JSON
-            if ((trimmed.startsWith("{") && trimmed.endsWith("}")) ||
-                (trimmed.startsWith("[") && trimmed.endsWith("]"))) {
-                try {
-                    // Parse and re-stringify
-                    const parsed = JSON.parse(trimmed);
-                    // Validate against schema if provided
-                    if (schema) {
-                        validateWithSchema(parsed, schema);
-                    }
-                    return JSON.stringify(parsed);
-                }
-                catch (parseError) {
-                    // If it's a validation error, re-throw it
-                    if (parseError?.code === constants.HTTP_CODES.UNPROCESSABLE_ENTITY) {
-                        throw parseError;
-                    }
-                    // Not valid JSON, try extraction
-                }
+        // Retry once on schema validation failure (see streaming branch for rationale).
+        // Each attempt is billed independently via recordCost so retry cost is visible.
+        const canRetryOnValidation = expectsJsonResponse && !!schema && !useNativeSchema;
+        let activeMessages = messagesToSend;
+        let attempt = 0;
+        while (true) {
+            assertQuotaAvailable();
+            const response = await llm.invoke(activeMessages);
+            recordCost(calculateTextCost(modelName, extractUsageFromInvoke(response), config));
+            const rawContent = response?.content || response;
+            // If not expecting JSON, return raw content directly
+            if (!expectsJsonResponse) {
+                return rawContent;
             }
-            // Try to extract JSON from markdown or other wrapping
-            const extracted = extractJSONFromResponse(rawContent);
-            if (extracted) {
-                // Validate against schema if provided
-                if (schema) {
-                    validateWithSchema(extracted, schema);
-                }
-                return JSON.stringify(extracted);
-            }
-        }
-        // Return raw content as last resort - but throw if JSON was expected
-        if (expectsJsonResponse) {
-            let preview = "";
-            if (typeof rawContent === "string") {
-                preview = rawContent.substring(0, 100);
-            }
-            else if (typeof rawContent === "object" && rawContent !== null) {
-                preview = JSON.stringify(rawContent).substring(0, 100);
-            }
-            else if (rawContent !== undefined && rawContent !== null) {
-                preview = String(rawContent).substring(0, 100);
+            try {
+                return parseAndValidateJSONResponse(rawContent, schema);
             }
-            else {
-                preview = "[empty response]";
+            catch (err) {
+                const isValidationFailure = err?.code === constants.HTTP_CODES.UNPROCESSABLE_ENTITY;
+                if (isValidationFailure &&
+                    canRetryOnValidation &&
+                    attempt < MAX_VALIDATION_RETRIES) {
+                    attempt++;
+                    logger.log(null, logger.levels.warn, `Schema validation failed for "${modelName}", retrying (${attempt}/${MAX_VALIDATION_RETRIES}): ${err.message}`);
+                    activeMessages = buildValidationRetryMessages(activeMessages, rawContent, err);
+                    continue;
+                }
+                throw err;
             }
-            throw errorHandlingHelper.create(constants.HTTP_CODES.INTERNAL_SERVER_ERROR, `LLM returned non-JSON response when JSON was expected: ${preview}...`);
         }
-        return rawContent;
     }
 };
 /**
@@ -872,7 +1059,9 @@ const generateImageWithOpenAI = async (modelName, config, prompt, options) => {
             requestParams.style = style;
         }
     }
+    assertQuotaAvailable();
     const response = await openai.images.generate(requestParams);
+    recordCost(calculateImageCost(modelName, response.data?.length || n, config));
     // Format response based on number of images
     if (n === 1) {
         const imageData = response.data[0];
@@ -970,6 +1159,7 @@ const generateImageWithGoogle = async (modelName, config, prompt, options) => {
             safetySettings: safetySettings,
         },
     };
+    assertQuotaAvailable();
     const response = await ai.models.generateContent(req);
     // Extract images from response
     const images = [];
@@ -988,6 +1178,7 @@ const generateImageWithGoogle = async (modelName, config, prompt, options) => {
     if (images.length === 0) {
         throw errorHandlingHelper.create(constants.HTTP_CODES.INTERNAL_SERVER_ERROR, `No images were generated by ${modelName}`);
     }
+    recordCost(calculateImageCost(modelName, images.length, config));
     if (numberOfImages === 1 || images.length === 1) {
         return images[0];
     }

package/package.json

@@ -3,7 +3,7 @@
   "publishConfig": {
     "access": "restricted"
   },
-  "version": "1.0.19",
+  "version": "1.0.21",
   "description": "",
   "main": "dist/cjs/index.js",
   "module": "dist/esm/index.js",