feat(experiment): cap per-turn output tokens and steer on truncation - #927
feat(experiment): cap per-turn output tokens and steer on truncation#927pmateusz wants to merge 2 commits into
Conversation
|
Documentation Changes Added
Actions
If neither actions are selected, on PR close/merge the docs branch in ReadMe will remain open. |
Kimchi Code Review
Summary📊 Review Score: 86/100 (overall code quality — 0 lowest, 100 highest) 🧪 Tests: yes — Good test coverage for the new behavior. 📝 Found 4 issue(s). See inline comments for details. What to expectKimchi will analyze the changes in this pull request and post:
The review typically completes within a few minutes. This comment will be updated once the review is ready. Interact with Kimchi
ConfigurationReviews are configured by your organization admin. Powered by Kimchi — AI-powered code review by CAST AI |
![kimchi-review[bot]](https://avatars.githubusercontent.com/in/3223788?s=60&v=4){kind=small}
There was a problem hiding this comment.
📊 Review Score: 86/100 (overall code quality — 0 lowest, 100 highest)
⏱️ Estimated effort to review: 2/5 (1 = trivial, 5 = very complex)
🧪 Tests: yes — Good test coverage for the new behavior. src/extensions/max-output-tokens.test.ts covers env parsing (resolveMaxOutputTokens), payload clamping (applyMaxOutputTokens), headroom arithmetic, fallback prompt-size estimation, and the truncation-steering state machine. src/extensions/agents/manager/agent-runner.test.ts was updated to account for the new factory. The main gap is direct coverage of the before_provider_request hook path that wires ctx.getContextUsage/ctx.model.contextWindow into applyMaxOutputTokens.
📝 Found 4 issue(s). See inline comments for details.
| * untouched — including any limit the caller already set, which is deliberate: | ||
| * "disabled" means "do not interfere", not "remove existing limits". | ||
| */ | ||
| export function applyMaxOutputTokens(payload: unknown, cap: number, limits: ContextLimits = {}): unknown { |
There was a problem hiding this comment.
applyMaxOutputTokens mutates the incoming payload object in place (payload[field] = effective, payload[DEFAULT_FIELD] = effective). In a hook chain the same object may be observed by telemetry or other extensions, and callers may be surprised that their input was modified.
💡 Suggestion: Return a shallow copy instead of mutating the input, e.g. create const result = { ...payload } at the top, mutate result, and return it.
| // the session rather than carrying a grudge from an earlier rough patch. | ||
| let consecutiveTruncations = 0 | ||
|
|
||
| pi.on("turn_start", async () => { |
There was a problem hiding this comment.
The turn_start handler resets consecutiveTruncations only when it already equals or exceeds MAX_CONSECUTIVE_TRUNCATION_STEERS, which contradicts the comment that any user-started turn is a fresh attempt. This causes truncation state to persist across interleaved user turns and can push the model to give up sooner than intended.
💡 Suggestion: Reset consecutiveTruncations = 0 unconditionally inside turn_start, or update the comment to describe the actual threshold-based behavior.
| export function applyMaxOutputTokens(payload: unknown, cap: number, limits: ContextLimits = {}): unknown { | ||
| if (cap <= 0 || !isPayload(payload)) return payload | ||
|
|
||
| const headroom = resolveHeadroom(payload, limits) |
There was a problem hiding this comment.
When headroom is below MIN_USABLE_CAP_TOKENS, applyMaxOutputTokens returns the payload untouched, leaving any pre-existing numeric max_tokens or max_completion_tokens unclamped. If another extension or caller has already set a large limit, the request can still exceed the context window and trigger the vLLM 400 the guard is meant to prevent.
💡 Suggestion: Before the early return, still iterate MAX_TOKEN_FIELDS and clamp or remove any existing numeric limit that is larger than the available headroom, so the guard covers all numeric limits.
| ) | ||
| }) | ||
|
|
||
| pi.on("before_provider_request", async (event, ctx) => { |
There was a problem hiding this comment.
ℹ️🧪 Testing
The before_provider_request hook integration — reading ctx.getContextUsage, falling back to ctx.model.contextWindow, resolving the env-driven cap, and returning a clamped payload — is not directly tested. A regression in how context usage is mapped to ContextLimits would only be caught indirectly.
💡 Suggestion: Add a test using a minimal ExtensionAPI stub that fires before_provider_request with a fake ctx and payload, and asserts the returned payload contains the expected max_tokens value.
Bounds per-turn generation via a before_provider_request extension, with automatic recovery when the cap is hit. Context: TB-2.1 analysis (3,682 trials) found ~19.3M output tokens generated then discarded — the GCLB BackendConfig (timeoutSec: 660) severs any request exceeding 11min regardless of activity. The largest discarded turn was 487k tokens. Nothing bounded this: the thinking dial is advisory and pi never sends an output limit on the openai-completions path (verified by pinning models to 600 in models.json with no effect). The cap is fitted to remaining context (vLLM 400s if prompt + max_tokens > context_window). Truncation is detected and the agent is steered into a new turn instead of silently ending. A consecutive-truncation limit prevents infinite looping. Registered on subagent sessions too — they do the long exploratory work. Default 32,000; KIMCHI_MAX_OUTPUT_TOKENS overrides, 0 disables. Deliberately below the largest legitimate turn observed (40k/106k successful turns) to produce measurable truncation cost; 64k was the zero-collateral value. Does not address the 660s cut-offs themselves — a token cap is not a time bound. 32k prevents ~18% of them.
fc4e783 to
e12e869
Compare
The output cap bounds everything a turn emits, so a limit strict enough to
stop runaway deliberation also severs tool calls. This counts streamed
reasoning and cuts the request when it passes KIMCHI_MAX_THINKING_TOKENS
(default 8k), which leaves acting turns untouched.
There is no server-side equivalent: the gateway accepts
`reasoning: {max_tokens}` and `thinking: {budget_tokens}` with HTTP 200 and
ignores both, and pi's thinkingBudgets is unimplemented for
openai-completions. The extension API has no stream hook either, and its one
lever — ctx.abort() — trips the run's own controller, ending the whole run
before the steering queue is polled (exit 1 under --print). So the wrapper
aborts a private controller and synthesizes a done/"length" event, leaving
the run intact and reading as an ordinary truncation.
Threshold picked from the thinking distribution across 37 stored
terminal-bench sessions: p50 229, p90 4.7k, p99 58k. 8k cuts ~6% of thinking
turns and recovers ~37% of the thinking that survives the 32k output cap.
Also switches the output-cap steer from followUp to steer + triggerTurn. The
follow-up queue drains only once the agent loop runs dry, so under ferment the
guidance waited behind long tool loops and arrived far too late or never. The
steering queue drains after every turn; triggerTurn covers the case where a
truncation left no tool calls and the run would otherwise end.
The truncation counter now resets on user input rather than turn_start:
agent_start opens every agent loop, including the one triggerTurn starts to
carry the steer, so resetting there would zero the counter on the run it is
counting and the give-up would never fire.
Verified against the live gateway on glm-5.2-fp8, minimax-m3 and kimi-k2.7 —
all three stream reasoning_content deltas, truncate at the budget, steer, and
exit 0 with the run intact.
1 participant
Linked issue
Closes #0
What does this PR do?
Cap per-turn output tokens and steer on truncation
Problem
TB-2.1 benchmark analysis found output tokens generated and then discarded because
the BackendConfig (timeoutSec: 660) severs any request exceeding 11 minutes regardless of activity.
The largest single discarded turn was 487k tokens. Nothing bounded this - the thinking dial is advisory,
and pi never sends an output limit on the openai-completions path.
Solution
A before_provider_request extension that:
so a flat cap would introduce a hard 400 that doesn't exist today)
instead of the agent silently treating the truncated response as a finished answer
Registered on subagent sessions as well — they do the long exploratory work most likely to run away.
Configuration
32k is deliberately below the largest legitimate turn observed (40,042 tokens across 106,594 successful
turns) to produce measurable truncation cost in the experiment. 64k was the zero-collateral value.
Checklist
pnpm run test)pnpm run check)