@stackfactor/agent-utils 1.2.8 → 1.2.11

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

@@ -1 +1 @@
-{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":"AA2BA;;;;;GAKG;AACH,MAAM,MAAM,YAAY,GAAG;IACzB,IAAI,EAAE,MAAM,CAAC;IACb,MAAM,EAAE;QAAE,CAAC,QAAQ,EAAE,MAAM,GAAG,MAAM,CAAA;KAAE,CAAC;CACxC,CAAC;;0CAE2C,GAAG,KAAG,IAAI;wBA+jB/C,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBA0BG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,iBACb,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;oCA4XF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,iBACE,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;sDAusBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,iBACE,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;0CA3+B8B,GAAG,KAAG,MAAM;+CApjB9C,YAAY,GAAG,IAAI,GAAG,SAAS,aAC7B,MAAM,kBACD,MAAM,UACd,GAAG,KACV,IAAI;mCAokBU,MAAM,mBACJ,MAAM,EAAE,KACxB,MAAM;;AAkhCT,wBASE"}
+{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":"AAgCA;;;;;GAKG;AACH,MAAM,MAAM,YAAY,GAAG;IACzB,IAAI,EAAE,MAAM,CAAC;IACb,MAAM,EAAE;QAAE,CAAC,QAAQ,EAAE,MAAM,GAAG,MAAM,CAAA;KAAE,CAAC;CACxC,CAAC;;0CAE2C,GAAG,KAAG,IAAI;wBA6lB/C,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBA0BG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,iBACb,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;oCA4XF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,iBACE,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;sDAstBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,iBACE,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;0CA1/B8B,GAAG,KAAG,MAAM;+CA/jB9C,YAAY,GAAG,IAAI,GAAG,SAAS,aAC7B,MAAM,kBACD,MAAM,UACd,GAAG,KACV,IAAI;mCA+kBU,MAAM,mBACJ,MAAM,EAAE,KACxB,MAAM;;AAiiCT,wBASE"}

package/dist/cjs/langChain.js

@@ -56,7 +56,19 @@ const updateUsageTracker = (tracker, modelName, usage, config) => {
     if (!tracker.tokens || typeof tracker.tokens !== "object")
         tracker.tokens = {};
     const inputTokens = usage.input_tokens || 0;
-    const outputTokens = usage.output_tokens || 0;
+    const visibleOutputTokens = usage.output_tokens || 0;
+    const totalTokens = usage.total_tokens || 0;
+    // Recover reasoning/"thinking" tokens the provider leaves out of
+    // `output_tokens`. Gemini reports its thoughts only in `totalTokenCount`, so
+    // `total - input - visibleOutput` is the thinking output the caller was still
+    // billed for. Providers that already fold reasoning into `output_tokens`
+    // (Anthropic, OpenAI) report `total == input + output`, so this adds 0.
+    // Reasoning is charged at the output rate, so we treat it as output for both
+    // the token counters and the cost.
+    const reasoningTokens = totalTokens > 0
+        ? Math.max(0, totalTokens - inputTokens - visibleOutputTokens)
+        : 0;
+    const outputTokens = visibleOutputTokens + reasoningTokens;
     const inputRate = getModelRate(modelName, config, "input-token");
     const outputRate = getModelRate(modelName, config, "output-token");
     const addedCost = (inputTokens / 1_000_000) * inputRate +
@@ -122,7 +134,12 @@ const updateImageUsageTracker = (tracker, modelName, response, config, provider)
         imageOutputTokens = Number(um.candidatesTokenCount) || 0;
     }
     const textInputRate = getModelRate(modelName, config, "input-token");
-    const imageInputRate = getModelRate(modelName, config, "image-input-token");
+    // Reference-image input tokens are billed at the model's input rate. Configs
+    // that don't define a dedicated `<model>-image-input-token-costs` (the common
+    // case — image models charge all input at one rate and only differ on output)
+    // fall back to the text input rate rather than silently billing image input
+    // at $0.
+    const imageInputRate = getModelRate(modelName, config, "image-input-token") || textInputRate;
     const imageOutputRate = getModelRate(modelName, config, "image-output-token");
     const addedCost = (textInputTokens / 1_000_000) * textInputRate +
         (imageInputTokens / 1_000_000) * imageInputRate +
@@ -176,6 +193,7 @@ const extractUsageFromInvoke = (response) => {
         return {
             input_tokens: um.input_tokens || 0,
             output_tokens: um.output_tokens || 0,
+            total_tokens: um.total_tokens || 0,
         };
     }
     const rm = response.response_metadata;
@@ -183,12 +201,14 @@ const extractUsageFromInvoke = (response) => {
         return {
             input_tokens: rm.usage.input_tokens || rm.usage.prompt_tokens || 0,
             output_tokens: rm.usage.output_tokens || rm.usage.completion_tokens || 0,
+            total_tokens: rm.usage.total_tokens || rm.usage.total_token_count || 0,
         };
     }
     if (rm?.tokenUsage) {
         return {
             input_tokens: rm.tokenUsage.promptTokens || 0,
             output_tokens: rm.tokenUsage.completionTokens || 0,
+            total_tokens: rm.tokenUsage.totalTokens || 0,
         };
     }
     return null;
@@ -205,6 +225,7 @@ const accumulateChunkUsage = (acc, chunk) => {
     if (um) {
         acc.input_tokens += um.input_tokens || 0;
         acc.output_tokens += um.output_tokens || 0;
+        acc.total_tokens += um.total_tokens || 0;
         return acc;
     }
     const rm = chunk.response_metadata;
@@ -212,6 +233,8 @@ const accumulateChunkUsage = (acc, chunk) => {
         acc.input_tokens += rm.usage.input_tokens || rm.usage.prompt_tokens || 0;
         acc.output_tokens +=
             rm.usage.output_tokens || rm.usage.completion_tokens || 0;
+        acc.total_tokens +=
+            rm.usage.total_tokens || rm.usage.total_token_count || 0;
     }
     return acc;
 };
@@ -224,13 +247,18 @@ const sumAgentResponseUsage = (response) => {
     const messages = response?.messages;
     if (!Array.isArray(messages) || messages.length === 0)
         return null;
-    const total = { input_tokens: 0, output_tokens: 0 };
+    const total = {
+        input_tokens: 0,
+        output_tokens: 0,
+        total_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;
+            total.total_tokens += um.total_tokens || 0;
             found = true;
         }
     }
@@ -1054,7 +1082,11 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
         while (true) {
             let rawContent = "";
             let chunkCount = 0;
-            let streamUsage = { input_tokens: 0, output_tokens: 0 };
+            let streamUsage = {
+                input_tokens: 0,
+                output_tokens: 0,
+                total_tokens: 0,
+            };
             // Inner loop: wait + retry on 429 around stream setup and consumption.
             // Usage is only recorded on a successful stream — partial streams that
             // error out with a rate limit are not counted. A 429 fired mid-stream
@@ -1063,7 +1095,7 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
             while (true) {
                 rawContent = "";
                 chunkCount = 0;
-                streamUsage = { input_tokens: 0, output_tokens: 0 };
+                streamUsage = { input_tokens: 0, output_tokens: 0, total_tokens: 0 };
                 try {
                     // Honour caller cancellation: passing the signal tears down the
                     // upstream HTTP request so a cancelled call stops billing tokens.
@@ -1339,20 +1371,21 @@ const generateImageWithOpenAI = async (modelName, config, prompt, options, usage
  * Google models do not support alpha channels. Validates the `aspectRatio` against
  * allowed values (`"1:1"`, `"3:4"`, `"4:3"`, `"9:16"`, `"16:9"`). Applies model-
  * specific generation config: Imagen models receive `numberOfImages`, `aspectRatio`,
- * `outputMimeType`, and optional `negativePrompt`; Gemini models receive `temperature`
- * and `topP`. Safety settings are disabled for image generation. Throws an error if no
- * images are returned.
+ * `outputMimeType`, and optional `negativePrompt`; Gemini models receive `temperature`,
+ * `topP`, and an `imageConfig` carrying `aspectRatio` plus an optional `imageSize`
+ * resolution tier (e.g. "1K"/"2K"/"4K", honoured by gemini-3-pro-image). Safety
+ * settings are disabled for image generation. Throws an error if no images are returned.
  * @param modelName - The Google model identifier (e.g. `"imagen-4.0-generate-001"`,
  *   `"gemini-3.0-pro-image"`)
  * @param config - Configuration object; must include `googleAPIKey`
  * @param prompt - The text prompt describing the image to generate
- * @param options - Generation options including `aspectRatio`, `numberOfImages`, and
- *   optional `negativePrompt` (Imagen only)
+ * @param options - Generation options including `aspectRatio`, optional `imageSize`
+ *   (Gemini resolution tier), `numberOfImages`, and optional `negativePrompt` (Imagen only)
  * @returns A single image descriptor `{ b64_json, mimeType }` when one image is
  *   requested, or `{ images: [...] }` for multiple images
  */
 const generateImageWithGoogle = async (modelName, config, prompt, options, usageTracker = null) => {
-    const { aspectRatio = "1:1", numberOfImages = 1, negativePrompt = "", } = options;
+    const { aspectRatio = "1:1", imageSize, numberOfImages = 1, negativePrompt = "", } = options;
     if (!config.googleAPIKey) {
         throw errorHandling_js_1.default.create(const_js_1.default.HTTP_CODES.BAD_REQUEST, "Google API key is required for Google image generation");
     }
@@ -1389,6 +1422,15 @@ const generateImageWithGoogle = async (modelName, config, prompt, options, usage
                 // Gemini-specific config
                 temperature: 1,
                 topP: 0.95,
+                // Gemini controls output dimensions through `imageConfig`, not the
+                // top-level `aspectRatio` field that Imagen uses. `imageSize`
+                // (e.g. "1K"/"2K"/"4K") is only honoured by resolution-capable
+                // models such as gemini-3-pro-image, so it is passed through only
+                // when the caller explicitly requests it.
+                imageConfig: {
+                    aspectRatio: aspectRatio,
+                    ...(imageSize ? { imageSize: imageSize } : {}),
+                },
             }),
     };
     // Safety settings (disable for image generation)

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

@@ -1 +1 @@
-{"version":3,"file":"serve.d.ts","sourceRoot":"","sources":["../../src/serve.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,IAAI,MAAM,eAAe,CAAC;AA8LtC;;;;;GAKG;AACH,MAAM,MAAM,SAAS,GAAG,CAAC,MAAM,CAAC,EAAE,WAAW,KAAK,OAAO,CAAC,GAAG,CAAC,CAAC;AAE/D;;;;GAIG;AACH,MAAM,MAAM,aAAa,GAAG,CAC1B,OAAO,EAAE,GAAG,KACT,OAAO,CAAC;IAAE,IAAI,EAAE,GAAG,CAAC;IAAC,MAAM,EAAE,GAAG,CAAA;CAAE,CAAC,CAAC;AAEzC,MAAM,WAAW,YAAY;IAC3B,4EAA4E;IAC5E,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,2CAA2C;IAC3C,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,wFAAwF;IACxF,YAAY,CAAC,EAAE,aAAa,CAAC;CAC9B;AAkDD;;;;;;;;;;GAUG;AACH,eAAO,MAAM,KAAK,GAChB,MAAM,SAAS,EACf,UAAS,YAAiB,KACzB,IAAI,CAAC,MAuKP,CAAC"}
+{"version":3,"file":"serve.d.ts","sourceRoot":"","sources":["../../src/serve.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,IAAI,MAAM,eAAe,CAAC;AA8LtC;;;;;GAKG;AACH,MAAM,MAAM,SAAS,GAAG,CAAC,MAAM,CAAC,EAAE,WAAW,KAAK,OAAO,CAAC,GAAG,CAAC,CAAC;AAE/D;;;;GAIG;AACH,MAAM,MAAM,aAAa,GAAG,CAC1B,OAAO,EAAE,GAAG,KACT,OAAO,CAAC;IAAE,IAAI,EAAE,GAAG,CAAC;IAAC,MAAM,EAAE,GAAG,CAAA;CAAE,CAAC,CAAC;AAEzC,MAAM,WAAW,YAAY;IAC3B,4EAA4E;IAC5E,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,2CAA2C;IAC3C,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,wFAAwF;IACxF,YAAY,CAAC,EAAE,aAAa,CAAC;CAC9B;AAyHD;;;;;;;;;;GAUG;AACH,eAAO,MAAM,KAAK,GAChB,MAAM,SAAS,EACf,UAAS,YAAiB,KACzB,IAAI,CAAC,MA4KP,CAAC"}

package/dist/cjs/serve.js

@@ -199,23 +199,78 @@ const runSelfCheck = async (checkCode, config, session) => {
 /** Marks an auth failure so the handler can map it to gRPC UNAUTHENTICATED. */
 class UnauthenticatedError extends Error {
 }
+/**
+ * Marks a *transient* failure to reach the StackFactor session-validation
+ * endpoint (5xx / request timeout / transport error). Unlike
+ * {@link UnauthenticatedError} the token was never actually rejected — the
+ * backend was momentarily unreachable — so the handler maps this to the
+ * retryable gRPC UNAVAILABLE instead of UNAUTHENTICATED, and the message no
+ * longer slanders a perfectly valid token.
+ */
+class BackendUnavailableError extends Error {
+}
+/**
+ * Session-validation retry policy. A GKE front-door blip (a backend pod briefly
+ * refusing connections during a restart or an event-loop stall) surfaces here as
+ * a 502/503 or a transport error, not a 401 — and retrying on a fresh connection
+ * almost always lands on a healthy replica. Env-overridable for prod tuning.
+ */
+const AUTH_MAX_ATTEMPTS = Math.max(1, Number(process.env.STACKFACTOR_AUTH_MAX_ATTEMPTS) || 4);
+const AUTH_RETRY_BASE_MS = Math.max(0, Number(process.env.STACKFACTOR_AUTH_RETRY_BASE_MS) || 500);
+const delay = (ms) => new Promise((res) => setTimeout(res, ms));
+/** A 401/403 is the backend genuinely rejecting the token. */
+const isAuthRejection = (error) => {
+    const status = error?.response?.status;
+    return status === 401 || status === 403;
+};
+/**
+ * True when the failure is the validation endpoint being unreachable rather than
+ * the token being bad: a 5xx/408/429 HTTP response, or a transport-level error
+ * (no HTTP response but an axios/errno signature). Plain programming errors
+ * (no response and no axios/errno marker) are NOT treated as transient.
+ */
+const isTransient = (error) => {
+    const status = error?.response?.status;
+    if (status)
+        return status >= 500 || status === 408 || status === 429;
+    return Boolean(error?.isAxiosError || error?.code);
+};
 /**
  * Default authenticator: require a token on the request and resolve it to a
  * session via the StackFactor API. Requires `BACKEND_URL` (or `REACT_APP_NODE_ENV`)
- * to point client-api at the right backend. Throws {@link UnauthenticatedError}
- * when the token is missing or rejected.
+ * to point client-api at the right backend.
+ *
+ * A 401/403 (or any non-transient error) throws {@link UnauthenticatedError}
+ * immediately. A transient failure (5xx / timeout / network) is retried with
+ * exponential backoff; if every attempt fails it throws
+ * {@link BackendUnavailableError} so the caller sees a retryable UNAVAILABLE
+ * rather than a misleading "invalid token".
  */
 const defaultAuthenticate = async (request) => {
     const token = request?.authToken ?? request?.authorization;
     if (!token) {
         throw new UnauthenticatedError("Missing StackFactor auth token");
     }
-    try {
-        return await client_api_1.session.getSession(token);
-    }
-    catch (error) {
-        throw new UnauthenticatedError(`Invalid StackFactor auth token: ${error?.message ?? error}`);
+    let lastError;
+    for (let attempt = 1; attempt <= AUTH_MAX_ATTEMPTS; attempt++) {
+        try {
+            return await client_api_1.session.getSession(token);
+        }
+        catch (error) {
+            lastError = error;
+            // Real rejection, or an unexpected non-transient error: fail fast.
+            if (isAuthRejection(error) || !isTransient(error)) {
+                throw new UnauthenticatedError(`Invalid StackFactor auth token: ${error?.message ?? error}`);
+            }
+            // Transient: the endpoint is momentarily unreachable, not the token bad.
+            if (attempt < AUTH_MAX_ATTEMPTS) {
+                const backoff = AUTH_RETRY_BASE_MS * 2 ** (attempt - 1);
+                logger_js_1.default.log(null, logger_js_1.default.levels.warn, `session validation transient failure (attempt ${attempt}/${AUTH_MAX_ATTEMPTS}); retrying in ${backoff}ms: ${error?.message ?? error}`);
+                await delay(backoff);
+            }
+        }
     }
+    throw new BackendUnavailableError(`StackFactor session validation unavailable after ${AUTH_MAX_ATTEMPTS} attempts: ${lastError?.message ?? lastError}`);
 };
 const safeParse = (value, fallback) => {
     if (!value)
@@ -283,9 +338,12 @@ const serve = (main, options = {}) => {
             }
             catch (error) {
                 const message = error?.message ?? "Unauthenticated";
-                logger_js_1.default.log(null, logger_js_1.default.levels.warn, `agent Execute refused: ${message}`);
+                const unavailable = error instanceof BackendUnavailableError;
+                logger_js_1.default.log(null, logger_js_1.default.levels.warn, `agent Execute refused (${unavailable ? "unavailable" : "unauthenticated"}): ${message}`);
                 call.emit("error", {
-                    code: grpc.status.UNAUTHENTICATED,
+                    code: unavailable
+                        ? grpc.status.UNAVAILABLE
+                        : grpc.status.UNAUTHENTICATED,
                     details: message,
                 });
                 return;

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

@@ -1 +1 @@
-{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":"AA2BA;;;;;GAKG;AACH,MAAM,MAAM,YAAY,GAAG;IACzB,IAAI,EAAE,MAAM,CAAC;IACb,MAAM,EAAE;QAAE,CAAC,QAAQ,EAAE,MAAM,GAAG,MAAM,CAAA;KAAE,CAAC;CACxC,CAAC;;0CAE2C,GAAG,KAAG,IAAI;wBA+jB/C,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBA0BG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,iBACb,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;oCA4XF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,iBACE,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;sDAusBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,iBACE,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;0CA3+B8B,GAAG,KAAG,MAAM;+CApjB9C,YAAY,GAAG,IAAI,GAAG,SAAS,aAC7B,MAAM,kBACD,MAAM,UACd,GAAG,KACV,IAAI;mCAokBU,MAAM,mBACJ,MAAM,EAAE,KACxB,MAAM;;AAkhCT,wBASE"}
+{"version":3,"file":"langChain.d.ts","sourceRoot":"","sources":["../../src/langChain.ts"],"names":[],"mappings":"AAgCA;;;;;GAKG;AACH,MAAM,MAAM,YAAY,GAAG;IACzB,IAAI,EAAE,MAAM,CAAC;IACb,MAAM,EAAE;QAAE,CAAC,QAAQ,EAAE,MAAM,GAAG,MAAM,CAAA;KAAE,CAAC;CACxC,CAAC;;0CAE2C,GAAG,KAAG,IAAI;wBA6lB/C,MAAM,aACD,MAAM,gBACH,MAAM,SACb,GAAG,EAAE,kBACI,GAAG,UACX,GAAG,KACV,GAAG;sBA0BG,GAAG,UACF,MAAM,UACN,GAAG,eACC,QAAQ,GAAG,IAAI,iBACb,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;oCA4XF,MAAM,UACT,GAAG,UACH,GAAG,oBACO,GAAG,eACT,MAAM,eACN,MAAM,wBACG,OAAO,WACpB,GAAG,cACA,MAAM,UACV,GAAG,EAAE,iBACE,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;sDAstBF,MAAM,UACT,GAAG,UACH,MAAM,YACL,GAAG,iBACE,YAAY,GAAG,IAAI,KAChC,OAAO,CAAC,GAAG,CAAC;0CA1/B8B,GAAG,KAAG,MAAM;+CA/jB9C,YAAY,GAAG,IAAI,GAAG,SAAS,aAC7B,MAAM,kBACD,MAAM,UACd,GAAG,KACV,IAAI;mCA+kBU,MAAM,mBACJ,MAAM,EAAE,KACxB,MAAM;;AAiiCT,wBASE"}

package/dist/esm/langChain.js

@@ -51,7 +51,19 @@ const updateUsageTracker = (tracker, modelName, usage, config) => {
     if (!tracker.tokens || typeof tracker.tokens !== "object")
         tracker.tokens = {};
     const inputTokens = usage.input_tokens || 0;
-    const outputTokens = usage.output_tokens || 0;
+    const visibleOutputTokens = usage.output_tokens || 0;
+    const totalTokens = usage.total_tokens || 0;
+    // Recover reasoning/"thinking" tokens the provider leaves out of
+    // `output_tokens`. Gemini reports its thoughts only in `totalTokenCount`, so
+    // `total - input - visibleOutput` is the thinking output the caller was still
+    // billed for. Providers that already fold reasoning into `output_tokens`
+    // (Anthropic, OpenAI) report `total == input + output`, so this adds 0.
+    // Reasoning is charged at the output rate, so we treat it as output for both
+    // the token counters and the cost.
+    const reasoningTokens = totalTokens > 0
+        ? Math.max(0, totalTokens - inputTokens - visibleOutputTokens)
+        : 0;
+    const outputTokens = visibleOutputTokens + reasoningTokens;
     const inputRate = getModelRate(modelName, config, "input-token");
     const outputRate = getModelRate(modelName, config, "output-token");
     const addedCost = (inputTokens / 1_000_000) * inputRate +
@@ -117,7 +129,12 @@ const updateImageUsageTracker = (tracker, modelName, response, config, provider)
         imageOutputTokens = Number(um.candidatesTokenCount) || 0;
     }
     const textInputRate = getModelRate(modelName, config, "input-token");
-    const imageInputRate = getModelRate(modelName, config, "image-input-token");
+    // Reference-image input tokens are billed at the model's input rate. Configs
+    // that don't define a dedicated `<model>-image-input-token-costs` (the common
+    // case — image models charge all input at one rate and only differ on output)
+    // fall back to the text input rate rather than silently billing image input
+    // at $0.
+    const imageInputRate = getModelRate(modelName, config, "image-input-token") || textInputRate;
     const imageOutputRate = getModelRate(modelName, config, "image-output-token");
     const addedCost = (textInputTokens / 1_000_000) * textInputRate +
         (imageInputTokens / 1_000_000) * imageInputRate +
@@ -171,6 +188,7 @@ const extractUsageFromInvoke = (response) => {
         return {
             input_tokens: um.input_tokens || 0,
             output_tokens: um.output_tokens || 0,
+            total_tokens: um.total_tokens || 0,
         };
     }
     const rm = response.response_metadata;
@@ -178,12 +196,14 @@ const extractUsageFromInvoke = (response) => {
         return {
             input_tokens: rm.usage.input_tokens || rm.usage.prompt_tokens || 0,
             output_tokens: rm.usage.output_tokens || rm.usage.completion_tokens || 0,
+            total_tokens: rm.usage.total_tokens || rm.usage.total_token_count || 0,
         };
     }
     if (rm?.tokenUsage) {
         return {
             input_tokens: rm.tokenUsage.promptTokens || 0,
             output_tokens: rm.tokenUsage.completionTokens || 0,
+            total_tokens: rm.tokenUsage.totalTokens || 0,
         };
     }
     return null;
@@ -200,6 +220,7 @@ const accumulateChunkUsage = (acc, chunk) => {
     if (um) {
         acc.input_tokens += um.input_tokens || 0;
         acc.output_tokens += um.output_tokens || 0;
+        acc.total_tokens += um.total_tokens || 0;
         return acc;
     }
     const rm = chunk.response_metadata;
@@ -207,6 +228,8 @@ const accumulateChunkUsage = (acc, chunk) => {
         acc.input_tokens += rm.usage.input_tokens || rm.usage.prompt_tokens || 0;
         acc.output_tokens +=
             rm.usage.output_tokens || rm.usage.completion_tokens || 0;
+        acc.total_tokens +=
+            rm.usage.total_tokens || rm.usage.total_token_count || 0;
     }
     return acc;
 };
@@ -219,13 +242,18 @@ const sumAgentResponseUsage = (response) => {
     const messages = response?.messages;
     if (!Array.isArray(messages) || messages.length === 0)
         return null;
-    const total = { input_tokens: 0, output_tokens: 0 };
+    const total = {
+        input_tokens: 0,
+        output_tokens: 0,
+        total_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;
+            total.total_tokens += um.total_tokens || 0;
             found = true;
         }
     }
@@ -1049,7 +1077,11 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
         while (true) {
             let rawContent = "";
             let chunkCount = 0;
-            let streamUsage = { input_tokens: 0, output_tokens: 0 };
+            let streamUsage = {
+                input_tokens: 0,
+                output_tokens: 0,
+                total_tokens: 0,
+            };
             // Inner loop: wait + retry on 429 around stream setup and consumption.
             // Usage is only recorded on a successful stream — partial streams that
             // error out with a rate limit are not counted. A 429 fired mid-stream
@@ -1058,7 +1090,7 @@ const runPromptWithModel = async (modelName, config, prompt, onProgressReport, m
             while (true) {
                 rawContent = "";
                 chunkCount = 0;
-                streamUsage = { input_tokens: 0, output_tokens: 0 };
+                streamUsage = { input_tokens: 0, output_tokens: 0, total_tokens: 0 };
                 try {
                     // Honour caller cancellation: passing the signal tears down the
                     // upstream HTTP request so a cancelled call stops billing tokens.
@@ -1334,20 +1366,21 @@ const generateImageWithOpenAI = async (modelName, config, prompt, options, usage
  * Google models do not support alpha channels. Validates the `aspectRatio` against
  * allowed values (`"1:1"`, `"3:4"`, `"4:3"`, `"9:16"`, `"16:9"`). Applies model-
  * specific generation config: Imagen models receive `numberOfImages`, `aspectRatio`,
- * `outputMimeType`, and optional `negativePrompt`; Gemini models receive `temperature`
- * and `topP`. Safety settings are disabled for image generation. Throws an error if no
- * images are returned.
+ * `outputMimeType`, and optional `negativePrompt`; Gemini models receive `temperature`,
+ * `topP`, and an `imageConfig` carrying `aspectRatio` plus an optional `imageSize`
+ * resolution tier (e.g. "1K"/"2K"/"4K", honoured by gemini-3-pro-image). Safety
+ * settings are disabled for image generation. Throws an error if no images are returned.
  * @param modelName - The Google model identifier (e.g. `"imagen-4.0-generate-001"`,
  *   `"gemini-3.0-pro-image"`)
  * @param config - Configuration object; must include `googleAPIKey`
  * @param prompt - The text prompt describing the image to generate
- * @param options - Generation options including `aspectRatio`, `numberOfImages`, and
- *   optional `negativePrompt` (Imagen only)
+ * @param options - Generation options including `aspectRatio`, optional `imageSize`
+ *   (Gemini resolution tier), `numberOfImages`, and optional `negativePrompt` (Imagen only)
  * @returns A single image descriptor `{ b64_json, mimeType }` when one image is
  *   requested, or `{ images: [...] }` for multiple images
  */
 const generateImageWithGoogle = async (modelName, config, prompt, options, usageTracker = null) => {
-    const { aspectRatio = "1:1", numberOfImages = 1, negativePrompt = "", } = options;
+    const { aspectRatio = "1:1", imageSize, numberOfImages = 1, negativePrompt = "", } = options;
     if (!config.googleAPIKey) {
         throw errorHandlingHelper.create(constants.HTTP_CODES.BAD_REQUEST, "Google API key is required for Google image generation");
     }
@@ -1384,6 +1417,15 @@ const generateImageWithGoogle = async (modelName, config, prompt, options, usage
                 // Gemini-specific config
                 temperature: 1,
                 topP: 0.95,
+                // Gemini controls output dimensions through `imageConfig`, not the
+                // top-level `aspectRatio` field that Imagen uses. `imageSize`
+                // (e.g. "1K"/"2K"/"4K") is only honoured by resolution-capable
+                // models such as gemini-3-pro-image, so it is passed through only
+                // when the caller explicitly requests it.
+                imageConfig: {
+                    aspectRatio: aspectRatio,
+                    ...(imageSize ? { imageSize: imageSize } : {}),
+                },
             }),
     };
     // Safety settings (disable for image generation)

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

@@ -1 +1 @@
-{"version":3,"file":"serve.d.ts","sourceRoot":"","sources":["../../src/serve.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,IAAI,MAAM,eAAe,CAAC;AA8LtC;;;;;GAKG;AACH,MAAM,MAAM,SAAS,GAAG,CAAC,MAAM,CAAC,EAAE,WAAW,KAAK,OAAO,CAAC,GAAG,CAAC,CAAC;AAE/D;;;;GAIG;AACH,MAAM,MAAM,aAAa,GAAG,CAC1B,OAAO,EAAE,GAAG,KACT,OAAO,CAAC;IAAE,IAAI,EAAE,GAAG,CAAC;IAAC,MAAM,EAAE,GAAG,CAAA;CAAE,CAAC,CAAC;AAEzC,MAAM,WAAW,YAAY;IAC3B,4EAA4E;IAC5E,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,2CAA2C;IAC3C,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,wFAAwF;IACxF,YAAY,CAAC,EAAE,aAAa,CAAC;CAC9B;AAkDD;;;;;;;;;;GAUG;AACH,eAAO,MAAM,KAAK,GAChB,MAAM,SAAS,EACf,UAAS,YAAiB,KACzB,IAAI,CAAC,MAuKP,CAAC"}
+{"version":3,"file":"serve.d.ts","sourceRoot":"","sources":["../../src/serve.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,IAAI,MAAM,eAAe,CAAC;AA8LtC;;;;;GAKG;AACH,MAAM,MAAM,SAAS,GAAG,CAAC,MAAM,CAAC,EAAE,WAAW,KAAK,OAAO,CAAC,GAAG,CAAC,CAAC;AAE/D;;;;GAIG;AACH,MAAM,MAAM,aAAa,GAAG,CAC1B,OAAO,EAAE,GAAG,KACT,OAAO,CAAC;IAAE,IAAI,EAAE,GAAG,CAAC;IAAC,MAAM,EAAE,GAAG,CAAA;CAAE,CAAC,CAAC;AAEzC,MAAM,WAAW,YAAY;IAC3B,4EAA4E;IAC5E,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,2CAA2C;IAC3C,IAAI,CAAC,EAAE,MAAM,CAAC;IACd,wFAAwF;IACxF,YAAY,CAAC,EAAE,aAAa,CAAC;CAC9B;AAyHD;;;;;;;;;;GAUG;AACH,eAAO,MAAM,KAAK,GAChB,MAAM,SAAS,EACf,UAAS,YAAiB,KACzB,IAAI,CAAC,MA4KP,CAAC"}

package/dist/esm/serve.js

@@ -160,23 +160,78 @@ const runSelfCheck = async (checkCode, config, session) => {
 /** Marks an auth failure so the handler can map it to gRPC UNAUTHENTICATED. */
 class UnauthenticatedError extends Error {
 }
+/**
+ * Marks a *transient* failure to reach the StackFactor session-validation
+ * endpoint (5xx / request timeout / transport error). Unlike
+ * {@link UnauthenticatedError} the token was never actually rejected — the
+ * backend was momentarily unreachable — so the handler maps this to the
+ * retryable gRPC UNAVAILABLE instead of UNAUTHENTICATED, and the message no
+ * longer slanders a perfectly valid token.
+ */
+class BackendUnavailableError extends Error {
+}
+/**
+ * Session-validation retry policy. A GKE front-door blip (a backend pod briefly
+ * refusing connections during a restart or an event-loop stall) surfaces here as
+ * a 502/503 or a transport error, not a 401 — and retrying on a fresh connection
+ * almost always lands on a healthy replica. Env-overridable for prod tuning.
+ */
+const AUTH_MAX_ATTEMPTS = Math.max(1, Number(process.env.STACKFACTOR_AUTH_MAX_ATTEMPTS) || 4);
+const AUTH_RETRY_BASE_MS = Math.max(0, Number(process.env.STACKFACTOR_AUTH_RETRY_BASE_MS) || 500);
+const delay = (ms) => new Promise((res) => setTimeout(res, ms));
+/** A 401/403 is the backend genuinely rejecting the token. */
+const isAuthRejection = (error) => {
+    const status = error?.response?.status;
+    return status === 401 || status === 403;
+};
+/**
+ * True when the failure is the validation endpoint being unreachable rather than
+ * the token being bad: a 5xx/408/429 HTTP response, or a transport-level error
+ * (no HTTP response but an axios/errno signature). Plain programming errors
+ * (no response and no axios/errno marker) are NOT treated as transient.
+ */
+const isTransient = (error) => {
+    const status = error?.response?.status;
+    if (status)
+        return status >= 500 || status === 408 || status === 429;
+    return Boolean(error?.isAxiosError || error?.code);
+};
 /**
  * Default authenticator: require a token on the request and resolve it to a
  * session via the StackFactor API. Requires `BACKEND_URL` (or `REACT_APP_NODE_ENV`)
- * to point client-api at the right backend. Throws {@link UnauthenticatedError}
- * when the token is missing or rejected.
+ * to point client-api at the right backend.
+ *
+ * A 401/403 (or any non-transient error) throws {@link UnauthenticatedError}
+ * immediately. A transient failure (5xx / timeout / network) is retried with
+ * exponential backoff; if every attempt fails it throws
+ * {@link BackendUnavailableError} so the caller sees a retryable UNAVAILABLE
+ * rather than a misleading "invalid token".
  */
 const defaultAuthenticate = async (request) => {
     const token = request?.authToken ?? request?.authorization;
     if (!token) {
         throw new UnauthenticatedError("Missing StackFactor auth token");
     }
-    try {
-        return await clientSession.getSession(token);
-    }
-    catch (error) {
-        throw new UnauthenticatedError(`Invalid StackFactor auth token: ${error?.message ?? error}`);
+    let lastError;
+    for (let attempt = 1; attempt <= AUTH_MAX_ATTEMPTS; attempt++) {
+        try {
+            return await clientSession.getSession(token);
+        }
+        catch (error) {
+            lastError = error;
+            // Real rejection, or an unexpected non-transient error: fail fast.
+            if (isAuthRejection(error) || !isTransient(error)) {
+                throw new UnauthenticatedError(`Invalid StackFactor auth token: ${error?.message ?? error}`);
+            }
+            // Transient: the endpoint is momentarily unreachable, not the token bad.
+            if (attempt < AUTH_MAX_ATTEMPTS) {
+                const backoff = AUTH_RETRY_BASE_MS * 2 ** (attempt - 1);
+                logger.log(null, logger.levels.warn, `session validation transient failure (attempt ${attempt}/${AUTH_MAX_ATTEMPTS}); retrying in ${backoff}ms: ${error?.message ?? error}`);
+                await delay(backoff);
+            }
+        }
     }
+    throw new BackendUnavailableError(`StackFactor session validation unavailable after ${AUTH_MAX_ATTEMPTS} attempts: ${lastError?.message ?? lastError}`);
 };
 const safeParse = (value, fallback) => {
     if (!value)
@@ -244,9 +299,12 @@ export const serve = (main, options = {}) => {
             }
             catch (error) {
                 const message = error?.message ?? "Unauthenticated";
-                logger.log(null, logger.levels.warn, `agent Execute refused: ${message}`);
+                const unavailable = error instanceof BackendUnavailableError;
+                logger.log(null, logger.levels.warn, `agent Execute refused (${unavailable ? "unavailable" : "unauthenticated"}): ${message}`);
                 call.emit("error", {
-                    code: grpc.status.UNAUTHENTICATED,
+                    code: unavailable
+                        ? grpc.status.UNAVAILABLE
+                        : grpc.status.UNAUTHENTICATED,
                     details: message,
                 });
                 return;

package/package.json

@@ -3,7 +3,7 @@
   "publishConfig": {
     "access": "public"
   },
-  "version": "1.2.8",
+  "version": "1.2.11",
   "description": "",
   "main": "dist/cjs/index.js",
   "module": "dist/esm/index.js",