@juspay/neurolink 10.2.1 → 10.3.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -21,12 +21,24 @@ export declare const DEFAULT_OUTPUT_RESERVE_RATIO = 0.35;
21
21
  * The "_default" key is the fallback for unknown models within a provider.
22
22
  */
23
23
  export declare const MODEL_CONTEXT_WINDOWS: Record<string, Record<string, number>>;
24
+ /**
25
+ * Register a runtime-discovered context window for a provider/model pair.
26
+ * Later registrations overwrite earlier ones (rediscovery refreshes values).
27
+ * Non-positive/non-finite windows are ignored so a malformed discovery source
28
+ * can never shrink a budget to zero.
29
+ */
30
+ export declare function registerRuntimeContextWindow(provider: string, model: string, contextWindow: number): void;
31
+ /** Test hook: clear runtime-discovered windows (state is module-global). */
32
+ export declare function clearRuntimeContextWindows(): void;
24
33
  /**
25
34
  * Resolve context window size for a provider/model combination.
26
35
  *
27
36
  * Priority:
28
37
  * 0. Dynamic model registry (DynamicModelProvider) — resolves cross-provider
29
38
  * models (e.g. Claude on Vertex) that the static table cannot handle
39
+ * 0.5 Runtime-discovered windows (registerRuntimeContextWindow) — real
40
+ * per-model limits fetched from the serving infrastructure (LiteLLM
41
+ * `/model/info`)
30
42
  * 1. Exact model match under provider in static registry
31
43
  * 2. Prefix match under provider in static registry
32
44
  * 3. Provider's _default in static registry
@@ -392,12 +392,43 @@ function normalizeProviderForLookup(provider) {
392
392
  // to DEFAULT_CONTEXT_WINDOW.
393
393
  return PROVIDER_ALIAS_MAP[stripped] ?? stripped;
394
394
  }
395
+ /**
396
+ * Runtime-discovered context windows, keyed `${provider}:${model}`.
397
+ *
398
+ * Populated asynchronously by providers that can discover real per-model
399
+ * limits at runtime (e.g. the LiteLLM provider reads `max_input_tokens` from
400
+ * the proxy's `/model/info`), and read synchronously by
401
+ * {@link getContextWindowSize} — the same async-populate/sync-read contract as
402
+ * the DynamicModelProvider registry. Keys use the RAW provider string callers
403
+ * pass into budget calculations (see the alias-map comment above: normalized
404
+ * provider names never reach these lookups).
405
+ */
406
+ const RUNTIME_CONTEXT_WINDOWS = new Map();
407
+ /**
408
+ * Register a runtime-discovered context window for a provider/model pair.
409
+ * Later registrations overwrite earlier ones (rediscovery refreshes values).
410
+ * Non-positive/non-finite windows are ignored so a malformed discovery source
411
+ * can never shrink a budget to zero.
412
+ */
413
+ export function registerRuntimeContextWindow(provider, model, contextWindow) {
414
+ if (!Number.isFinite(contextWindow) || contextWindow <= 0) {
415
+ return;
416
+ }
417
+ RUNTIME_CONTEXT_WINDOWS.set(`${provider}:${model}`, contextWindow);
418
+ }
419
+ /** Test hook: clear runtime-discovered windows (state is module-global). */
420
+ export function clearRuntimeContextWindows() {
421
+ RUNTIME_CONTEXT_WINDOWS.clear();
422
+ }
395
423
  /**
396
424
  * Resolve context window size for a provider/model combination.
397
425
  *
398
426
  * Priority:
399
427
  * 0. Dynamic model registry (DynamicModelProvider) — resolves cross-provider
400
428
  * models (e.g. Claude on Vertex) that the static table cannot handle
429
+ * 0.5 Runtime-discovered windows (registerRuntimeContextWindow) — real
430
+ * per-model limits fetched from the serving infrastructure (LiteLLM
431
+ * `/model/info`)
401
432
  * 1. Exact model match under provider in static registry
402
433
  * 2. Prefix match under provider in static registry
403
434
  * 3. Provider's _default in static registry
@@ -421,6 +452,13 @@ export function getContextWindowSize(provider, model) {
421
452
  // Dynamic registry not initialized yet — fall through to static lookup
422
453
  }
423
454
  }
455
+ // Step 0.5: Runtime-discovered window for this exact provider/model.
456
+ if (model) {
457
+ const discovered = RUNTIME_CONTEXT_WINDOWS.get(`${provider}:${model}`);
458
+ if (discovered !== undefined) {
459
+ return discovered;
460
+ }
461
+ }
424
462
  // Static fallback chain — normalize aliases first so "lmstudio" / "llama.cpp" /
425
463
  // "nvidianim" find their canonical entries instead of falling back to default.
426
464
  const canonical = normalizeProviderForLookup(provider);
@@ -0,0 +1,47 @@
1
+ /**
2
+ * Per-step context budget guard for the AI-SDK agent loop.
3
+ *
4
+ * Pre-call budgeting (`checkContextBudget` + compaction) runs ONCE before
5
+ * dispatch and only sees the input/session conversation. The AI-SDK tool loop
6
+ * then appends assistant turns and tool results on every step — growth the
7
+ * pre-call pipeline never sees, which is how long agentic runs overflow the
8
+ * model's real context window mid-loop (provider 400s after dozens of tool
9
+ * calls). The googleVertex native loops already guard this via
10
+ * `createContextGuard`; this module brings the AI-SDK path (every provider
11
+ * that delegates to `generateText`) to parity — and goes one step further:
12
+ * instead of stopping the loop, it deterministically reclaims budget so the
13
+ * loop can CONTINUE.
14
+ *
15
+ * Wired in `GenerationHandler.callGenerateText` through
16
+ * `experimental_prepareStep`, whose result may replace the step's `messages`.
17
+ * The guard operates on `ModelMessage[]` natively (no lossy ChatMessage
18
+ * round-trip) and never makes LLM calls:
19
+ *
20
+ * Stage 1 — truncate OLD tool outputs to head/tail previews
21
+ * (`generateToolOutputPreview`), oldest first, outside the
22
+ * protected recent tail.
23
+ * Stage 2 — drop the oldest complete tool exchanges (assistant tool-call
24
+ * message + its following tool-result messages, as a unit, so
25
+ * call/result pairing stays intact), replacing them with a single
26
+ * elision note.
27
+ *
28
+ * The system prompt and tool definitions ride OUTSIDE the step messages (the
29
+ * handler hoists system into generateText's `system` option), so their cost is
30
+ * passed in as `fixedOverheadTokens`. The first user message (the task) and
31
+ * the most recent messages are never touched.
32
+ */
33
+ import type { ModelMessage, StepBudgetGuardConfig } from "../types/index.js";
34
+ /** Estimate the token cost of a step's message array. */
35
+ export declare function estimateStepMessagesTokens(messages: readonly ModelMessage[], provider?: string): number;
36
+ /**
37
+ * Estimate the fixed per-request overhead: hoisted system prompt + tool
38
+ * definitions. Mirrors `checkContextBudget`'s categories for the pieces that
39
+ * do not live in the step messages.
40
+ */
41
+ export declare function estimateFixedOverheadTokens(system: unknown, tools: Record<string, unknown> | undefined, provider?: string): number;
42
+ /**
43
+ * Create a per-step budget guard. Returns a function that, given the step's
44
+ * messages, returns a compacted replacement array when the projected request
45
+ * exceeds the threshold — or `undefined` when no change is needed.
46
+ */
47
+ export declare function createStepBudgetGuard(config: StepBudgetGuardConfig): (messages: readonly ModelMessage[]) => ModelMessage[] | undefined;
@@ -0,0 +1,289 @@
1
+ /**
2
+ * Per-step context budget guard for the AI-SDK agent loop.
3
+ *
4
+ * Pre-call budgeting (`checkContextBudget` + compaction) runs ONCE before
5
+ * dispatch and only sees the input/session conversation. The AI-SDK tool loop
6
+ * then appends assistant turns and tool results on every step — growth the
7
+ * pre-call pipeline never sees, which is how long agentic runs overflow the
8
+ * model's real context window mid-loop (provider 400s after dozens of tool
9
+ * calls). The googleVertex native loops already guard this via
10
+ * `createContextGuard`; this module brings the AI-SDK path (every provider
11
+ * that delegates to `generateText`) to parity — and goes one step further:
12
+ * instead of stopping the loop, it deterministically reclaims budget so the
13
+ * loop can CONTINUE.
14
+ *
15
+ * Wired in `GenerationHandler.callGenerateText` through
16
+ * `experimental_prepareStep`, whose result may replace the step's `messages`.
17
+ * The guard operates on `ModelMessage[]` natively (no lossy ChatMessage
18
+ * round-trip) and never makes LLM calls:
19
+ *
20
+ * Stage 1 — truncate OLD tool outputs to head/tail previews
21
+ * (`generateToolOutputPreview`), oldest first, outside the
22
+ * protected recent tail.
23
+ * Stage 2 — drop the oldest complete tool exchanges (assistant tool-call
24
+ * message + its following tool-result messages, as a unit, so
25
+ * call/result pairing stays intact), replacing them with a single
26
+ * elision note.
27
+ *
28
+ * The system prompt and tool definitions ride OUTSIDE the step messages (the
29
+ * handler hoists system into generateText's `system` option), so their cost is
30
+ * passed in as `fixedOverheadTokens`. The first user message (the task) and
31
+ * the most recent messages are never touched.
32
+ */
33
+ import { DEFAULT_CONTEXT_GUARD_RATIO } from "../core/constants.js";
34
+ import { getAvailableInputTokens } from "../constants/contextWindows.js";
35
+ import { estimateTokens, TOKENS_PER_MESSAGE, } from "../utils/tokenEstimation.js";
36
+ import { generateToolOutputPreview } from "./toolOutputLimits.js";
37
+ import { logger } from "../utils/logger.js";
38
+ /** Estimated tokens for a tool definition that fails to serialize. */
39
+ const TOKENS_PER_TOOL_DEFINITION = 200;
40
+ /** Messages at the end of the conversation the guard never modifies. */
41
+ const PROTECTED_TAIL_MESSAGES = 4;
42
+ /** Stage-1 preview budget for an old tool output (bytes). */
43
+ const OLD_TOOL_OUTPUT_PREVIEW_BYTES = 2_048;
44
+ /** Stage-1 preview budget for an old tool output (lines). */
45
+ const OLD_TOOL_OUTPUT_PREVIEW_LINES = 60;
46
+ /**
47
+ * Serialize any ModelMessage content to text for estimation. Tool-call args
48
+ * and tool-result outputs are JSON-stringified; unserializable values fall
49
+ * back to a fixed-size placeholder so estimation never throws.
50
+ */
51
+ function contentToText(content) {
52
+ if (typeof content === "string") {
53
+ return content;
54
+ }
55
+ try {
56
+ return JSON.stringify(content) ?? "";
57
+ }
58
+ catch {
59
+ return "x".repeat(TOKENS_PER_TOOL_DEFINITION * 4);
60
+ }
61
+ }
62
+ /** Estimate the token cost of a step's message array. */
63
+ export function estimateStepMessagesTokens(messages, provider) {
64
+ let total = 0;
65
+ for (const message of messages) {
66
+ total +=
67
+ estimateTokens(contentToText(message.content), provider) +
68
+ TOKENS_PER_MESSAGE;
69
+ }
70
+ return total;
71
+ }
72
+ /**
73
+ * Estimate the fixed per-request overhead: hoisted system prompt + tool
74
+ * definitions. Mirrors `checkContextBudget`'s categories for the pieces that
75
+ * do not live in the step messages.
76
+ */
77
+ export function estimateFixedOverheadTokens(system, tools, provider) {
78
+ let total = system
79
+ ? estimateTokens(contentToText(system), provider) + TOKENS_PER_MESSAGE
80
+ : 0;
81
+ for (const tool of Object.values(tools ?? {})) {
82
+ try {
83
+ total += estimateTokens(JSON.stringify(tool) ?? "", provider);
84
+ }
85
+ catch {
86
+ total += TOKENS_PER_TOOL_DEFINITION;
87
+ }
88
+ }
89
+ return total;
90
+ }
91
+ /**
92
+ * Serialize a ToolResultOutput to the text the MODEL should see in a preview.
93
+ * Variant-aware: `text`/`error-text` carry their payload in `.value` directly —
94
+ * stringifying the wrapper would put escaped `{"type":"text","value":…}` JSON
95
+ * in front of the model instead of the actual output. `json`/`error-json`/
96
+ * `content` serialize their value; unknown shapes fall back to the wrapper.
97
+ */
98
+ function toolResultOutputToText(output) {
99
+ const variant = output;
100
+ if (variant && typeof variant === "object" && "type" in variant) {
101
+ if ((variant.type === "text" || variant.type === "error-text") &&
102
+ typeof variant.value === "string") {
103
+ return variant.value;
104
+ }
105
+ if (variant.type === "json" ||
106
+ variant.type === "error-json" ||
107
+ variant.type === "content") {
108
+ return contentToText(variant.value);
109
+ }
110
+ }
111
+ return contentToText(output);
112
+ }
113
+ /** True when the message is an assistant message that issues tool calls. */
114
+ function isToolCallAssistantMessage(message) {
115
+ return (message.role === "assistant" &&
116
+ Array.isArray(message.content) &&
117
+ message.content.some((part) => part?.type === "tool-call"));
118
+ }
119
+ /**
120
+ * Stage 1: replace large tool-result outputs outside the protected tail with
121
+ * head/tail previews. Returns the new array plus how many outputs shrank.
122
+ */
123
+ function truncateOldToolOutputs(messages) {
124
+ const cutoff = Math.max(0, messages.length - PROTECTED_TAIL_MESSAGES);
125
+ let truncated = 0;
126
+ const next = messages.map((message, index) => {
127
+ if (index >= cutoff || message.role !== "tool") {
128
+ return message;
129
+ }
130
+ if (!Array.isArray(message.content)) {
131
+ return message;
132
+ }
133
+ let changed = false;
134
+ const content = message.content.map((part) => {
135
+ const resultPart = part;
136
+ if (resultPart?.type !== "tool-result") {
137
+ return part;
138
+ }
139
+ const serialized = toolResultOutputToText(resultPart.output);
140
+ if (serialized.length <= OLD_TOOL_OUTPUT_PREVIEW_BYTES) {
141
+ return part;
142
+ }
143
+ const { preview } = generateToolOutputPreview(serialized, {
144
+ maxBytes: OLD_TOOL_OUTPUT_PREVIEW_BYTES,
145
+ maxLines: OLD_TOOL_OUTPUT_PREVIEW_LINES,
146
+ });
147
+ changed = true;
148
+ truncated += 1;
149
+ return {
150
+ ...resultPart,
151
+ output: { type: "text", value: preview },
152
+ };
153
+ });
154
+ return changed ? { ...message, content } : message;
155
+ });
156
+ return { messages: next, truncated };
157
+ }
158
+ /**
159
+ * Stage 2: drop the oldest complete tool exchanges — an assistant tool-call
160
+ * message together with ALL directly-following `tool` messages — until the
161
+ * estimate fits or only the protected head/tail remains. The first
162
+ * non-assistant message run (the task) is never dropped. A single elision
163
+ * note replaces everything removed so the model knows history was elided.
164
+ */
165
+ function dropOldestToolExchanges(messages, budgetTokens, fixedOverheadTokens, provider) {
166
+ const result = [...messages];
167
+ let droppedExchanges = 0;
168
+ // Running-total accounting: estimate each message ONCE, keep the estimates
169
+ // array in lockstep with `result`, and subtract dropped blocks — instead of
170
+ // re-estimating the whole array on every iteration (O(n²) with many drops).
171
+ const estimates = result.map((message) => estimateTokens(contentToText(message.content), provider) +
172
+ TOKENS_PER_MESSAGE);
173
+ let currentTokens = fixedOverheadTokens + estimates.reduce((sum, tokens) => sum + tokens, 0);
174
+ while (currentTokens > budgetTokens) {
175
+ // Find the FIRST (oldest) droppable exchange outside the protected tail.
176
+ const tailStart = Math.max(0, result.length - PROTECTED_TAIL_MESSAGES);
177
+ let exchangeStart = -1;
178
+ for (let i = 0; i < tailStart; i++) {
179
+ if (isToolCallAssistantMessage(result[i])) {
180
+ exchangeStart = i;
181
+ break;
182
+ }
183
+ }
184
+ if (exchangeStart === -1) {
185
+ break; // nothing left the guard is allowed to drop
186
+ }
187
+ let exchangeEnd = exchangeStart + 1;
188
+ while (exchangeEnd < result.length && result[exchangeEnd].role === "tool") {
189
+ exchangeEnd++;
190
+ }
191
+ if (exchangeEnd > tailStart) {
192
+ // The oldest remaining exchange bleeds into the protected tail. Because
193
+ // the scan is oldest-first, every exchange after this one STARTS inside
194
+ // the tail (this one's result chain reaches it), and every exchange
195
+ // before it was already dropped by earlier iterations — so there is
196
+ // nothing else the guard may remove. Stop.
197
+ break;
198
+ }
199
+ const dropped = estimates
200
+ .slice(exchangeStart, exchangeEnd)
201
+ .reduce((sum, tokens) => sum + tokens, 0);
202
+ result.splice(exchangeStart, exchangeEnd - exchangeStart);
203
+ estimates.splice(exchangeStart, exchangeEnd - exchangeStart);
204
+ currentTokens -= dropped;
205
+ droppedExchanges++;
206
+ }
207
+ if (droppedExchanges > 0) {
208
+ // Insert one elision note where history was removed: after the leading
209
+ // non-exchange messages (typically the first user/task message), but
210
+ // never after the protected tail — when every droppable exchange was
211
+ // removed, an uncapped scan would append the note at the END, where the
212
+ // "history was removed" cue lands after the content it refers to.
213
+ let noteIndex = 0;
214
+ while (noteIndex < result.length &&
215
+ !isToolCallAssistantMessage(result[noteIndex])) {
216
+ noteIndex++;
217
+ }
218
+ const tailBoundary = Math.max(0, result.length - PROTECTED_TAIL_MESSAGES);
219
+ result.splice(Math.min(noteIndex, tailBoundary), 0, {
220
+ role: "user",
221
+ content: [
222
+ {
223
+ type: "text",
224
+ text: `[context truncated: ${droppedExchanges} earlier tool exchange(s) were removed to fit the model's context window. Continue from the remaining context.]`,
225
+ },
226
+ ],
227
+ });
228
+ }
229
+ return { messages: result, droppedExchanges };
230
+ }
231
+ /**
232
+ * Create a per-step budget guard. Returns a function that, given the step's
233
+ * messages, returns a compacted replacement array when the projected request
234
+ * exceeds the threshold — or `undefined` when no change is needed.
235
+ */
236
+ export function createStepBudgetGuard(config) {
237
+ const { provider, model, maxTokens, fixedOverheadTokens = 0, getFixedOverheadTokens, thresholdRatio = DEFAULT_CONTEXT_GUARD_RATIO, } = config;
238
+ const availableInput = getAvailableInputTokens(provider, model, maxTokens);
239
+ const thresholdTokens = Math.floor(availableInput * thresholdRatio);
240
+ return function guardStepMessages(messages) {
241
+ // Resolve overhead per invocation: the tool set can GROW mid-loop
242
+ // (search_tools hydration adds discovered tools between steps), so a
243
+ // once-captured value would undercount later steps.
244
+ const overheadTokens = getFixedOverheadTokens?.() ?? fixedOverheadTokens;
245
+ const estimate = overheadTokens + estimateStepMessagesTokens(messages, provider);
246
+ // Logger Guard: per-step diagnostics for debugging why a long run does
247
+ // (or does not) trigger compaction — gated so nothing is serialized when
248
+ // debug logging is off.
249
+ if (logger.shouldLog("debug")) {
250
+ logger.debug("[StepBudgetGuard] step estimate", {
251
+ provider,
252
+ model,
253
+ messageCount: messages.length,
254
+ estimatedTokens: estimate,
255
+ thresholdTokens,
256
+ willCompact: estimate > thresholdTokens,
257
+ });
258
+ }
259
+ if (estimate <= thresholdTokens) {
260
+ return undefined;
261
+ }
262
+ // Stage 1: shrink old tool outputs to previews.
263
+ const stage1 = truncateOldToolOutputs([...messages]);
264
+ let compacted = stage1.messages;
265
+ let newEstimate = overheadTokens + estimateStepMessagesTokens(compacted, provider);
266
+ // Stage 2: drop oldest complete tool exchanges if still over.
267
+ let droppedExchanges = 0;
268
+ if (newEstimate > thresholdTokens) {
269
+ const stage2 = dropOldestToolExchanges(compacted, thresholdTokens, overheadTokens, provider);
270
+ compacted = stage2.messages;
271
+ droppedExchanges = stage2.droppedExchanges;
272
+ newEstimate =
273
+ overheadTokens + estimateStepMessagesTokens(compacted, provider);
274
+ }
275
+ if (stage1.truncated === 0 && droppedExchanges === 0) {
276
+ return undefined; // nothing actionable (already all-protected)
277
+ }
278
+ logger.info("[StepBudgetGuard] Compacted agent-loop step messages", {
279
+ provider,
280
+ model,
281
+ estimatedTokens: estimate,
282
+ thresholdTokens,
283
+ afterTokens: newEstimate,
284
+ toolOutputsTruncated: stage1.truncated,
285
+ exchangesDropped: droppedExchanges,
286
+ });
287
+ return compacted;
288
+ };
289
+ }
@@ -21,6 +21,7 @@ import { calculateCost } from "../../utils/pricing.js";
21
21
  import { withProviderRetry } from "../../utils/providerRetry.js";
22
22
  import { calculateCacheSavingsPercent, extractCacheCreationTokens, extractCacheReadTokens, extractTokenUsage, } from "../../utils/tokenUtils.js";
23
23
  import { DEFAULT_MAX_STEPS } from "../constants.js";
24
+ import { createStepBudgetGuard, estimateFixedOverheadTokens, } from "../../context/stepBudgetGuard.js";
24
25
  import { isTemperatureDeprecatedError, isSchemaComplexityError, isToolsSchemaConflictError, isToolsSchemaExclusionInForce, } from "./structuredOutputPolicy.js";
25
26
  import { coerceJsonToSchema } from "../../utils/json/coerce.js";
26
27
  import { NoObjectGeneratedError } from "../../utils/generationErrors.js";
@@ -116,6 +117,37 @@ export class GenerationHandler {
116
117
  // rather than passing them inside `messages` (deprecated by the AI SDK,
117
118
  // rejected in v7). See extractSystemMessages for the rationale. (#1024)
118
119
  const { system, messages: nonSystemMessages } = extractSystemMessages(messages);
120
+ // Per-step context budget guard: the tool loop appends assistant turns and
121
+ // tool results on every step — growth the pre-call budget check never
122
+ // sees. Estimate each step's projected request and deterministically
123
+ // reclaim budget (truncate old tool outputs, then drop oldest exchanges)
124
+ // so long agentic runs cannot overflow the model's window mid-loop.
125
+ // Parity with the googleVertex native loops' createContextGuard, upgraded
126
+ // from stop-only to compact-and-continue. The caller's prepareStep result
127
+ // wins on conflicts; the guard only contributes `messages`.
128
+ //
129
+ // Overhead is resolved PER STEP because `toolsWithCache` is deliberately
130
+ // mutable (search_tools hydration adds discovered tools mid-loop) — a
131
+ // once-captured estimate would undercount later steps. Tools are only
132
+ // ever added, so memoizing on tool count keeps the common step O(1).
133
+ let cachedOverhead = { toolCount: -1, tokens: 0 };
134
+ const stepBudgetGuard = createStepBudgetGuard({
135
+ provider: this.providerName ?? "unknown",
136
+ model: this.modelName,
137
+ maxTokens: options.maxTokens,
138
+ getFixedOverheadTokens: () => {
139
+ const toolCount = shouldUseTools
140
+ ? Object.keys(toolsWithCache).length
141
+ : 0;
142
+ if (toolCount !== cachedOverhead.toolCount) {
143
+ cachedOverhead = {
144
+ toolCount,
145
+ tokens: estimateFixedOverheadTokens(system, shouldUseTools ? toolsWithCache : undefined, this.providerName),
146
+ };
147
+ }
148
+ return cachedOverhead.tokens;
149
+ },
150
+ });
119
151
  return await generateText({
120
152
  model,
121
153
  ...(system && { system }),
@@ -125,11 +157,27 @@ export class GenerationHandler {
125
157
  stopWhen: stepCountIs(options.maxSteps ?? DEFAULT_MAX_STEPS),
126
158
  ...(shouldUseTools &&
127
159
  options.toolChoice && { toolChoice: options.toolChoice }),
128
- ...(prepareStep && {
129
- experimental_prepareStep: ((stepOptions) => prepareStep({
130
- ...stepOptions,
131
- maxSteps: options.maxSteps ?? DEFAULT_MAX_STEPS,
132
- })),
160
+ experimental_prepareStep: (async (stepOptions) => {
161
+ // Public contract preserved: a caller-supplied prepareStep receives
162
+ // the ORIGINAL AI-SDK step options, exactly as before the guard
163
+ // existed callers that inspect message history see the real thing.
164
+ const callerResult = prepareStep
165
+ ? await prepareStep({
166
+ ...stepOptions,
167
+ maxSteps: options.maxSteps ?? DEFAULT_MAX_STEPS,
168
+ })
169
+ : undefined;
170
+ // The guard runs LAST, on the messages that will actually be sent:
171
+ // the caller's override when one was returned (out-of-contract for
172
+ // NeuroLink's public prepareStep type, but possible at runtime), else
173
+ // the step's own messages. It never replaces a caller's content
174
+ // choices — it only reclaims budget from whatever was chosen.
175
+ const callerMessages = callerResult?.messages;
176
+ const compacted = stepBudgetGuard(callerMessages ?? stepOptions.messages);
177
+ if (!compacted) {
178
+ return callerResult;
179
+ }
180
+ return { ...(callerResult ?? {}), messages: compacted };
133
181
  }),
134
182
  temperature: options.temperature,
135
183
  maxOutputTokens: options.maxTokens,
@@ -21,12 +21,24 @@ export declare const DEFAULT_OUTPUT_RESERVE_RATIO = 0.35;
21
21
  * The "_default" key is the fallback for unknown models within a provider.
22
22
  */
23
23
  export declare const MODEL_CONTEXT_WINDOWS: Record<string, Record<string, number>>;
24
+ /**
25
+ * Register a runtime-discovered context window for a provider/model pair.
26
+ * Later registrations overwrite earlier ones (rediscovery refreshes values).
27
+ * Non-positive/non-finite windows are ignored so a malformed discovery source
28
+ * can never shrink a budget to zero.
29
+ */
30
+ export declare function registerRuntimeContextWindow(provider: string, model: string, contextWindow: number): void;
31
+ /** Test hook: clear runtime-discovered windows (state is module-global). */
32
+ export declare function clearRuntimeContextWindows(): void;
24
33
  /**
25
34
  * Resolve context window size for a provider/model combination.
26
35
  *
27
36
  * Priority:
28
37
  * 0. Dynamic model registry (DynamicModelProvider) — resolves cross-provider
29
38
  * models (e.g. Claude on Vertex) that the static table cannot handle
39
+ * 0.5 Runtime-discovered windows (registerRuntimeContextWindow) — real
40
+ * per-model limits fetched from the serving infrastructure (LiteLLM
41
+ * `/model/info`)
30
42
  * 1. Exact model match under provider in static registry
31
43
  * 2. Prefix match under provider in static registry
32
44
  * 3. Provider's _default in static registry
@@ -392,12 +392,43 @@ function normalizeProviderForLookup(provider) {
392
392
  // to DEFAULT_CONTEXT_WINDOW.
393
393
  return PROVIDER_ALIAS_MAP[stripped] ?? stripped;
394
394
  }
395
+ /**
396
+ * Runtime-discovered context windows, keyed `${provider}:${model}`.
397
+ *
398
+ * Populated asynchronously by providers that can discover real per-model
399
+ * limits at runtime (e.g. the LiteLLM provider reads `max_input_tokens` from
400
+ * the proxy's `/model/info`), and read synchronously by
401
+ * {@link getContextWindowSize} — the same async-populate/sync-read contract as
402
+ * the DynamicModelProvider registry. Keys use the RAW provider string callers
403
+ * pass into budget calculations (see the alias-map comment above: normalized
404
+ * provider names never reach these lookups).
405
+ */
406
+ const RUNTIME_CONTEXT_WINDOWS = new Map();
407
+ /**
408
+ * Register a runtime-discovered context window for a provider/model pair.
409
+ * Later registrations overwrite earlier ones (rediscovery refreshes values).
410
+ * Non-positive/non-finite windows are ignored so a malformed discovery source
411
+ * can never shrink a budget to zero.
412
+ */
413
+ export function registerRuntimeContextWindow(provider, model, contextWindow) {
414
+ if (!Number.isFinite(contextWindow) || contextWindow <= 0) {
415
+ return;
416
+ }
417
+ RUNTIME_CONTEXT_WINDOWS.set(`${provider}:${model}`, contextWindow);
418
+ }
419
+ /** Test hook: clear runtime-discovered windows (state is module-global). */
420
+ export function clearRuntimeContextWindows() {
421
+ RUNTIME_CONTEXT_WINDOWS.clear();
422
+ }
395
423
  /**
396
424
  * Resolve context window size for a provider/model combination.
397
425
  *
398
426
  * Priority:
399
427
  * 0. Dynamic model registry (DynamicModelProvider) — resolves cross-provider
400
428
  * models (e.g. Claude on Vertex) that the static table cannot handle
429
+ * 0.5 Runtime-discovered windows (registerRuntimeContextWindow) — real
430
+ * per-model limits fetched from the serving infrastructure (LiteLLM
431
+ * `/model/info`)
401
432
  * 1. Exact model match under provider in static registry
402
433
  * 2. Prefix match under provider in static registry
403
434
  * 3. Provider's _default in static registry
@@ -421,6 +452,13 @@ export function getContextWindowSize(provider, model) {
421
452
  // Dynamic registry not initialized yet — fall through to static lookup
422
453
  }
423
454
  }
455
+ // Step 0.5: Runtime-discovered window for this exact provider/model.
456
+ if (model) {
457
+ const discovered = RUNTIME_CONTEXT_WINDOWS.get(`${provider}:${model}`);
458
+ if (discovered !== undefined) {
459
+ return discovered;
460
+ }
461
+ }
424
462
  // Static fallback chain — normalize aliases first so "lmstudio" / "llama.cpp" /
425
463
  // "nvidianim" find their canonical entries instead of falling back to default.
426
464
  const canonical = normalizeProviderForLookup(provider);
@@ -0,0 +1,47 @@
1
+ /**
2
+ * Per-step context budget guard for the AI-SDK agent loop.
3
+ *
4
+ * Pre-call budgeting (`checkContextBudget` + compaction) runs ONCE before
5
+ * dispatch and only sees the input/session conversation. The AI-SDK tool loop
6
+ * then appends assistant turns and tool results on every step — growth the
7
+ * pre-call pipeline never sees, which is how long agentic runs overflow the
8
+ * model's real context window mid-loop (provider 400s after dozens of tool
9
+ * calls). The googleVertex native loops already guard this via
10
+ * `createContextGuard`; this module brings the AI-SDK path (every provider
11
+ * that delegates to `generateText`) to parity — and goes one step further:
12
+ * instead of stopping the loop, it deterministically reclaims budget so the
13
+ * loop can CONTINUE.
14
+ *
15
+ * Wired in `GenerationHandler.callGenerateText` through
16
+ * `experimental_prepareStep`, whose result may replace the step's `messages`.
17
+ * The guard operates on `ModelMessage[]` natively (no lossy ChatMessage
18
+ * round-trip) and never makes LLM calls:
19
+ *
20
+ * Stage 1 — truncate OLD tool outputs to head/tail previews
21
+ * (`generateToolOutputPreview`), oldest first, outside the
22
+ * protected recent tail.
23
+ * Stage 2 — drop the oldest complete tool exchanges (assistant tool-call
24
+ * message + its following tool-result messages, as a unit, so
25
+ * call/result pairing stays intact), replacing them with a single
26
+ * elision note.
27
+ *
28
+ * The system prompt and tool definitions ride OUTSIDE the step messages (the
29
+ * handler hoists system into generateText's `system` option), so their cost is
30
+ * passed in as `fixedOverheadTokens`. The first user message (the task) and
31
+ * the most recent messages are never touched.
32
+ */
33
+ import type { ModelMessage, StepBudgetGuardConfig } from "../types/index.js";
34
+ /** Estimate the token cost of a step's message array. */
35
+ export declare function estimateStepMessagesTokens(messages: readonly ModelMessage[], provider?: string): number;
36
+ /**
37
+ * Estimate the fixed per-request overhead: hoisted system prompt + tool
38
+ * definitions. Mirrors `checkContextBudget`'s categories for the pieces that
39
+ * do not live in the step messages.
40
+ */
41
+ export declare function estimateFixedOverheadTokens(system: unknown, tools: Record<string, unknown> | undefined, provider?: string): number;
42
+ /**
43
+ * Create a per-step budget guard. Returns a function that, given the step's
44
+ * messages, returns a compacted replacement array when the projected request
45
+ * exceeds the threshold — or `undefined` when no change is needed.
46
+ */
47
+ export declare function createStepBudgetGuard(config: StepBudgetGuardConfig): (messages: readonly ModelMessage[]) => ModelMessage[] | undefined;