@juspay/neurolink 10.2.2 → 10.3.0

package/dist/context/stepBudgetGuard.d.ts

@@ -0,0 +1,47 @@
+/**
+ * Per-step context budget guard for the AI-SDK agent loop.
+ *
+ * Pre-call budgeting (`checkContextBudget` + compaction) runs ONCE before
+ * dispatch and only sees the input/session conversation. The AI-SDK tool loop
+ * then appends assistant turns and tool results on every step — growth the
+ * pre-call pipeline never sees, which is how long agentic runs overflow the
+ * model's real context window mid-loop (provider 400s after dozens of tool
+ * calls). The googleVertex native loops already guard this via
+ * `createContextGuard`; this module brings the AI-SDK path (every provider
+ * that delegates to `generateText`) to parity — and goes one step further:
+ * instead of stopping the loop, it deterministically reclaims budget so the
+ * loop can CONTINUE.
+ *
+ * Wired in `GenerationHandler.callGenerateText` through
+ * `experimental_prepareStep`, whose result may replace the step's `messages`.
+ * The guard operates on `ModelMessage[]` natively (no lossy ChatMessage
+ * round-trip) and never makes LLM calls:
+ *
+ *   Stage 1 — truncate OLD tool outputs to head/tail previews
+ *             (`generateToolOutputPreview`), oldest first, outside the
+ *             protected recent tail.
+ *   Stage 2 — drop the oldest complete tool exchanges (assistant tool-call
+ *             message + its following tool-result messages, as a unit, so
+ *             call/result pairing stays intact), replacing them with a single
+ *             elision note.
+ *
+ * The system prompt and tool definitions ride OUTSIDE the step messages (the
+ * handler hoists system into generateText's `system` option), so their cost is
+ * passed in as `fixedOverheadTokens`. The first user message (the task) and
+ * the most recent messages are never touched.
+ */
+import type { ModelMessage, StepBudgetGuardConfig } from "../types/index.js";
+/** Estimate the token cost of a step's message array. */
+export declare function estimateStepMessagesTokens(messages: readonly ModelMessage[], provider?: string): number;
+/**
+ * Estimate the fixed per-request overhead: hoisted system prompt + tool
+ * definitions. Mirrors `checkContextBudget`'s categories for the pieces that
+ * do not live in the step messages.
+ */
+export declare function estimateFixedOverheadTokens(system: unknown, tools: Record<string, unknown> | undefined, provider?: string): number;
+/**
+ * Create a per-step budget guard. Returns a function that, given the step's
+ * messages, returns a compacted replacement array when the projected request
+ * exceeds the threshold — or `undefined` when no change is needed.
+ */
+export declare function createStepBudgetGuard(config: StepBudgetGuardConfig): (messages: readonly ModelMessage[]) => ModelMessage[] | undefined;

package/dist/context/stepBudgetGuard.js

@@ -0,0 +1,289 @@
+/**
+ * Per-step context budget guard for the AI-SDK agent loop.
+ *
+ * Pre-call budgeting (`checkContextBudget` + compaction) runs ONCE before
+ * dispatch and only sees the input/session conversation. The AI-SDK tool loop
+ * then appends assistant turns and tool results on every step — growth the
+ * pre-call pipeline never sees, which is how long agentic runs overflow the
+ * model's real context window mid-loop (provider 400s after dozens of tool
+ * calls). The googleVertex native loops already guard this via
+ * `createContextGuard`; this module brings the AI-SDK path (every provider
+ * that delegates to `generateText`) to parity — and goes one step further:
+ * instead of stopping the loop, it deterministically reclaims budget so the
+ * loop can CONTINUE.
+ *
+ * Wired in `GenerationHandler.callGenerateText` through
+ * `experimental_prepareStep`, whose result may replace the step's `messages`.
+ * The guard operates on `ModelMessage[]` natively (no lossy ChatMessage
+ * round-trip) and never makes LLM calls:
+ *
+ *   Stage 1 — truncate OLD tool outputs to head/tail previews
+ *             (`generateToolOutputPreview`), oldest first, outside the
+ *             protected recent tail.
+ *   Stage 2 — drop the oldest complete tool exchanges (assistant tool-call
+ *             message + its following tool-result messages, as a unit, so
+ *             call/result pairing stays intact), replacing them with a single
+ *             elision note.
+ *
+ * The system prompt and tool definitions ride OUTSIDE the step messages (the
+ * handler hoists system into generateText's `system` option), so their cost is
+ * passed in as `fixedOverheadTokens`. The first user message (the task) and
+ * the most recent messages are never touched.
+ */
+import { DEFAULT_CONTEXT_GUARD_RATIO } from "../core/constants.js";
+import { getAvailableInputTokens } from "../constants/contextWindows.js";
+import { estimateTokens, TOKENS_PER_MESSAGE, } from "../utils/tokenEstimation.js";
+import { generateToolOutputPreview } from "./toolOutputLimits.js";
+import { logger } from "../utils/logger.js";
+/** Estimated tokens for a tool definition that fails to serialize. */
+const TOKENS_PER_TOOL_DEFINITION = 200;
+/** Messages at the end of the conversation the guard never modifies. */
+const PROTECTED_TAIL_MESSAGES = 4;
+/** Stage-1 preview budget for an old tool output (bytes). */
+const OLD_TOOL_OUTPUT_PREVIEW_BYTES = 2_048;
+/** Stage-1 preview budget for an old tool output (lines). */
+const OLD_TOOL_OUTPUT_PREVIEW_LINES = 60;
+/**
+ * Serialize any ModelMessage content to text for estimation. Tool-call args
+ * and tool-result outputs are JSON-stringified; unserializable values fall
+ * back to a fixed-size placeholder so estimation never throws.
+ */
+function contentToText(content) {
+    if (typeof content === "string") {
+        return content;
+    }
+    try {
+        return JSON.stringify(content) ?? "";
+    }
+    catch {
+        return "x".repeat(TOKENS_PER_TOOL_DEFINITION * 4);
+    }
+}
+/** Estimate the token cost of a step's message array. */
+export function estimateStepMessagesTokens(messages, provider) {
+    let total = 0;
+    for (const message of messages) {
+        total +=
+            estimateTokens(contentToText(message.content), provider) +
+                TOKENS_PER_MESSAGE;
+    }
+    return total;
+}
+/**
+ * Estimate the fixed per-request overhead: hoisted system prompt + tool
+ * definitions. Mirrors `checkContextBudget`'s categories for the pieces that
+ * do not live in the step messages.
+ */
+export function estimateFixedOverheadTokens(system, tools, provider) {
+    let total = system
+        ? estimateTokens(contentToText(system), provider) + TOKENS_PER_MESSAGE
+        : 0;
+    for (const tool of Object.values(tools ?? {})) {
+        try {
+            total += estimateTokens(JSON.stringify(tool) ?? "", provider);
+        }
+        catch {
+            total += TOKENS_PER_TOOL_DEFINITION;
+        }
+    }
+    return total;
+}
+/**
+ * Serialize a ToolResultOutput to the text the MODEL should see in a preview.
+ * Variant-aware: `text`/`error-text` carry their payload in `.value` directly —
+ * stringifying the wrapper would put escaped `{"type":"text","value":…}` JSON
+ * in front of the model instead of the actual output. `json`/`error-json`/
+ * `content` serialize their value; unknown shapes fall back to the wrapper.
+ */
+function toolResultOutputToText(output) {
+    const variant = output;
+    if (variant && typeof variant === "object" && "type" in variant) {
+        if ((variant.type === "text" || variant.type === "error-text") &&
+            typeof variant.value === "string") {
+            return variant.value;
+        }
+        if (variant.type === "json" ||
+            variant.type === "error-json" ||
+            variant.type === "content") {
+            return contentToText(variant.value);
+        }
+    }
+    return contentToText(output);
+}
+/** True when the message is an assistant message that issues tool calls. */
+function isToolCallAssistantMessage(message) {
+    return (message.role === "assistant" &&
+        Array.isArray(message.content) &&
+        message.content.some((part) => part?.type === "tool-call"));
+}
+/**
+ * Stage 1: replace large tool-result outputs outside the protected tail with
+ * head/tail previews. Returns the new array plus how many outputs shrank.
+ */
+function truncateOldToolOutputs(messages) {
+    const cutoff = Math.max(0, messages.length - PROTECTED_TAIL_MESSAGES);
+    let truncated = 0;
+    const next = messages.map((message, index) => {
+        if (index >= cutoff || message.role !== "tool") {
+            return message;
+        }
+        if (!Array.isArray(message.content)) {
+            return message;
+        }
+        let changed = false;
+        const content = message.content.map((part) => {
+            const resultPart = part;
+            if (resultPart?.type !== "tool-result") {
+                return part;
+            }
+            const serialized = toolResultOutputToText(resultPart.output);
+            if (serialized.length <= OLD_TOOL_OUTPUT_PREVIEW_BYTES) {
+                return part;
+            }
+            const { preview } = generateToolOutputPreview(serialized, {
+                maxBytes: OLD_TOOL_OUTPUT_PREVIEW_BYTES,
+                maxLines: OLD_TOOL_OUTPUT_PREVIEW_LINES,
+            });
+            changed = true;
+            truncated += 1;
+            return {
+                ...resultPart,
+                output: { type: "text", value: preview },
+            };
+        });
+        return changed ? { ...message, content } : message;
+    });
+    return { messages: next, truncated };
+}
+/**
+ * Stage 2: drop the oldest complete tool exchanges — an assistant tool-call
+ * message together with ALL directly-following `tool` messages — until the
+ * estimate fits or only the protected head/tail remains. The first
+ * non-assistant message run (the task) is never dropped. A single elision
+ * note replaces everything removed so the model knows history was elided.
+ */
+function dropOldestToolExchanges(messages, budgetTokens, fixedOverheadTokens, provider) {
+    const result = [...messages];
+    let droppedExchanges = 0;
+    // Running-total accounting: estimate each message ONCE, keep the estimates
+    // array in lockstep with `result`, and subtract dropped blocks — instead of
+    // re-estimating the whole array on every iteration (O(n²) with many drops).
+    const estimates = result.map((message) => estimateTokens(contentToText(message.content), provider) +
+        TOKENS_PER_MESSAGE);
+    let currentTokens = fixedOverheadTokens + estimates.reduce((sum, tokens) => sum + tokens, 0);
+    while (currentTokens > budgetTokens) {
+        // Find the FIRST (oldest) droppable exchange outside the protected tail.
+        const tailStart = Math.max(0, result.length - PROTECTED_TAIL_MESSAGES);
+        let exchangeStart = -1;
+        for (let i = 0; i < tailStart; i++) {
+            if (isToolCallAssistantMessage(result[i])) {
+                exchangeStart = i;
+                break;
+            }
+        }
+        if (exchangeStart === -1) {
+            break; // nothing left the guard is allowed to drop
+        }
+        let exchangeEnd = exchangeStart + 1;
+        while (exchangeEnd < result.length && result[exchangeEnd].role === "tool") {
+            exchangeEnd++;
+        }
+        if (exchangeEnd > tailStart) {
+            // The oldest remaining exchange bleeds into the protected tail. Because
+            // the scan is oldest-first, every exchange after this one STARTS inside
+            // the tail (this one's result chain reaches it), and every exchange
+            // before it was already dropped by earlier iterations — so there is
+            // nothing else the guard may remove. Stop.
+            break;
+        }
+        const dropped = estimates
+            .slice(exchangeStart, exchangeEnd)
+            .reduce((sum, tokens) => sum + tokens, 0);
+        result.splice(exchangeStart, exchangeEnd - exchangeStart);
+        estimates.splice(exchangeStart, exchangeEnd - exchangeStart);
+        currentTokens -= dropped;
+        droppedExchanges++;
+    }
+    if (droppedExchanges > 0) {
+        // Insert one elision note where history was removed: after the leading
+        // non-exchange messages (typically the first user/task message), but
+        // never after the protected tail — when every droppable exchange was
+        // removed, an uncapped scan would append the note at the END, where the
+        // "history was removed" cue lands after the content it refers to.
+        let noteIndex = 0;
+        while (noteIndex < result.length &&
+            !isToolCallAssistantMessage(result[noteIndex])) {
+            noteIndex++;
+        }
+        const tailBoundary = Math.max(0, result.length - PROTECTED_TAIL_MESSAGES);
+        result.splice(Math.min(noteIndex, tailBoundary), 0, {
+            role: "user",
+            content: [
+                {
+                    type: "text",
+                    text: `[context truncated: ${droppedExchanges} earlier tool exchange(s) were removed to fit the model's context window. Continue from the remaining context.]`,
+                },
+            ],
+        });
+    }
+    return { messages: result, droppedExchanges };
+}
+/**
+ * Create a per-step budget guard. Returns a function that, given the step's
+ * messages, returns a compacted replacement array when the projected request
+ * exceeds the threshold — or `undefined` when no change is needed.
+ */
+export function createStepBudgetGuard(config) {
+    const { provider, model, maxTokens, fixedOverheadTokens = 0, getFixedOverheadTokens, thresholdRatio = DEFAULT_CONTEXT_GUARD_RATIO, } = config;
+    const availableInput = getAvailableInputTokens(provider, model, maxTokens);
+    const thresholdTokens = Math.floor(availableInput * thresholdRatio);
+    return function guardStepMessages(messages) {
+        // Resolve overhead per invocation: the tool set can GROW mid-loop
+        // (search_tools hydration adds discovered tools between steps), so a
+        // once-captured value would undercount later steps.
+        const overheadTokens = getFixedOverheadTokens?.() ?? fixedOverheadTokens;
+        const estimate = overheadTokens + estimateStepMessagesTokens(messages, provider);
+        // Logger Guard: per-step diagnostics for debugging why a long run does
+        // (or does not) trigger compaction — gated so nothing is serialized when
+        // debug logging is off.
+        if (logger.shouldLog("debug")) {
+            logger.debug("[StepBudgetGuard] step estimate", {
+                provider,
+                model,
+                messageCount: messages.length,
+                estimatedTokens: estimate,
+                thresholdTokens,
+                willCompact: estimate > thresholdTokens,
+            });
+        }
+        if (estimate <= thresholdTokens) {
+            return undefined;
+        }
+        // Stage 1: shrink old tool outputs to previews.
+        const stage1 = truncateOldToolOutputs([...messages]);
+        let compacted = stage1.messages;
+        let newEstimate = overheadTokens + estimateStepMessagesTokens(compacted, provider);
+        // Stage 2: drop oldest complete tool exchanges if still over.
+        let droppedExchanges = 0;
+        if (newEstimate > thresholdTokens) {
+            const stage2 = dropOldestToolExchanges(compacted, thresholdTokens, overheadTokens, provider);
+            compacted = stage2.messages;
+            droppedExchanges = stage2.droppedExchanges;
+            newEstimate =
+                overheadTokens + estimateStepMessagesTokens(compacted, provider);
+        }
+        if (stage1.truncated === 0 && droppedExchanges === 0) {
+            return undefined; // nothing actionable (already all-protected)
+        }
+        logger.info("[StepBudgetGuard] Compacted agent-loop step messages", {
+            provider,
+            model,
+            estimatedTokens: estimate,
+            thresholdTokens,
+            afterTokens: newEstimate,
+            toolOutputsTruncated: stage1.truncated,
+            exchangesDropped: droppedExchanges,
+        });
+        return compacted;
+    };
+}

package/dist/core/modules/GenerationHandler.js

@@ -21,6 +21,7 @@ import { calculateCost } from "../../utils/pricing.js";
 import { withProviderRetry } from "../../utils/providerRetry.js";
 import { calculateCacheSavingsPercent, extractCacheCreationTokens, extractCacheReadTokens, extractTokenUsage, } from "../../utils/tokenUtils.js";
 import { DEFAULT_MAX_STEPS } from "../constants.js";
+import { createStepBudgetGuard, estimateFixedOverheadTokens, } from "../../context/stepBudgetGuard.js";
 import { isTemperatureDeprecatedError, isSchemaComplexityError, isToolsSchemaConflictError, isToolsSchemaExclusionInForce, } from "./structuredOutputPolicy.js";
 import { coerceJsonToSchema } from "../../utils/json/coerce.js";
 import { NoObjectGeneratedError } from "../../utils/generationErrors.js";
@@ -116,6 +117,37 @@ export class GenerationHandler {
         // rather than passing them inside `messages` (deprecated by the AI SDK,
         // rejected in v7). See extractSystemMessages for the rationale. (#1024)
         const { system, messages: nonSystemMessages } = extractSystemMessages(messages);
+        // Per-step context budget guard: the tool loop appends assistant turns and
+        // tool results on every step — growth the pre-call budget check never
+        // sees. Estimate each step's projected request and deterministically
+        // reclaim budget (truncate old tool outputs, then drop oldest exchanges)
+        // so long agentic runs cannot overflow the model's window mid-loop.
+        // Parity with the googleVertex native loops' createContextGuard, upgraded
+        // from stop-only to compact-and-continue. The caller's prepareStep result
+        // wins on conflicts; the guard only contributes `messages`.
+        //
+        // Overhead is resolved PER STEP because `toolsWithCache` is deliberately
+        // mutable (search_tools hydration adds discovered tools mid-loop) — a
+        // once-captured estimate would undercount later steps. Tools are only
+        // ever added, so memoizing on tool count keeps the common step O(1).
+        let cachedOverhead = { toolCount: -1, tokens: 0 };
+        const stepBudgetGuard = createStepBudgetGuard({
+            provider: this.providerName ?? "unknown",
+            model: this.modelName,
+            maxTokens: options.maxTokens,
+            getFixedOverheadTokens: () => {
+                const toolCount = shouldUseTools
+                    ? Object.keys(toolsWithCache).length
+                    : 0;
+                if (toolCount !== cachedOverhead.toolCount) {
+                    cachedOverhead = {
+                        toolCount,
+                        tokens: estimateFixedOverheadTokens(system, shouldUseTools ? toolsWithCache : undefined, this.providerName),
+                    };
+                }
+                return cachedOverhead.tokens;
+            },
+        });
         return await generateText({
             model,
             ...(system && { system }),
@@ -125,11 +157,27 @@ export class GenerationHandler {
             stopWhen: stepCountIs(options.maxSteps ?? DEFAULT_MAX_STEPS),
             ...(shouldUseTools &&
                 options.toolChoice && { toolChoice: options.toolChoice }),
-            ...(prepareStep && {
-                experimental_prepareStep: ((stepOptions) => prepareStep({
-                    ...stepOptions,
-                    maxSteps: options.maxSteps ?? DEFAULT_MAX_STEPS,
-                })),
+            experimental_prepareStep: (async (stepOptions) => {
+                // Public contract preserved: a caller-supplied prepareStep receives
+                // the ORIGINAL AI-SDK step options, exactly as before the guard
+                // existed — callers that inspect message history see the real thing.
+                const callerResult = prepareStep
+                    ? await prepareStep({
+                        ...stepOptions,
+                        maxSteps: options.maxSteps ?? DEFAULT_MAX_STEPS,
+                    })
+                    : undefined;
+                // The guard runs LAST, on the messages that will actually be sent:
+                // the caller's override when one was returned (out-of-contract for
+                // NeuroLink's public prepareStep type, but possible at runtime), else
+                // the step's own messages. It never replaces a caller's content
+                // choices — it only reclaims budget from whatever was chosen.
+                const callerMessages = callerResult?.messages;
+                const compacted = stepBudgetGuard(callerMessages ?? stepOptions.messages);
+                if (!compacted) {
+                    return callerResult;
+                }
+                return { ...(callerResult ?? {}), messages: compacted };
             }),
             temperature: options.temperature,
             maxOutputTokens: options.maxTokens,

package/dist/lib/context/stepBudgetGuard.d.ts

@@ -0,0 +1,47 @@
+/**
+ * Per-step context budget guard for the AI-SDK agent loop.
+ *
+ * Pre-call budgeting (`checkContextBudget` + compaction) runs ONCE before
+ * dispatch and only sees the input/session conversation. The AI-SDK tool loop
+ * then appends assistant turns and tool results on every step — growth the
+ * pre-call pipeline never sees, which is how long agentic runs overflow the
+ * model's real context window mid-loop (provider 400s after dozens of tool
+ * calls). The googleVertex native loops already guard this via
+ * `createContextGuard`; this module brings the AI-SDK path (every provider
+ * that delegates to `generateText`) to parity — and goes one step further:
+ * instead of stopping the loop, it deterministically reclaims budget so the
+ * loop can CONTINUE.
+ *
+ * Wired in `GenerationHandler.callGenerateText` through
+ * `experimental_prepareStep`, whose result may replace the step's `messages`.
+ * The guard operates on `ModelMessage[]` natively (no lossy ChatMessage
+ * round-trip) and never makes LLM calls:
+ *
+ *   Stage 1 — truncate OLD tool outputs to head/tail previews
+ *             (`generateToolOutputPreview`), oldest first, outside the
+ *             protected recent tail.
+ *   Stage 2 — drop the oldest complete tool exchanges (assistant tool-call
+ *             message + its following tool-result messages, as a unit, so
+ *             call/result pairing stays intact), replacing them with a single
+ *             elision note.
+ *
+ * The system prompt and tool definitions ride OUTSIDE the step messages (the
+ * handler hoists system into generateText's `system` option), so their cost is
+ * passed in as `fixedOverheadTokens`. The first user message (the task) and
+ * the most recent messages are never touched.
+ */
+import type { ModelMessage, StepBudgetGuardConfig } from "../types/index.js";
+/** Estimate the token cost of a step's message array. */
+export declare function estimateStepMessagesTokens(messages: readonly ModelMessage[], provider?: string): number;
+/**
+ * Estimate the fixed per-request overhead: hoisted system prompt + tool
+ * definitions. Mirrors `checkContextBudget`'s categories for the pieces that
+ * do not live in the step messages.
+ */
+export declare function estimateFixedOverheadTokens(system: unknown, tools: Record<string, unknown> | undefined, provider?: string): number;
+/**
+ * Create a per-step budget guard. Returns a function that, given the step's
+ * messages, returns a compacted replacement array when the projected request
+ * exceeds the threshold — or `undefined` when no change is needed.
+ */
+export declare function createStepBudgetGuard(config: StepBudgetGuardConfig): (messages: readonly ModelMessage[]) => ModelMessage[] | undefined;