@juspay/neurolink 11.14.0 → 11.15.1

package/dist/core/geminiLoopAdapter.js

@@ -101,13 +101,19 @@ export function createGeminiLoopAdapter(config) {
         },
         async executeStep(request, channel, signal) {
             const rawStream = await config.sendStep(request.raw, signal);
-            // `collectStreamChunksIncremental` wants a StreamChannel but only ever
-            // calls `.push`, so the engine's push-only channel satisfies it. The
-            // helper is reused rather than reimplemented precisely because it also
-            // owns usage extraction and thought-signature preservation.
-            const collected = await collectStreamChunksIncremental(rawStream, {
-                push: (chunk) => channel.push(chunk),
-            });
+            // The shared helper by default: it owns usage extraction and
+            // thought-signature preservation, and it wants a StreamChannel but only
+            // ever calls `.push`, so the engine's push-only channel satisfies it.
+            //
+            // A provider whose drain genuinely differs supplies `collectStep`
+            // instead. Vertex does: it folds cumulative usage counts as deltas in
+            // its own loop, and that behaviour is characterized, so it keeps its
+            // collector rather than being quietly switched to this one.
+            const collected = config.collectStep
+                ? await config.collectStep(rawStream, channel)
+                : await collectStreamChunksIncremental(rawStream, {
+                    push: (chunk) => channel.push(chunk),
+                });
             // The provider's context guard calibrates from real per-step counts;
             // `inputTokens` is this step's full prompt size, which is what it
             // projects the next request from.

package/dist/providers/googleVertex/client.js

@@ -575,6 +575,101 @@ function reclaimVertexAnthropicContext(messages, modelName, observedPromptTokens
     messages.push(...rebuilt);
     return true;
 }
+/**
+ * Fold one Vertex Gemini step's stream into the shape the loop engine reports.
+ *
+ * Lifted from the inline drain in executeNativeGemini3Stream so the loop can
+ * later be handed to createGeminiLoopAdapter as `collectStep`. Vertex does not
+ * share googleNativeGemini3's collector: it reads parts straight off each
+ * candidate — avoiding the SDK warning that `chunk.text` raises when
+ * thoughtSignature or functionCall parts are present — and that behaviour is
+ * characterized.
+ *
+ * `onUsageDelta` fires PER CHUNK and is not an optimisation to fold away. The
+ * drain updates the turn totals incrementally so they are correct at every
+ * point mid-stream: a step killed by an abort, the turn deadline or the stall
+ * watchdog still bills the tokens it already reported. Returning a step total
+ * for the caller to add would be arithmetically identical and operationally
+ * wrong, because a killed step never returns.
+ */
+async function collectVertexStreamChunks(stream, channel, hooks = {}) {
+    const rawResponseParts = [];
+    const stepFunctionCalls = [];
+    let lastFinishReason;
+    let stepInputTokens = 0;
+    let stepOutputTokens = 0;
+    let stepCacheReadTokens = 0;
+    let stepReasoningTokens = 0;
+    for await (const chunk of stream) {
+        hooks.onProgress?.();
+        // Extract raw parts from candidates FIRST
+        // This avoids using chunk.text which triggers SDK warning when
+        // non-text parts (thoughtSignature, functionCall) are present
+        const chunkRecord = chunk;
+        const candidates = chunkRecord.candidates;
+        const firstCandidate = candidates?.[0];
+        // Capture the SDK finish reason (Bug 2: previously dropped). Last
+        // non-empty value across chunks wins.
+        const chunkFinishReason = firstCandidate?.finishReason;
+        if (typeof chunkFinishReason === "string" && chunkFinishReason) {
+            lastFinishReason = chunkFinishReason;
+        }
+        const chunkContent = firstCandidate?.content;
+        if (chunkContent && Array.isArray(chunkContent.parts)) {
+            for (const part of chunkContent.parts) {
+                rawResponseParts.push(part);
+                if (typeof part.text === "string" && part.text.length > 0) {
+                    channel.push({ content: part.text });
+                }
+            }
+        }
+        if (chunk.functionCalls) {
+            stepFunctionCalls.push(...chunk.functionCalls);
+        }
+        // Extract usage metadata from chunk
+        // promptTokenCount is typically in the final chunk, candidatesTokenCount accumulates
+        const usageMetadata = chunkRecord.usageMetadata;
+        if (usageMetadata) {
+            // Take the latest promptTokenCount (usually only in final chunk)
+            if (usageMetadata.promptTokenCount !== undefined &&
+                usageMetadata.promptTokenCount > 0) {
+                hooks.onUsageDelta?.("input", usageMetadata.promptTokenCount - stepInputTokens);
+                stepInputTokens = usageMetadata.promptTokenCount;
+                // Feed the context guard the REAL prompt size of this call.
+                hooks.onUsage?.(usageMetadata.promptTokenCount, usageMetadata.candidatesTokenCount ?? 0);
+                // cachedContentTokenCount is OVERLAPPING (a subset already inside
+                // 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);
+                hooks.onUsageDelta?.("cacheRead", chunkCacheReadTokens - stepCacheReadTokens);
+                stepCacheReadTokens = chunkCacheReadTokens;
+            }
+            // Take the latest candidatesTokenCount (accumulates through chunks)
+            if (usageMetadata.candidatesTokenCount !== undefined &&
+                usageMetadata.candidatesTokenCount > 0) {
+                hooks.onUsageDelta?.("output", 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) {
+                hooks.onUsageDelta?.("reasoning", usageMetadata.thoughtsTokenCount - stepReasoningTokens);
+                stepReasoningTokens = usageMetadata.thoughtsTokenCount;
+            }
+        }
+    }
+    return {
+        rawResponseParts,
+        stepFunctionCalls,
+        ...(lastFinishReason ? { finishReason: lastFinishReason } : {}),
+        inputTokens: stepInputTokens,
+        outputTokens: stepOutputTokens,
+        ...(stepCacheReadTokens ? { cacheReadTokens: stepCacheReadTokens } : {}),
+        ...(stepReasoningTokens ? { reasoningTokens: stepReasoningTokens } : {}),
+    };
+}
 export class GoogleVertexProvider extends BaseProvider {
     projectId;
     location;
@@ -1515,74 +1610,39 @@ export class GoogleVertexProvider extends BaseProvider {
                     // 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
-                        // This avoids using chunk.text which triggers SDK warning when
-                        // non-text parts (thoughtSignature, functionCall) are present
-                        const chunkRecord = chunk;
-                        const candidates = chunkRecord.candidates;
-                        const firstCandidate = candidates?.[0];
-                        // Capture the SDK finish reason (Bug 2: previously dropped). Last
-                        // non-empty value across chunks wins.
-                        const chunkFinishReason = firstCandidate?.finishReason;
-                        if (typeof chunkFinishReason === "string" && chunkFinishReason) {
-                            lastFinishReason = chunkFinishReason;
-                            stepFinishReason = chunkFinishReason;
-                        }
-                        const chunkContent = firstCandidate?.content;
-                        if (chunkContent && Array.isArray(chunkContent.parts)) {
-                            for (const part of chunkContent.parts) {
-                                rawResponseParts.push(part);
-                                if (typeof part.text === "string" && part.text.length > 0) {
-                                    incrementalTextChunks.push(part.text);
-                                }
+                    // The drain now lives in collectVertexStreamChunks. The hooks below
+                    // are the two couplings it had to this loop, plus the per-chunk
+                    // usage deltas — those must stay per-chunk so a step killed
+                    // mid-stream still bills what it reported.
+                    const collected = await collectVertexStreamChunks(stream, {
+                        push: (chunk) => {
+                            if (chunk.content) {
+                                incrementalTextChunks.push(chunk.content);
                             }
-                        }
-                        if (chunk.functionCalls) {
-                            stepFunctionCalls.push(...chunk.functionCalls);
-                        }
-                        // Extract usage metadata from chunk
-                        // promptTokenCount is typically in the final chunk, candidatesTokenCount accumulates
-                        const usageMetadata = chunkRecord.usageMetadata;
-                        if (usageMetadata) {
-                            // Take the latest promptTokenCount (usually only in final chunk)
-                            if (usageMetadata.promptTokenCount !== undefined &&
-                                usageMetadata.promptTokenCount > 0) {
-                                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 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;
+                        },
+                    }, {
+                        onProgress: () => turnClock.noteProgress(),
+                        onUsage: (input, output) => contextGuard.noteUsage(input, output),
+                        onUsageDelta: (counter, delta) => {
+                            if (counter === "input") {
+                                totalInputTokens += delta;
                             }
-                            // Take the latest candidatesTokenCount (accumulates through chunks)
-                            if (usageMetadata.candidatesTokenCount !== undefined &&
-                                usageMetadata.candidatesTokenCount > 0) {
-                                totalOutputTokens +=
-                                    usageMetadata.candidatesTokenCount - stepOutputTokens;
-                                stepOutputTokens = usageMetadata.candidatesTokenCount;
+                            else if (counter === "output") {
+                                totalOutputTokens += delta;
                             }
-                            // 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;
+                            else if (counter === "cacheRead") {
+                                totalCacheReadTokens += delta;
                             }
-                        }
+                            else {
+                                totalReasoningTokens += delta;
+                            }
+                        },
+                    });
+                    rawResponseParts.push(...collected.rawResponseParts);
+                    stepFunctionCalls.push(...collected.stepFunctionCalls);
+                    if (collected.finishReason) {
+                        stepFinishReason = collected.finishReason;
+                        lastFinishReason = collected.finishReason;
                     }
                     // Extract text from raw parts after stream completes
                     // This avoids SDK warning about non-text parts (thoughtSignature, functionCall)
@@ -2508,70 +2568,35 @@ export class GoogleVertexProvider extends BaseProvider {
                     // 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();
-                        // Extract raw parts from candidates FIRST
-                        // This avoids using chunk.text which triggers SDK warning when
-                        // non-text parts (thoughtSignature, functionCall) are present
-                        const chunkRecord = chunk;
-                        const candidates = chunkRecord.candidates;
-                        const firstCandidate = candidates?.[0];
-                        // Capture the SDK finish reason (Bug 2: previously dropped). Last
-                        // non-empty value across chunks wins.
-                        const chunkFinishReason = firstCandidate?.finishReason;
-                        if (typeof chunkFinishReason === "string" && chunkFinishReason) {
-                            lastFinishReason = chunkFinishReason;
-                            stepFinishReason = chunkFinishReason;
-                        }
-                        const chunkContent = firstCandidate?.content;
-                        if (chunkContent && Array.isArray(chunkContent.parts)) {
-                            rawResponseParts.push(...chunkContent.parts);
-                        }
-                        if (chunk.functionCalls) {
-                            stepFunctionCalls.push(...chunk.functionCalls);
-                        }
-                        // Extract usage metadata from chunk
-                        // promptTokenCount is typically in the final chunk, candidatesTokenCount accumulates
-                        const usageMetadata = chunkRecord.usageMetadata;
-                        if (usageMetadata) {
-                            // Take the latest promptTokenCount (usually only in final chunk)
-                            if (usageMetadata.promptTokenCount !== undefined &&
-                                usageMetadata.promptTokenCount > 0) {
-                                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 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;
+                    // Same collector as the streaming twin. generate() returns one
+                    // result rather than streaming, so the channel is a no-op — the
+                    // stream loop uses it to fill incrementalTextChunks for replay, and
+                    // there is nothing to replay here. Everything else is identical,
+                    // including the per-chunk usage deltas that keep the turn totals
+                    // correct for a step killed mid-stream.
+                    const collected = await collectVertexStreamChunks(stream, { push: () => { } }, {
+                        onProgress: () => turnClock.noteProgress(),
+                        onUsage: (input, output) => contextGuard.noteUsage(input, output),
+                        onUsageDelta: (counter, delta) => {
+                            if (counter === "input") {
+                                totalInputTokens += delta;
                             }
-                            // Take the latest candidatesTokenCount (accumulates through chunks)
-                            if (usageMetadata.candidatesTokenCount !== undefined &&
-                                usageMetadata.candidatesTokenCount > 0) {
-                                totalOutputTokens +=
-                                    usageMetadata.candidatesTokenCount - stepOutputTokens;
-                                stepOutputTokens = usageMetadata.candidatesTokenCount;
+                            else if (counter === "output") {
+                                totalOutputTokens += delta;
                             }
-                            // 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;
+                            else if (counter === "cacheRead") {
+                                totalCacheReadTokens += delta;
                             }
-                        }
+                            else {
+                                totalReasoningTokens += delta;
+                            }
+                        },
+                    });
+                    rawResponseParts.push(...collected.rawResponseParts);
+                    stepFunctionCalls.push(...collected.stepFunctionCalls);
+                    if (collected.finishReason) {
+                        stepFinishReason = collected.finishReason;
+                        lastFinishReason = collected.finishReason;
                     }
                     // Extract text from raw parts after stream completes
                     // This avoids SDK warning about non-text parts (thoughtSignature, functionCall)

package/dist/types/loopEngine.d.ts

@@ -1,6 +1,6 @@
 import type Anthropic from "@anthropic-ai/sdk";
 import type { Tool } from "./tools.js";
-import type { NativeFunctionCall, NativeToolDeclarationsResult } from "./providers.js";
+import type { CollectedChunkResult, NativeFunctionCall, NativeToolDeclarationsResult } from "./providers.js";
 /**
  * One chunk on the engine's stream.
  *
@@ -240,6 +240,22 @@ export type GeminiLoopAdapterCoreConfig = {
      * step with that step's real token counts.
      */
     noteUsage?: (inputTokens: number, outputTokens: number) => void;
+    /**
+     * Fold one step's raw stream into the shape the adapter reports.
+     *
+     * Defaults to `collectStreamChunksIncremental`, which is what AI Studio and
+     * any provider sharing the googleNativeGemini3 helpers want. Vertex does
+     * NOT share them: its loop drains the stream itself, folding cumulative
+     * usage counts as deltas and capturing thought signatures in its own way,
+     * and that behaviour is characterized rather than incidental.
+     *
+     * So the collector is a hook rather than a hard-coded call. A provider
+     * whose drain differs supplies its own and keeps its measured behaviour;
+     * one that matches the shared helper passes nothing.
+     */
+    collectStep?: (stream: unknown, channel: {
+        push(chunk: AgenticLoopChunk): void;
+    }) => Promise<CollectedChunkResult>;
 };
 /**
  * Opt in to the single MALFORMED_FUNCTION_CALL retry.

package/dist/types/providers.d.ts

@@ -1735,6 +1735,14 @@ export type NativeFunctionResponse = {
     };
 };
 /** Result from collectStreamChunks. */
+/**
+ * Which turn-level counter a per-chunk Vertex usage delta belongs to.
+ *
+ * Vertex updates its turn totals incrementally so they stay correct mid-stream
+ * — a step killed by an abort, the turn deadline or the stall watchdog still
+ * bills the tokens it already reported.
+ */
+export type VertexUsageCounter = "input" | "output" | "cacheRead" | "reasoning";
 export type CollectedChunkResult = {
     rawResponseParts: unknown[];
     stepFunctionCalls: NativeFunctionCall[];