fix(chat): budget ask-mode history by tokens - #412
Conversation
Budget ask-mode and built-in generation messages against 80% of the active model context window. Preserve system guidance and the newest request, prioritize recent current-turn context, retain complete tool-call turns, and signal truncated or omitted content. Load tokenization asynchronously with bounded retries and conservative fallbacks, protect fenced and rule context from partial truncation, and skip pruning when a provider context window is inferred rather than configured. Agent tool loops remain unchanged.
pjdoland
left a comment
There was a problem hiding this comment.
Thank you for this. I reviewed 6cc3599b against merge-base fa6b175 in a clean worktree, at the "blockers only" bar. I did not find a blocker, and I have no changes to request. Since a bare approval is not very useful on a 900-line module landing on the main chat path, here is what I actually exercised, so you and the maintainers can judge how much weight the result carries.
Gates
1551 passed in tests/ plus 68 passed in tests/test_claude_client.py. No TypeScript, TSX, or CSS is touched, so there is no frontend surface to drive with Playwright and I did not fabricate one.
Property testing of budget_chat_messages
Reading a pruning algorithm rarely settles whether it can emit a sequence a provider rejects, so I fuzzed it against invariants instead, with the tokenizer warmed synchronously so the real tiktoken path was under test rather than the byte fallback.
Run 1, 8,000 randomized well-formed histories (multi-round tool turns, image-free, windows from 200 to 64,000 tokens), of which 3,812 actually took a pruning path:
- no orphaned
toolmessage, and noassistantwithtool_callsleft unanswered - output always terminates in a
usermessage systemmessages always remain at the front- no pruned output exceeded
CHAT_INPUT_BUDGET_RATIO * window - zero violations
Run 2, 6,000 trials with mcp_prompt_message_count and required_context_message_count non-zero:
- required-context messages survive verbatim on every pruned path (467 such runs)
- the MCP prompt block, when retained, stays contiguous, in order, and immediately before the newest request (500 such runs)
- relative order of every retained original is preserved
- zero violations
Latency on the request path
| window | messages | payload | real tokens | budgeting cost |
|---|---|---|---|---|
| 128,000 | 82 | 0.16 MB | 27,009 | 3.5 ms |
| 128,000 | 202 | 0.80 MB | 132,317 | 16.6 ms |
| 200,000 | 402 | 1.60 MB | 264,793 | 32.8 ms |
| 1,000,000 | 802 | 4.80 MB | 790,509 | 96.4 ms |
Negligible against the provider round trip it precedes. The byte-count screening pass keeps the common case off the tokenizer entirely, and the second pass re-checks with real token counts before pruning, so the cheap estimate can only cost work, never correctness. That layering is the right way around.
Wiring I checked specifically
The persisted_request_history / request_history split in the GenerateCode branch is the subtle part, and it is correct in a way worth calling out. ChatHistory.get_history returns the live list rather than a copy, so before this change AIServiceManager.handle_chat_request appended the bare prompt straight into persisted session state, leaving a duplicate behind on every inline edit. Slicing to [:-1] incidentally closes that. Keeping the full persisted_request_history for _inline_system_prompt_token_budget, where the length must match what is actually sent, while passing the sliced list as the request history, is exactly right and easy to get backwards.
I traced required_context_message_count=1 through to the message it protects and confirmed the existing-code message is the last of the three context messages in every prefix/suffix combination, so the tail slice lands on it. I confirmed mcp_prompt_message_count still indexes correctly through generate_code_cell and generate_markdown_for_code, where the pop()-then-append() shape preserves the tail structure the slice depends on. I also confirmed that persisted assistant messages carry only content and reasoning_content, never tool_calls, so _partition_turns' tool handling is defensive rather than load-bearing on the ask path, which is the conservative direction.
Relatedly, moving the tiktoken.encoding_for_model('gpt-4o') call out of module scope in both extension.py and rule_injector.py means a host that cannot reach the BPE download now loads the extension in a degraded state rather than failing to import. That is a strict improvement over main, and I mention it because it is a real operational win the summary undersells.
Things I deliberately did not raise
For completeness, since "no findings" should be auditable: the one-byte-per-token fallback runs about four times conservative, so a cold or unreachable tokenizer cache prunes history that would have fit. It is transient behind the warm-up, it errs toward under-filling rather than overflowing, and on main that same deployment did not start at all, so it is not a regression. The flat 4,096-token image allowance and the unchanged agent tool loop are both disclosed scope choices rather than defects. None of these met the bar you set, and I am noting them only so the absence of findings is not mistaken for a shallow pass.
Nice work. The fail-open wrapper, the atomic treatment of fenced and nonced context, and the decision to skip pruning entirely when a provider's context window is inferred rather than configured all read as someone who thought carefully about what happens when the estimate is wrong.
Summary
Validation