Repository navigation
feat(context): per-step budget guard compacts the AI-SDK tool loop before overflow - #1221
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
📝 WalkthroughWalkthroughAdds a per-step context budget guard that estimates token usage, truncates older tool outputs, removes older tool exchanges when needed, and integrates compaction into ChangesStep budget guard
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant GenerationHandler
participant prepareStep
participant StepBudgetGuard
participant Model
GenerationHandler->>prepareStep: prepare step with original options
prepareStep-->>GenerationHandler: return optional messages
GenerationHandler->>StepBudgetGuard: estimate and compact outgoing messages
StepBudgetGuard-->>GenerationHandler: return guarded messages or undefined
GenerationHandler->>Model: generate text with final messages
Possibly related PRs
Suggested labels: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 ESLint
src/lib/context/stepBudgetGuard.tsParsing error: Unable to parse the specified 'tsconfig' file. Ensure it's correct and has valid syntax. error TS5012: Cannot read file '/.svelte-kit/tsconfig.json': ENOENT: no such file or directory, open '/.svelte-kit/tsconfig.json'. src/lib/core/modules/GenerationHandler.tsParsing error: Unable to parse the specified 'tsconfig' file. Ensure it's correct and has valid syntax. error TS5012: Cannot read file '/.svelte-kit/tsconfig.json': ENOENT: no such file or directory, open '/.svelte-kit/tsconfig.json'. src/lib/types/context.tsParsing error: Unable to parse the specified 'tsconfig' file. Ensure it's correct and has valid syntax. error TS5012: Cannot read file '/.svelte-kit/tsconfig.json': ENOENT: no such file or directory, open '/.svelte-kit/tsconfig.json'. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
✅ Single Commit Policy - COMPLIANTStatus: Policy requirements met • 1 commit • Valid format • Ready for merge 📊 View validation details📝 Commit Details
✅ Validation Results
🤖 Automated validation by NeuroLink Single Commit Enforcement |
🤖 AI Review & Build Compliance ✅Status: AI analysis complete • Build rules validated • Ready for review 📊 View detailed analysis results🛡️ Analysis Complete
📋 Ready for Merge When
🤖 AI analysis complete - check individual code comments for specific feedback |
Tara-ag
left a comment
Tara-ag
left a comment
There was a problem hiding this comment.
Review Summary
Files reviewed: 5 changed files (package.json, src/lib/context/stepBudgetGuard.ts, src/lib/core/modules/GenerationHandler.ts, src/lib/types/context.ts, test/stepBudgetGuard.test.ts)
New issues raised: 7
- 🔒 CRITICAL: 0
⚠️ MAJOR: 2 (both onGenerationHandler.ts:220; related but distinct)- 💡 MINOR: 3
- 💬 SUGGESTION: 2
Decision: REQUEST_CHANGES
The implementation is solid conceptually and the tests cover the core compaction logic well, but the experimental_prepareStep wrapper in GenerationHandler.ts changes the public prepareStep contract in a backward-incompatible way:
-
Backward-compatibility break (blocking): The wrapper now passes guard-compacted
messagesinto the caller’sprepareStepcallback. Callers that inspect step messages (even if untyped) will see different inputs than before. Per CLAUDE.md Critical Rule 5, the public SDK API must not break existing callers without a documented migration path. The fix is straightforward: run the caller’sprepareStepon the originalstepOptions.messages, then merge the guard’s compaction afterward. -
Silent override risk: The merge
{ ...(callerResult ?? {}), messages: compacted }discards anymessagesthe caller returns without warning. Combined with the first issue, this makes the interaction between caller-supplied and guard-supplied message mutation unclear.
Other non-blocking notes:
- Stage-2 token estimation is O(n²); can be made linear with a running total.
- The elision note can be appended after the protected tail when all droppable exchanges are removed.
- The new Vitest test file conflicts with CONTRIBUTING.md’s stated convention that
*.test.tsfiles are not picked up; consider adding acontinuous-test-suite-step-budget-guard.tsorchestrator as the canonical test.
Please address the prepareStep wrapping order so callers continue to receive the original step messages, and re-request review.
7c126c1 to
5b8fd7d
Compare
🤖 AI Review & Build Compliance ✅Status: AI analysis complete • Build rules validated • Ready for review 📊 View detailed analysis results🛡️ Analysis Complete
📋 Ready for Merge When
🤖 AI analysis complete - check individual code comments for specific feedback |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
test/stepBudgetGuard.test.ts (1)
219-243: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssertions never execute — test is vacuous.
With 4 messages and
PROTECTED_TAIL_MESSAGES = 4,tailStartis 0 in both compaction stages, so nothing is ever eligible for truncation/dropping andguardStepMessagesalways returnsundefined(thestage1.truncated===0 && droppedExchanges===0short-circuit). Theif (compacted) {...}block therefore never runs, and this test always passes regardless of whether the underlying protected-tail logic is correct.✅ Proposed fix: assert the deterministic no-op outcome directly
const compacted = guard(messages); - // The last exchange is within the 4-message protected tail — the guard - // must not mutate it, so either nothing changes or nothing structural does. - if (compacted) { - expect(compacted.length).toBe(messages.length); - const lastTool = compacted[compacted.length - 1] as { - content: Array<{ output: { value: string } }>; - }; - expect(lastTool.content[0].output.value.length).toBe(50_000); - } + // The whole exchange sits inside the 4-message protected tail — the + // guard has nothing it is allowed to touch, so it must no-op. + expect(compacted).toBeUndefined();🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/stepBudgetGuard.test.ts` around lines 219 - 243, Update the test case around the guard invocation to assert the deterministic no-op result directly: verify that guard(messages) returns undefined, and remove the unreachable if (compacted) block and its nested structural/content assertions. Keep the existing four-message protected-tail setup unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/lib/context/stepBudgetGuard.ts`:
- Around line 139-165: Update the tool-result handling in the message content
mapping to serialize the actual output value using the appropriate variant-aware
serialization, rather than stringifying the structured result wrapper. Ensure
the resulting payload text is passed to generateToolOutputPreview while
preserving the existing size checks, truncation counters, and replacement
structure.
In `@src/lib/core/modules/GenerationHandler.ts`:
- Around line 194-213: Update the step budget guard setup around
createStepBudgetGuard so fixedOverheadTokens is recalculated from the current
toolsWithCache on every experimental_prepareStep invocation. Ensure newly
hydrated search_tools are included before enforcing the per-step budget, while
preserving the existing guard behavior and prepareStep message precedence.
---
Nitpick comments:
In `@test/stepBudgetGuard.test.ts`:
- Around line 219-243: Update the test case around the guard invocation to
assert the deterministic no-op result directly: verify that guard(messages)
returns undefined, and remove the unreachable if (compacted) block and its
nested structural/content assertions. Keep the existing four-message
protected-tail setup unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: da25a213-15b5-4ba6-937a-e3c6ab31d506
📒 Files selected for processing (5)
package.jsonsrc/lib/context/stepBudgetGuard.tssrc/lib/core/modules/GenerationHandler.tssrc/lib/types/context.tstest/stepBudgetGuard.test.ts
5b8fd7d to
b638eb3
Compare
🤖 AI Review & Build Compliance ✅Status: AI analysis complete • Build rules validated • Ready for review 📊 View detailed analysis results🛡️ Analysis Complete
📋 Ready for Merge When
🤖 AI analysis complete - check individual code comments for specific feedback |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
test/stepBudgetGuard.test.ts (1)
287-296: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert that the guard explicitly bails out for protected tails.
The
if (compacted)block correctly catches unwanted mutations, but it misses the opportunity to assert the guard's actual bailout contract. When all exchanges fall within the protected tail and no budget can be safely reclaimed, the guard is designed to returnundefined.Replacing the conditional block with a strict
toBeUndefined()assertion makes the test simpler and enforces the expected return type directly.🛠️ Proposed fix for the assertion
const compacted = guard(messages); - // The last exchange is within the 4-message protected tail — the guard - // must not mutate it, so either nothing changes or nothing structural does. - if (compacted) { - expect(compacted.length).toBe(messages.length); - const lastTool = compacted[compacted.length - 1] as { - content: Array<{ output: { value: string } }>; - }; - expect(lastTool.content[0].output.value.length).toBe(50_000); - } + // The last exchange is within the protected tail. The guard must not mutate + // it, and should return `undefined` to signal that no changes were made. + expect(compacted).toBeUndefined();🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/stepBudgetGuard.test.ts` around lines 287 - 296, Update the protected-tail test around guard(messages) to assert directly that the result is undefined when all exchanges are protected and no budget can be reclaimed. Remove the conditional structural/content assertions, preserving the existing setup and message expectations.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@test/stepBudgetGuard.test.ts`:
- Around line 152-170: Update the pairing check in the result validation loop to
remove each toolCallId from the tracking set when its matching tool-result is
encountered, while still asserting the result has a preceding call. After
processing all messages, assert that the set is empty so retained tool-calls
cannot remain orphaned.
---
Nitpick comments:
In `@test/stepBudgetGuard.test.ts`:
- Around line 287-296: Update the protected-tail test around guard(messages) to
assert directly that the result is undefined when all exchanges are protected
and no budget can be reclaimed. Remove the conditional structural/content
assertions, preserving the existing setup and message expectations.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 93fcfeba-5fb4-4146-83c5-02ba58caaf4e
📒 Files selected for processing (5)
package.jsonsrc/lib/context/stepBudgetGuard.tssrc/lib/core/modules/GenerationHandler.tssrc/lib/types/context.tstest/stepBudgetGuard.test.ts
🚧 Files skipped from review as they are similar to previous changes (4)
- src/lib/types/context.ts
- package.json
- src/lib/context/stepBudgetGuard.ts
- src/lib/core/modules/GenerationHandler.ts
| // Pairing integrity: every remaining tool message's toolCallId has a | ||
| // preceding assistant tool-call with the same id. | ||
| const callIds = new Set<string>(); | ||
| for (const message of result) { | ||
| if (!Array.isArray(message.content)) { | ||
| continue; | ||
| } | ||
| for (const part of message.content as Array<{ | ||
| type?: string; | ||
| toolCallId?: string; | ||
| }>) { | ||
| if (part.type === "tool-call" && part.toolCallId) { | ||
| callIds.add(part.toolCallId); | ||
| } | ||
| if (part.type === "tool-result" && part.toolCallId) { | ||
| expect(callIds.has(part.toolCallId)).toBe(true); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Verify that no orphaned tool-calls remain.
The current pairing integrity check verifies that every retained tool-result has a preceding tool-call, but it doesn't verify the reverse: that every retained tool-call has a matching tool-result. If the guard dropped a tool-result but left its tool-call hanging, this test would falsely pass.
Remove matched IDs from the set and assert that the set is empty at the end to ensure complete bidirectional pairing.
🛠️ Proposed fix to ensure bidirectional pairing
if (part.type === "tool-result" && part.toolCallId) {
expect(callIds.has(part.toolCallId)).toBe(true);
+ callIds.delete(part.toolCallId);
}
}
}
+
+ // Ensure no tool-calls were left without matching tool-results.
+ expect(callIds.size).toBe(0);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Pairing integrity: every remaining tool message's toolCallId has a | |
| // preceding assistant tool-call with the same id. | |
| const callIds = new Set<string>(); | |
| for (const message of result) { | |
| if (!Array.isArray(message.content)) { | |
| continue; | |
| } | |
| for (const part of message.content as Array<{ | |
| type?: string; | |
| toolCallId?: string; | |
| }>) { | |
| if (part.type === "tool-call" && part.toolCallId) { | |
| callIds.add(part.toolCallId); | |
| } | |
| if (part.type === "tool-result" && part.toolCallId) { | |
| expect(callIds.has(part.toolCallId)).toBe(true); | |
| } | |
| } | |
| } | |
| // Pairing integrity: every remaining tool message's toolCallId has a | |
| // preceding assistant tool-call with the same id. | |
| const callIds = new Set<string>(); | |
| for (const message of result) { | |
| if (!Array.isArray(message.content)) { | |
| continue; | |
| } | |
| for (const part of message.content as Array<{ | |
| type?: string; | |
| toolCallId?: string; | |
| }>) { | |
| if (part.type === "tool-call" && part.toolCallId) { | |
| callIds.add(part.toolCallId); | |
| } | |
| if (part.type === "tool-result" && part.toolCallId) { | |
| expect(callIds.has(part.toolCallId)).toBe(true); | |
| callIds.delete(part.toolCallId); | |
| } | |
| } | |
| } | |
| // Ensure no tool-calls were left without matching tool-results. | |
| expect(callIds.size).toBe(0); |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/stepBudgetGuard.test.ts` around lines 152 - 170, Update the pairing
check in the result validation loop to remove each toolCallId from the tracking
set when its matching tool-result is encountered, while still asserting the
result has a preceding call. After processing all messages, assert that the set
is empty so retained tool-calls cannot remain orphaned.
|
🎉 This PR is included in version 10.3.0 🎉 The release is available on: Your semantic-release bot 📦🚀 |
This branch was successfully deployed
3 participants
Problem
Pre-call budgeting (
checkContextBudget+ compaction inneurolink.ts) runs ONCE before dispatch and only sees the input/session conversation. The AI-SDK agent loop (GenerationHandler.callGenerateText, up toDEFAULT_MAX_STEPS=200steps) then appends assistant turns and tool results on every step — growth the pre-call pipeline never sees. Long agentic runs overflow the model's real window mid-loop: juspay/yama#61 CI accumulated ~246K input tokens against a 262K window (and double NeuroLink's own 128K litellm default) before the backend 400'd withContextWindowExceededError.The googleVertex native loops already guard this (
createContextGuard, checked every iteration,DEFAULT_CONTEXT_GUARD_RATIO0.85) — but as stop-only. The AI-SDK path, anthropic/bedrock native loops, and the chat-completions stream loop have nothing.Fix
context/stepBudgetGuard.ts+ wiring incallGenerateTextviaexperimental_prepareStep(whose result may replace the step'smessages— the sanctioned AI-SDK hook). Every step:getAvailableInputTokens() × DEFAULT_CONTEXT_GUARD_RATIO— same ratio as the Vertex guard.generateToolOutputPreview); Stage 2 drops the oldest complete tool exchanges (assistant tool-call + its tool-result messages as a unit, preserving call/result pairing) and inserts one elision note.Operates on
ModelMessage[]natively (no lossy ChatMessage round-trip), makes no LLM calls, and composes with a caller-suppliedprepareStep(caller result wins; the guard only contributesmessages).Tests
test/stepBudgetGuard.test.ts(7 tests, wired intotest:unit): estimation scaling, overhead accounting, no-op under threshold, stage-1 truncation with protected tail, stage-2 pair-integral drops with elision note, fixed-overhead trigger, protected-tail immunity.test:anthropic-tools-policyandtest:mcp:limitspass unchanged.🤖 Generated with Claude Code
Summary by CodeRabbit