Skip to content

fix(chat): budget ask-mode history by tokens - #412

Open
herikwebb wants to merge 1 commit into
plmbr:mainfrom
herikwebb:token-aware-chat-history-upstream
Open

fix(chat): budget ask-mode history by tokens#412
herikwebb wants to merge 1 commit into
plmbr:mainfrom
herikwebb:token-aware-chat-history-upstream

Conversation

@herikwebb

Copy link
Copy Markdown
Contributor

Summary

  • budget ask-mode and built-in generation messages against 80% of the active model context window
  • preserve system guidance, injected rules, required inline-edit source, MCP prompt sequences, and the newest request without splitting tool-call turns
  • truncate safe text context, omit protected fenced or multimodal context atomically, and include omission notices
  • load tokenization asynchronously with bounded retries and conservative fallbacks
  • skip pruning when a provider context window is inferred rather than configured; agent tool loops remain unchanged

Validation

  • 1,619 tests passed
  • focused chat-history, dispatch, inline-edit, rules, and provider regression suites passed
  • Cora approved exact commit 6cc3599 with both Codex and Claude
  • one squashed Angular-style commit
  • fork-only review workflows and scripts are excluded

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 pjdoland added the bug Something isn't working label Aug 27, 2026

@pjdoland pjdoland left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 tool message, and no assistant with tool_calls left unanswered
  • output always terminates in a user message
  • system messages 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants