@juspay/neurolink 10.6.4 → 10.6.5

package/dist/lib/providers/googleVertex.js

@@ -1469,11 +1469,14 @@ export class GoogleVertexProvider extends BaseProvider {
         // into a provider "prompt too long" rejection mid-loop.
         const contextGuard = createContextGuard(getContextWindowSize("vertex", modelName));
         let hitContextLimit = false;
-        // Track token usage across all steps
-        // promptTokenCount is typically in the final chunk, candidatesTokenCount accumulates
+        // Track token usage across all steps. Every loop step is a separate
+        // billed API call, so per-step usage is folded into these turn totals
+        // after each step's chunk drain — assigning them directly from chunk
+        // metadata would record only the last step of a multi-step tool loop.
         let totalInputTokens = 0;
         let totalOutputTokens = 0;
         let totalCacheReadTokens = 0;
+        let totalReasoningTokens = 0;
         // Track text parts as they arrive from the SDK so the returned async
         // iterable yields multiple chunks instead of a single buffered chunk.
         // The CLI's chunk-count smoke test asserts > 1 stream chunks for any
@@ -1557,6 +1560,18 @@ export class GoogleVertexProvider extends BaseProvider {
                     // lastFinishReason) — drives the single MALFORMED_FUNCTION_CALL
                     // retry below.
                     let stepFinishReason;
+                    // Per-step usage trackers for WRITE-THROUGH accumulation: within a
+                    // step's chunk stream the counts are latest-wins (promptTokenCount
+                    // arrives once in the final chunk; candidates/thoughts counts are
+                    // cumulative across chunks), so each chunk folds only the DELTA
+                    // over this step's previous value into the turn totals. The totals
+                    // are therefore correct at every point mid-drain — a step killed
+                    // mid-stream (abort / turn deadline / stall watchdog) still counts
+                    // the billed tokens it already reported.
+                    let stepInputTokens = 0;
+                    let stepOutputTokens = 0;
+                    let stepCacheReadTokens = 0;
+                    let stepReasoningTokens = 0;
                     for await (const chunk of stream) {
                         turnClock.noteProgress();
                         // Extract raw parts from candidates FIRST
@@ -1591,18 +1606,34 @@ export class GoogleVertexProvider extends BaseProvider {
                             // Take the latest promptTokenCount (usually only in final chunk)
                             if (usageMetadata.promptTokenCount !== undefined &&
                                 usageMetadata.promptTokenCount > 0) {
-                                totalInputTokens = usageMetadata.promptTokenCount;
+                                totalInputTokens +=
+                                    usageMetadata.promptTokenCount - stepInputTokens;
+                                stepInputTokens = usageMetadata.promptTokenCount;
                                 // Feed the context guard the REAL prompt size of this call.
                                 contextGuard.noteUsage(usageMetadata.promptTokenCount, usageMetadata.candidatesTokenCount ?? 0);
                                 // cachedContentTokenCount is OVERLAPPING (a subset already inside
-                                // promptTokenCount). Clamp to the prompt count so a later uncached
-                                // step resets it to 0 instead of leaving a stale cached value.
-                                totalCacheReadTokens = Math.min(usageMetadata.cachedContentTokenCount ?? 0, usageMetadata.promptTokenCount);
+                                // promptTokenCount). Clamp to the prompt count so an uncached
+                                // step reports 0 instead of a stale cached value.
+                                const chunkCacheReadTokens = Math.min(usageMetadata.cachedContentTokenCount ?? 0, usageMetadata.promptTokenCount);
+                                totalCacheReadTokens +=
+                                    chunkCacheReadTokens - stepCacheReadTokens;
+                                stepCacheReadTokens = chunkCacheReadTokens;
                             }
                             // Take the latest candidatesTokenCount (accumulates through chunks)
                             if (usageMetadata.candidatesTokenCount !== undefined &&
                                 usageMetadata.candidatesTokenCount > 0) {
-                                totalOutputTokens = usageMetadata.candidatesTokenCount;
+                                totalOutputTokens +=
+                                    usageMetadata.candidatesTokenCount - stepOutputTokens;
+                                stepOutputTokens = usageMetadata.candidatesTokenCount;
+                            }
+                            // thoughtsTokenCount (thinking tokens, billed at the output
+                            // rate) is NOT part of candidatesTokenCount — Gemini reports
+                            // totalTokenCount = prompt + candidates + thoughts.
+                            if (usageMetadata.thoughtsTokenCount !== undefined &&
+                                usageMetadata.thoughtsTokenCount > 0) {
+                                totalReasoningTokens +=
+                                    usageMetadata.thoughtsTokenCount - stepReasoningTokens;
+                                stepReasoningTokens = usageMetadata.thoughtsTokenCount;
                             }
                         }
                     }
@@ -1991,6 +2022,7 @@ export class GoogleVertexProvider extends BaseProvider {
                         incrementalTextChunks.push(synth.text);
                         totalInputTokens += synth.inputTokens;
                         totalOutputTokens += synth.outputTokens;
+                        totalReasoningTokens += synth.reasoningTokens ?? 0;
                         if (synth.finishReason) {
                             lastFinishReason = synth.finishReason;
                         }
@@ -2069,11 +2101,22 @@ export class GoogleVertexProvider extends BaseProvider {
             rawFinishReason: lastFinishReason,
             usage: {
                 input: adjustedInputTokens,
-                output: totalOutputTokens,
-                total: adjustedInputTokens + totalCacheReadTokens + totalOutputTokens,
+                // Thinking tokens are billed at the output rate but Gemini does NOT
+                // include them in candidatesTokenCount (totalTokenCount = prompt +
+                // candidates + thoughts), so they are folded into `output` — that is
+                // what calculateCost bills at the output rate — with `reasoning`
+                // reporting the thinking subset.
+                output: totalOutputTokens + totalReasoningTokens,
+                total: adjustedInputTokens +
+                    totalCacheReadTokens +
+                    totalOutputTokens +
+                    totalReasoningTokens,
                 ...(totalCacheReadTokens > 0 && {
                     cacheReadTokens: totalCacheReadTokens,
                 }),
+                ...(totalReasoningTokens > 0 && {
+                    reasoning: totalReasoningTokens,
+                }),
             },
             toolCalls: externalToolCalls.map((tc) => ({
                 toolName: tc.toolName,
@@ -2409,11 +2452,14 @@ export class GoogleVertexProvider extends BaseProvider {
         // into a provider "prompt too long" rejection mid-loop.
         const contextGuard = createContextGuard(getContextWindowSize("vertex", modelName));
         let hitContextLimit = false;
-        // Track token usage across all steps
-        // promptTokenCount is typically in the final chunk, candidatesTokenCount accumulates
+        // Track token usage across all steps. Every loop step is a separate
+        // billed API call, so per-step usage is folded into these turn totals
+        // after each step's chunk drain — assigning them directly from chunk
+        // metadata would record only the last step of a multi-step tool loop.
         let totalInputTokens = 0;
         let totalOutputTokens = 0;
         let totalCacheReadTokens = 0;
+        let totalReasoningTokens = 0;
         // Abort scaffolding (mirrors executeNativeAnthropicStream). The native
         // Gemini SDK cancels via config.abortSignal, so drive an internal
         // AbortController: the caller's signal and the turn clock's watchdogs
@@ -2488,6 +2534,18 @@ export class GoogleVertexProvider extends BaseProvider {
                     // lastFinishReason) — drives the single MALFORMED_FUNCTION_CALL
                     // retry below.
                     let stepFinishReason;
+                    // Per-step usage trackers for WRITE-THROUGH accumulation: within a
+                    // step's chunk stream the counts are latest-wins (promptTokenCount
+                    // arrives once in the final chunk; candidates/thoughts counts are
+                    // cumulative across chunks), so each chunk folds only the DELTA
+                    // over this step's previous value into the turn totals. The totals
+                    // are therefore correct at every point mid-drain — a step killed
+                    // mid-stream (abort / turn deadline / stall watchdog) still counts
+                    // the billed tokens it already reported.
+                    let stepInputTokens = 0;
+                    let stepOutputTokens = 0;
+                    let stepCacheReadTokens = 0;
+                    let stepReasoningTokens = 0;
                     // Collect all chunks from stream
                     for await (const chunk of stream) {
                         turnClock.noteProgress();
@@ -2518,18 +2576,34 @@ export class GoogleVertexProvider extends BaseProvider {
                             // Take the latest promptTokenCount (usually only in final chunk)
                             if (usageMetadata.promptTokenCount !== undefined &&
                                 usageMetadata.promptTokenCount > 0) {
-                                totalInputTokens = usageMetadata.promptTokenCount;
+                                totalInputTokens +=
+                                    usageMetadata.promptTokenCount - stepInputTokens;
+                                stepInputTokens = usageMetadata.promptTokenCount;
                                 // Feed the context guard the REAL prompt size of this call.
                                 contextGuard.noteUsage(usageMetadata.promptTokenCount, usageMetadata.candidatesTokenCount ?? 0);
                                 // cachedContentTokenCount is OVERLAPPING (a subset already inside
-                                // promptTokenCount). Clamp to the prompt count so a later uncached
-                                // step resets it to 0 instead of leaving a stale cached value.
-                                totalCacheReadTokens = Math.min(usageMetadata.cachedContentTokenCount ?? 0, usageMetadata.promptTokenCount);
+                                // promptTokenCount). Clamp to the prompt count so an uncached
+                                // step reports 0 instead of a stale cached value.
+                                const chunkCacheReadTokens = Math.min(usageMetadata.cachedContentTokenCount ?? 0, usageMetadata.promptTokenCount);
+                                totalCacheReadTokens +=
+                                    chunkCacheReadTokens - stepCacheReadTokens;
+                                stepCacheReadTokens = chunkCacheReadTokens;
                             }
                             // Take the latest candidatesTokenCount (accumulates through chunks)
                             if (usageMetadata.candidatesTokenCount !== undefined &&
                                 usageMetadata.candidatesTokenCount > 0) {
-                                totalOutputTokens = usageMetadata.candidatesTokenCount;
+                                totalOutputTokens +=
+                                    usageMetadata.candidatesTokenCount - stepOutputTokens;
+                                stepOutputTokens = usageMetadata.candidatesTokenCount;
+                            }
+                            // thoughtsTokenCount (thinking tokens, billed at the output
+                            // rate) is NOT part of candidatesTokenCount — Gemini reports
+                            // totalTokenCount = prompt + candidates + thoughts.
+                            if (usageMetadata.thoughtsTokenCount !== undefined &&
+                                usageMetadata.thoughtsTokenCount > 0) {
+                                totalReasoningTokens +=
+                                    usageMetadata.thoughtsTokenCount - stepReasoningTokens;
+                                stepReasoningTokens = usageMetadata.thoughtsTokenCount;
                             }
                         }
                     }
@@ -2912,6 +2986,7 @@ export class GoogleVertexProvider extends BaseProvider {
                         finalText = synth.text;
                         totalInputTokens += synth.inputTokens;
                         totalOutputTokens += synth.outputTokens;
+                        totalReasoningTokens += synth.reasoningTokens ?? 0;
                         if (synth.finishReason) {
                             lastFinishReason = synth.finishReason;
                         }
@@ -2974,12 +3049,26 @@ export class GoogleVertexProvider extends BaseProvider {
             stepsUsed: step,
             usage: {
                 input: adjustedInputTokens,
-                output: totalOutputTokens,
-                total: adjustedInputTokens + totalCacheReadTokens + totalOutputTokens,
+                // Thinking tokens are billed at the output rate but Gemini does NOT
+                // include them in candidatesTokenCount (totalTokenCount = prompt +
+                // candidates + thoughts), so they are folded into `output` — that is
+                // what calculateCost bills at the output rate — with `reasoning`
+                // reporting the thinking subset.
+                output: totalOutputTokens + totalReasoningTokens,
+                total: adjustedInputTokens +
+                    totalCacheReadTokens +
+                    totalOutputTokens +
+                    totalReasoningTokens,
                 ...(totalCacheReadTokens > 0 && {
                     cacheReadTokens: totalCacheReadTokens,
                 }),
+                ...(totalReasoningTokens > 0 && {
+                    reasoning: totalReasoningTokens,
+                }),
             },
+            ...(totalReasoningTokens > 0 && {
+                reasoningTokens: totalReasoningTokens,
+            }),
             responseTime,
             toolsUsed: externalToolCalls.map((tc) => tc.toolName),
             toolExecutions: resolveToolExecutionRecords(options, externalToolExecutions),
@@ -3016,7 +3105,7 @@ export class GoogleVertexProvider extends BaseProvider {
         // Already aborted — never issue the synth request (would add +300s after a
         // blown budget). Return empty so the caller maps to the graceful message.
         if (abortSignal?.aborted) {
-            return { text: "", inputTokens: 0, outputTokens: 0 };
+            return { text: "", inputTokens: 0, outputTokens: 0, reasoningTokens: 0 };
         }
         try {
             // Shallow clone so the loop's config is never mutated; dropping the
@@ -3053,6 +3142,7 @@ export class GoogleVertexProvider extends BaseProvider {
                 let finishReason;
                 let inputTokens = 0;
                 let outputTokens = 0;
+                let reasoningTokens = 0;
                 for await (const chunk of stream) {
                     const chunkRecord = chunk;
                     const candidates = chunkRecord.candidates;
@@ -3075,18 +3165,28 @@ export class GoogleVertexProvider extends BaseProvider {
                             usageMetadata.candidatesTokenCount > 0) {
                             outputTokens = usageMetadata.candidatesTokenCount;
                         }
+                        if (usageMetadata.thoughtsTokenCount !== undefined &&
+                            usageMetadata.thoughtsTokenCount > 0) {
+                            reasoningTokens = usageMetadata.thoughtsTokenCount;
+                        }
                     }
                 }
                 const text = parts
                     .filter((part) => typeof part.text === "string")
                     .map((part) => part.text)
                     .join("");
-                return { text, finishReason, inputTokens, outputTokens };
+                return {
+                    text,
+                    finishReason,
+                    inputTokens,
+                    outputTokens,
+                    reasoningTokens,
+                };
             })(), timeoutMs, "Gemini synthesis call timed out");
         }
         catch (error) {
             logger.warn("[GoogleVertex] Tools-disabled synthesis call failed; falling back to placeholder", { error: error instanceof Error ? error.message : String(error) });
-            return { text: "", inputTokens: 0, outputTokens: 0 };
+            return { text: "", inputTokens: 0, outputTokens: 0, reasoningTokens: 0 };
         }
     }
     /**
@@ -3651,7 +3751,14 @@ export class GoogleVertexProvider extends BaseProvider {
                         turnCacheUsage.creation1h += stepCacheCreation1h;
                         usage.input += response.usage?.input_tokens || 0;
                         usage.output += response.usage?.output_tokens || 0;
-                        usage.total = usage.input + usage.output;
+                        // Anthropic's input_tokens is only the UNCACHED remainder; cache
+                        // reads/writes are billed tokens reported separately, so the
+                        // total must include them (matches proxyTracer + anthropic.ts).
+                        usage.total =
+                            usage.input +
+                                usage.output +
+                                turnCacheUsage.read +
+                                turnCacheUsage.creation;
                         lastStopReason = response.stop_reason;
                         // Feed the context guard the FULL prompt size of this call
                         // (uncached input + cache reads/writes).
@@ -4102,7 +4209,6 @@ export class GoogleVertexProvider extends BaseProvider {
                             activeStream = undefined;
                             usage.input += response.usage?.input_tokens || 0;
                             usage.output += response.usage?.output_tokens || 0;
-                            usage.total = usage.input + usage.output;
                             turnCacheUsage.read +=
                                 response.usage?.cache_read_input_tokens ?? 0;
                             turnCacheUsage.creation +=
@@ -4111,6 +4217,13 @@ export class GoogleVertexProvider extends BaseProvider {
                                 response.usage?.cache_creation?.ephemeral_5m_input_tokens ?? 0;
                             turnCacheUsage.creation1h +=
                                 response.usage?.cache_creation?.ephemeral_1h_input_tokens ?? 0;
+                            // Total counts cache reads/writes — input_tokens is only the
+                            // uncached remainder.
+                            usage.total =
+                                usage.input +
+                                    usage.output +
+                                    turnCacheUsage.read +
+                                    turnCacheUsage.creation;
                             lastStopReason = response.stop_reason;
                             const forcedFinalResult = response.content.find((block) => block.type === "tool_use" && block.name === "final_result");
                             if (forcedFinalResult) {
@@ -4196,7 +4309,6 @@ export class GoogleVertexProvider extends BaseProvider {
                                 activeStream = undefined;
                                 usage.input += response.usage?.input_tokens || 0;
                                 usage.output += response.usage?.output_tokens || 0;
-                                usage.total = usage.input + usage.output;
                                 turnCacheUsage.read +=
                                     response.usage?.cache_read_input_tokens ?? 0;
                                 turnCacheUsage.creation +=
@@ -4207,6 +4319,13 @@ export class GoogleVertexProvider extends BaseProvider {
                                 turnCacheUsage.creation1h +=
                                     response.usage?.cache_creation?.ephemeral_1h_input_tokens ??
                                         0;
+                                // Total counts cache reads/writes — input_tokens is only the
+                                // uncached remainder.
+                                usage.total =
+                                    usage.input +
+                                        usage.output +
+                                        turnCacheUsage.read +
+                                        turnCacheUsage.creation;
                                 lastStopReason = response.stop_reason;
                                 const backstopText = response.content
                                     .filter((block) => block.type === "text")
@@ -5408,7 +5527,13 @@ export class GoogleVertexProvider extends BaseProvider {
             usage: {
                 input: totalInputTokens,
                 output: totalOutputTokens,
-                total: totalInputTokens + totalOutputTokens,
+                // Anthropic's input_tokens is only the UNCACHED remainder; cache
+                // reads/writes are billed tokens reported separately, so the total
+                // must include them (matches proxyTracer + anthropic.ts).
+                total: totalInputTokens +
+                    totalOutputTokens +
+                    totalCacheReadTokens +
+                    totalCacheCreationTokens,
                 ...(totalCacheReadTokens > 0 && {
                     cacheReadTokens: totalCacheReadTokens,
                 }),

package/dist/lib/providers/litellm.js

@@ -283,12 +283,23 @@ export class LiteLLMProvider extends OpenAIChatCompletionsProvider {
         };
         return {
             onUsage: (usage) => {
-                span.setAttribute("gen_ai.usage.input_tokens", usage.promptTokens);
+                // promptTokens is the uncached remainder — the attribute reports the
+                // full prompt (uncached + cache read/creation), and the cache fields
+                // let calculateCost price tiers.
+                span.setAttribute("gen_ai.usage.input_tokens", usage.promptTokens +
+                    (usage.cacheReadTokens ?? 0) +
+                    (usage.cacheCreationTokens ?? 0));
                 span.setAttribute("gen_ai.usage.output_tokens", usage.completionTokens);
-                const cost = calculateCost(this.providerName, this.modelName, {
+                const cost = calculateCost(this.providerName, modelId, {
                     input: usage.promptTokens,
                     output: usage.completionTokens,
                     total: usage.totalTokens,
+                    ...(usage.cacheReadTokens
+                        ? { cacheReadTokens: usage.cacheReadTokens }
+                        : {}),
+                    ...(usage.cacheCreationTokens
+                        ? { cacheCreationTokens: usage.cacheCreationTokens }
+                        : {}),
                 });
                 if (cost && cost > 0) {
                     span.setAttribute("neurolink.cost", cost);

package/dist/lib/providers/openAI.js

@@ -157,12 +157,23 @@ export class OpenAIProvider extends OpenAIChatCompletionsProvider {
         };
         return {
             onUsage: (usage) => {
-                span.setAttribute("gen_ai.usage.input_tokens", usage.promptTokens);
+                // promptTokens is the uncached remainder — the attribute reports the
+                // full prompt (uncached + cache read/creation), and the cache fields
+                // let calculateCost price tiers.
+                span.setAttribute("gen_ai.usage.input_tokens", usage.promptTokens +
+                    (usage.cacheReadTokens ?? 0) +
+                    (usage.cacheCreationTokens ?? 0));
                 span.setAttribute("gen_ai.usage.output_tokens", usage.completionTokens);
-                const cost = calculateCost(this.providerName, this.modelName, {
+                const cost = calculateCost(this.providerName, modelId, {
                     input: usage.promptTokens,
                     output: usage.completionTokens,
                     total: usage.totalTokens,
+                    ...(usage.cacheReadTokens
+                        ? { cacheReadTokens: usage.cacheReadTokens }
+                        : {}),
+                    ...(usage.cacheCreationTokens
+                        ? { cacheCreationTokens: usage.cacheCreationTokens }
+                        : {}),
                 });
                 if (cost && cost > 0) {
                     span.setAttribute("neurolink.cost", cost);

package/dist/lib/providers/openaiChatCompletionsBase.js

@@ -505,8 +505,21 @@ export class OpenAIChatCompletionsProvider extends BaseProvider {
                     usage: {
                         inputTokens: {
                             total: json.usage?.prompt_tokens,
-                            noCache: json.usage?.prompt_tokens,
-                            cacheRead: undefined,
+                            // cached_tokens is OVERLAPPING (a subset of prompt_tokens), so
+                            // noCache is the remainder and cacheRead is clamped to the
+                            // prompt total — some gateways report inconsistent usage where
+                            // cached_tokens exceeds prompt_tokens (mirrors the
+                            // reasoning_tokens clamp on the output side below).
+                            noCache: json.usage?.prompt_tokens !== undefined &&
+                                json.usage?.prompt_tokens_details?.cached_tokens !== undefined
+                                ? Math.max(0, json.usage.prompt_tokens -
+                                    json.usage.prompt_tokens_details.cached_tokens)
+                                : json.usage?.prompt_tokens,
+                            cacheRead: json.usage?.prompt_tokens !== undefined &&
+                                json.usage?.prompt_tokens_details?.cached_tokens !== undefined
+                                ? Math.min(json.usage.prompt_tokens_details.cached_tokens, json.usage.prompt_tokens)
+                                : json.usage?.prompt_tokens_details?.cached_tokens,
+                            // OpenAI-style APIs do not report cache writes.
                             cacheWrite: undefined,
                         },
                         outputTokens: {
@@ -727,9 +740,33 @@ export class OpenAIChatCompletionsProvider extends BaseProvider {
     }
     async runStreamLoop(args) {
         const { maxSteps, modelId, url, fetchImpl, abortSignal, options, conversation, openAITools, openAIToolChoice, toolsRecord, toolNameFromWire, emitter, toolsUsed, toolExecutionSummaries, pushChunk, resolveUsage, resolveFinish, } = args;
+        // Hoisted above the try so the catch can resolve the usage accumulated
+        // by steps that completed BEFORE the failure — those steps were billed.
+        let stepFinish = null;
+        let stepUsage;
+        // Aggregated DeferredUsage from whatever accumulated in stepUsage.
+        // prompt_tokens is OVERLAPPING with cached_tokens (a subset), so the
+        // uncached remainder goes in promptTokens and the cached part in
+        // cacheReadTokens — the non-overlapping convention extractTokenUsage
+        // and calculateCost expect. reasoning_tokens stays a subset of
+        // completionTokens (informational).
+        const toDeferredUsage = () => {
+            const promptTokens = stepUsage?.prompt_tokens ?? 0;
+            const completionTokens = stepUsage?.completion_tokens ?? 0;
+            const cachedTokens = Math.min(stepUsage?.prompt_tokens_details?.cached_tokens ?? 0, promptTokens);
+            // Clamped like cached_tokens above: reasoning is a SUBSET of
+            // completion_tokens, but some gateways report inconsistent values.
+            const reasoningTokens = Math.min(Math.max(0, stepUsage?.completion_tokens_details?.reasoning_tokens ?? 0), completionTokens);
+            return {
+                promptTokens: promptTokens - cachedTokens,
+                completionTokens,
+                // Gateways that omit total_tokens must not collapse to 0.
+                totalTokens: stepUsage?.total_tokens || promptTokens + completionTokens,
+                ...(cachedTokens > 0 ? { cacheReadTokens: cachedTokens } : {}),
+                ...(reasoningTokens > 0 ? { reasoningTokens } : {}),
+            };
+        };
         try {
-            let stepFinish = null;
-            let stepUsage;
             // May grow mid-turn: hydrated tools with wire-unsafe names need
             // reverse-mapping even when the initial name set required none.
             let effectiveToolNameFromWire = toolNameFromWire;
@@ -787,11 +824,7 @@ export class OpenAIChatCompletionsProvider extends BaseProvider {
                     options,
                 });
             }
-            resolveUsage({
-                promptTokens: stepUsage?.prompt_tokens ?? 0,
-                completionTokens: stepUsage?.completion_tokens ?? 0,
-                totalTokens: stepUsage?.total_tokens ?? 0,
-            });
+            resolveUsage(toDeferredUsage());
             resolveFinish(stepFinish ?? "stop");
             pushChunk({ done: true });
             return {
@@ -803,7 +836,9 @@ export class OpenAIChatCompletionsProvider extends BaseProvider {
             logger.error(`${this.providerName}: Stream error`, {
                 error: err instanceof Error ? err.message : String(err),
             });
-            resolveUsage({ promptTokens: 0, completionTokens: 0, totalTokens: 0 });
+            // Steps that completed before the failure were billed — report them
+            // instead of zeroing the whole turn.
+            resolveUsage(toDeferredUsage());
             resolveFinish("error");
             pushChunk({ done: true });
             throw err;

package/dist/lib/providers/openaiChatCompletionsClient.d.ts

@@ -13,7 +13,7 @@
  * Nothing here imports from "ai" or "@ai-sdk/*". The whole point of this
  * module is to be the native replacement for the AI SDK's OpenAI wrapper.
  */
-import type { OpenAICompatBuildBodyArgs, OpenAICompatChatMessage, OpenAICompatChatRequest, OpenAICompatChatTool, OpenAICompatMessage, OpenAICompatMessageContent, OpenAICompatResponseFormat, OpenAICompatSSEResult, OpenAICompatStreamChunk, OpenAICompatToolChoiceWire, OpenAICompatV3CallToolChoice, OpenAICompatV3CallTools, DeferredUsage, Tool } from "../types/index.js";
+import type { OpenAICompatBuildBodyArgs, OpenAICompatChatMessage, OpenAICompatChatRequest, OpenAICompatChatTool, OpenAICompatMessage, OpenAICompatMessageContent, OpenAICompatResponseFormat, OpenAICompatSSEResult, OpenAICompatStreamChunk, OpenAICompatToolChoiceWire, OpenAICompatUsage, OpenAICompatV3CallToolChoice, OpenAICompatV3CallTools, DeferredUsage, Tool } from "../types/index.js";
 export declare const stripTrailingSlash: (s: string) => string;
 /**
  * Build a bijective original ↔ wire tool-name map. Returns undefined when
@@ -71,16 +71,4 @@ export declare const createChunkQueue: () => {
     pushChunk: (c: OpenAICompatStreamChunk) => void;
     nextChunk: () => Promise<OpenAICompatStreamChunk>;
 };
-export declare const mergeUsage: (a: {
-    prompt_tokens?: number;
-    completion_tokens?: number;
-    total_tokens?: number;
-} | undefined, b: {
-    prompt_tokens?: number;
-    completion_tokens?: number;
-    total_tokens?: number;
-} | undefined) => {
-    prompt_tokens?: number;
-    completion_tokens?: number;
-    total_tokens?: number;
-} | undefined;
+export declare const mergeUsage: (a: OpenAICompatUsage | undefined, b: OpenAICompatUsage | undefined) => OpenAICompatUsage | undefined;

package/dist/lib/providers/openaiChatCompletionsClient.js

@@ -673,10 +673,29 @@ export const mergeUsage = (a, b) => {
     if (!b) {
         return a;
     }
+    // Sum the nested details too — narrowing to the three flat scalars here
+    // silently dropped cached_tokens / reasoning_tokens the first time two
+    // steps were merged (single-step streams passed the object through
+    // untouched, so the loss only showed up in multi-step tool loops).
+    const cachedTokens = (a.prompt_tokens_details?.cached_tokens ?? 0) +
+        (b.prompt_tokens_details?.cached_tokens ?? 0);
+    const reasoningTokens = (a.completion_tokens_details?.reasoning_tokens ?? 0) +
+        (b.completion_tokens_details?.reasoning_tokens ?? 0);
+    // A step that omits total_tokens contributes its prompt+completion sum
+    // instead of 0 — otherwise one totals-reporting step next to one that
+    // omits it yields a nonzero-but-partial aggregate that downstream
+    // consumers trust over their own prompt+completion fallback.
+    const stepTotal = (u) => u.total_tokens || (u.prompt_tokens ?? 0) + (u.completion_tokens ?? 0);
     return {
         prompt_tokens: (a.prompt_tokens ?? 0) + (b.prompt_tokens ?? 0),
         completion_tokens: (a.completion_tokens ?? 0) + (b.completion_tokens ?? 0),
-        total_tokens: (a.total_tokens ?? 0) + (b.total_tokens ?? 0),
+        total_tokens: stepTotal(a) + stepTotal(b),
+        ...(cachedTokens > 0
+            ? { prompt_tokens_details: { cached_tokens: cachedTokens } }
+            : {}),
+        ...(reasoningTokens > 0
+            ? { completion_tokens_details: { reasoning_tokens: reasoningTokens } }
+            : {}),
     };
 };
 //# sourceMappingURL=openaiChatCompletionsClient.js.map

package/dist/lib/providers/sagemaker/language-model.js

@@ -7,7 +7,7 @@
 import { randomUUID } from "crypto";
 import { SageMakerRuntimeClient } from "./client.js";
 import { handleSageMakerError } from "./errors.js";
-import { estimateTokenUsage, createSageMakerStream } from "./streaming.js";
+import { estimateTokenUsage, createSageMakerStream, parseUsageFromResponseBody, } from "./streaming.js";
 import { createAdaptiveSemaphore } from "./adaptive-semaphore.js";
 import { logger } from "../../utils/logger.js";
 /**
@@ -140,8 +140,10 @@ export class SageMakerLanguageModel {
             const generatedText = this.extractTextFromResponse(responseBody);
             // Extract tool calls if present (Phase 4 enhancement)
             const toolCalls = this.extractToolCallsFromResponse(responseBody);
-            // Calculate token usage
-            const usage = estimateTokenUsage(promptText, generatedText);
+            // Prefer the endpoint's REAL token counts; only fall back to the
+            // char-heuristic estimate when the response reports none.
+            const usage = parseUsageFromResponseBody(responseBody) ??
+                estimateTokenUsage(promptText, generatedText);
             // Determine finish reason based on response content
             let finishReason = "stop";
             if (toolCalls && toolCalls.length > 0) {

package/dist/lib/providers/sagemaker/parsers.js

@@ -276,10 +276,13 @@ export class HuggingFaceStreamParser extends BaseStreamingParser {
             return undefined;
         }
         const tokens = details.tokens;
+        const promptTokens = Number(tokens.input) || 0;
+        const completionTokens = Number(tokens.generated) || 0;
         return {
-            promptTokens: Number(tokens.input) || 0,
-            completionTokens: Number(tokens.generated) || 0,
-            total: Number(tokens.total) || 0,
+            promptTokens,
+            completionTokens,
+            // Endpoints that omit the total must not report 0 next to real counts.
+            total: Number(tokens.total) || promptTokens + completionTokens,
         };
     }
     mapFinishReason(reason) {
@@ -447,10 +450,13 @@ export class LlamaStreamParser extends BaseStreamingParser {
         return toolCall;
     }
     parseLlamaUsage(usage) {
+        const promptTokens = Number(usage.prompt_tokens) || 0;
+        const completionTokens = Number(usage.completion_tokens) || 0;
         return {
-            promptTokens: Number(usage.prompt_tokens) || 0,
-            completionTokens: Number(usage.completion_tokens) || 0,
-            total: Number(usage.total_tokens) || 0,
+            promptTokens,
+            completionTokens,
+            // Endpoints that omit total_tokens must not report 0 next to real counts.
+            total: Number(usage.total_tokens) || promptTokens + completionTokens,
         };
     }
     mapFinishReason(reason) {
@@ -569,10 +575,13 @@ export class CustomStreamParser extends BaseStreamingParser {
     }
     parseCustomUsage(data) {
         const usage = (data.usage || data.tokens || {});
+        const promptTokens = Number(usage.prompt_tokens || usage.input_tokens) || 0;
+        const completionTokens = Number(usage.completion_tokens || usage.output_tokens) || 0;
         return {
-            promptTokens: Number(usage.prompt_tokens || usage.input_tokens) || 0,
-            completionTokens: Number(usage.completion_tokens || usage.output_tokens) || 0,
-            total: Number(usage.total_tokens) || 0,
+            promptTokens,
+            completionTokens,
+            // Endpoints that omit total_tokens must not report 0 next to real counts.
+            total: Number(usage.total_tokens) || promptTokens + completionTokens,
         };
     }
 }

package/dist/lib/providers/sagemaker/streaming.d.ts

@@ -29,6 +29,15 @@ export declare function createSyntheticStream(text: string, usage: SageMakerUsag
     onChunk?: (chunk: SageMakerStreamChunk) => void;
     onComplete?: (usage: SageMakerUsage) => void;
 }): Promise<ReadableStream<unknown>>;
+/**
+ * Extract REAL token usage from a complete (non-streaming) SageMaker
+ * response body when the endpoint reports it. Covers the OpenAI-compatible
+ * shape (vLLM / TGI-OpenAI: usage.prompt_tokens/completion_tokens), the
+ * generic input_tokens/output_tokens shape, and the HuggingFace TGI shape
+ * (details.tokens.input/generated). Returns undefined when no real counts
+ * are present so callers can fall back to estimateTokenUsage.
+ */
+export declare function parseUsageFromResponseBody(body: Record<string, unknown> | Array<Record<string, unknown>>): SageMakerUsage | undefined;
 /**
  * Estimate token usage from text content
  *