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

[herikwebb]

Summary

• budget ask-mode and built-in generation requests against 80% of the active model context window
• preserve system instructions, the newest user request, recent complete tool-call turns, and newest current-turn context
• truncate oversized mandatory and multimodal context with visible omission notices while preserving rule and fenced-context boundaries
• load tokenization outside the request path with bounded retries and conservative fallbacks
• prune only when providers expose a configured context window; agent tool loops remain unchanged

Validation

• .venv/bin/python -m pytest -q tests/test_chat_history_budget.py tests/test_base_chat_participant_integration.py tests/test_websocket_handler_integration.py tests/test_github_copilot_llm_provider.py tests/test_litellm_compatible_llm_provider.py tests/test_openai_compatible_llm_provider.py (122 passed)
• .venv/bin/python -m pytest -q (1,613 passed; one existing asyncio subprocess cleanup warning)
• git diff --check origin/main...HEAD
• fork and upstream feature patches are byte-identical

[github-actions]

Automated PR Review (Claude)

Model: claude-opus-5
Verdict: CHANGES_REQUESTED

Changed files:

CHANGELOG.md
notebook_intelligence/base_chat_participant.py
notebook_intelligence/chat_history_budget.py
notebook_intelligence/extension.py
tests/test_base_chat_participant_integration.py
tests/test_chat_history_budget.py

Review result:

Review

Overall the extraction of the token-budget helpers out of extension.py is a nice cleanup, and the new module is readable and unit-tested. A few issues in the new chat_history_budget.py look worth addressing before merge.

---

1. _complete_turns() can produce malformed message sequences for tool-using histories — Medium

notebook_intelligence/chat_history_budget.py

if role == "assistant":
    turns.append(current_turn)
    current_turn = []

A turn is closed at the first assistant message. For a history of the shape

user → assistant(tool_calls=[...]) → tool → tool → assistant(final text)

the function keeps [user, assistant(tool_calls)] as a "complete turn", then discards the tool messages and the final assistant answer as "orphans" (they don't start with role == "user"). Sending an assistant message containing tool_calls without the matching role: "tool" replies is rejected by OpenAI-compatible APIs (HTTP 400), and dropping the final assistant text silently loses the actual answer.

Even if today's ask-mode history never carries tool messages, chat_history is keyed per chatId and shared across mode switches (agent → ask in the same conversation), so this is reachable. Suggest closing a turn only on an assistant message without tool_calls, and treating role in ("tool", "assistant") continuations as part of the open turn.

Related: the same normalization drops any mid-history system message that isn't in the leading run, and it fires even when the history is under budget (needs_pruning is also true when normalized_past_messages != past_messages), so messages are removed and the omission notice is injected with no size pressure at all. That side effect is asserted by test_pruning_drops_an_orphaned_assistant_message, but it's surprising behavior for a "budget" function.

---

2. Eager tiktoken.encoding_for_model() at module import — Medium

notebook_intelligence/chat_history_budget.py

_encoding = tiktoken.encoding_for_model("gpt-4o")

base_chat_participant.py now imports this module, so the encoding is constructed during server-extension import. encoding_for_model downloads the BPE file over HTTPS on first use when the tiktoken cache is cold, so this can block or raise (offline/air-gapped/proxied deployments) at import time — and an exception here now takes down base_chat_participant, not just a helper. This also cuts against the lazy-provider-import work landing in the same release (Provider SDKs load on first use instead of at module import). Please make it lazy, e.g. an @lru_cache accessor called from text_token_count/truncate_text.

---

3. truncate_text() worst-case cost and a behavior change for tiny budgets — Low

notebook_intelligence/chat_history_budget.py

while prefix_size > 0:
    prefix = _encoding.decode(encoded[:prefix_size]).rstrip()
    candidate = prefix + marker
    if text_token_count(candidate) <= token_budget:
        return candidate
    prefix_size -= 1

Each iteration re-encodes the whole candidate, so the loop is O(n²) in the pathological case. It normally converges in one or two iterations, but a decrement of 1 with a full re-encode is an easy foot-gun on large context blobs; a small backoff step or binary search would bound it.

Also note the semantics change for extension.py's _truncate_context_content: the old implementation returned prefix + "\n...[truncated]" (possibly exceeding token_budget), the new one reserves marker tokens and returns "" when token_budget <= marker_tokens (~6 tokens). Small budgets that previously yielded a short snippet now yield nothing. Worth confirming no existing test/caller depends on the old output.

---

4. Silent over-budget result when the mandatory messages alone don't fit — Low

budget_chat_messages clamps remaining_budget at 0 but never truncates the system prompt or the newest user message, so an oversized single prompt still goes out over budget and fails at the provider with an opaque error. A log.warning when remaining_budget clamps to 0 would make this diagnosable (the module currently has no logger at all).

---

5. Scope of the fix vs. the CHANGELOG claim — Low

CHANGELOG.md, notebook_intelligence/base_chat_participant.py

The entry says "Chat requests now fit prior conversation and current context to the active model's input budget", but budget_chat_messages is only applied on the ask-mode paths (generate_code_cell, generate_markdown_for_code, newPythonFile, plain ask completion). Agent mode (handle_chat_request_with_tools) — the path most likely to blow the window, since tool results accumulate — is untouched. Either narrow the changelog wording to ask mode or state the agent path as follow-up.

---

6. Test coverage gaps — Low

tests/test_chat_history_budget.py

Good coverage of the plain text paths. Missing:

• multimodal content (list / {"type": "image_url"}) exercising IMAGE_TOKEN_ESTIMATE and the dict/list branches of _content_token_count;
• histories containing tool_calls / role: "tool" (see finding 1) — a test here would have caught the orphan behavior;
• the extension.py delegation (_truncate_context_content / _token_count) after the refactor, including the sub-marker-budget case.

IMAGE_TOKEN_ESTIMATE = 1024 is also an undercount for high-detail images on several providers; underestimating pushes the request over the real window. Not blocking, but a comment on the assumption would help.

---

Residual risk

I could not see handle_chat_request_with_tools or how ChatHistory records assistant/tool messages in this diff, so finding 1's exploitability depends on whether tool-role entries ever reach request.chat_history; the fix is cheap either way.

[github-actions]

Automated PR Review (Claude)

Model: claude-opus-5
Verdict: CHANGES_REQUESTED

Changed files:

CHANGELOG.md
notebook_intelligence/base_chat_participant.py
notebook_intelligence/chat_history_budget.py
notebook_intelligence/extension.py
tests/test_base_chat_participant_integration.py
tests/test_chat_history_budget.py
tests/test_websocket_handler_integration.py

Review result:

Review — Ask-mode chat history budgeting

Overall the new chat_history_budget module is well structured (turn-complete pruning, mandatory system + newest user, deterministic image estimate) and the unit tests cover the main branches. A few issues around failure handling and where the tokenizer is now first loaded look worth addressing before merge.

---

Findings

1. budget_chat_messages() is called outside the guard that guarantees response.finish() — Medium

File: notebook_intelligence/base_chat_participant.py (handle_ask_mode_chat_request)

messages = [...] + request.chat_history
messages = budget_chat_messages(messages, chat_model.context_window)

try:
    ...
    chat_model.completions(messages, response=response, cancel_token=...)
except Exception as e:
    ...
    response.finish()

Any exception raised inside budget_chat_messages (see #2 for the most likely cause, but also a malformed/non-dict history entry) escapes before the try: block, so response.finish() is never emitted and the chat UI is left spinning indefinitely. The same pattern applies to the newPythonFile branch and to generate_code_cell / generate_markdown_for_code, where a raise now aborts a /newNotebook flow after the notebook has already been created via run_ui_command.

Suggested fix: make the budgeter fail-open (wrap the body and return list(messages) on unexpected errors, logging once), and/or move the call inside the existing try:.

2. First tokenizer load moved from import time to request time, onto the IOLoop — Medium

Files: notebook_intelligence/chat_history_budget.py (_get_encoding), notebook_intelligence/extension.py

tiktoken.encoding_for_model("gpt-4o") was previously evaluated at module import; it is now lazily evaluated on first use. In extension.py, _token_count / _truncate_context_content are invoked from WebsocketCopilotHandler.on_message, which runs on the Tornado IOLoop. encoding_for_model performs blocking filesystem and, when the BPE file is not already cached, blocking HTTP fetch of the o200k_base vocabulary. That means:

• On an air-gapped or egress-restricted deployment the failure surface moved from "extension fails to load loudly at startup" to "the first chat request raises inside the handler" (which combines badly with security: reject forged upload paths outside the NBI upload directory #1), and
• On a cold cache the whole Jupyter server event loop can stall for the duration of the download.

Note this is the same class of problem the changelog entry for plmbr#370 called out for litellm's model-cost-map fetch. Consider warming _get_encoding() during extension startup (in a thread), and/or catching the load failure inside _get_encoding and degrading to a no-budgeting path.

3. Over-budget mandatory context is returned unmodified — Low

File: notebook_intelligence/chat_history_budget.py

When base_mandatory_tokens > input_budget the function logs a warning and then returns a message list that is knowingly larger than the model window, so the request still fails downstream with a provider-side context-length error. The changelog claims oversized sessions "degrade predictably instead of overflowing the model window", which isn't true for the single-huge-prompt case. Either truncate the newest user message as a last resort, or surface a user-visible error rather than letting the provider 400 bubble up as the generic "Oops!" message.

4. Cross-provider token estimation is gpt-4o-only — Low

text_token_count uses the o200k_base encoding for every provider (Anthropic, Gemini, Ollama models, …), where real token counts can run 10–30% higher. The 20% output reserve absorbs some of this, but the reserve is doing double duty (output headroom and tokenizer skew). Worth a comment/TODO at minimum so the ratio isn't tuned down later without accounting for it.

5. Repeated full-history encoding per request — Low (performance)

budget_chat_messages re-encodes content several times: once for total_tokens, again for base_mandatory_tokens, up to twice more for the two omission-notice candidates, once more for mandatory_messages, and then per-message for context/turns. Similarly truncate_text binary-searches with a full encode(prefix + marker) per iteration (~2–3× the old single-decode cost, now on the IOLoop path in extension.py). Memoizing per-message estimates (e.g. an id()-keyed dict local to the call) would cut this materially for long histories with large attached file contexts.

6. Test-coverage gaps — Low

• No test exercises the new budgeting in generate_code_cell, generate_markdown_for_code, or the newPythonFile branch — only handle_ask_mode_chat_request.
• No test for tokenizer-load failure / fail-open behavior (directly relevant to security: reject forged upload paths outside the NBI upload directory #1 and fix(server): reject forged upload context paths #2).
• No test that the emitted sequence remains API-valid for an interleaved edge case such as [assistant(tool_calls), user, tool], where the trailing tool message is mandatory but its originating assistant turn is discarded by _partition_turns — that would produce a tool message with no preceding tool_calls and a provider 400.

---

Nits

• CHANGELOG.md: the new bullet is the only one in the section without an issue/PR reference; the surrounding entries all carry (#NNN).
• chat_history_budget.py: _partition_turns silently drops leading non-user messages and unfinished tool exchanges. This is intentional and tested, but the "over budget only" precondition is what keeps it safe — a short comment on budget_chat_messages noting that no normalization happens on the under-budget fast path would help future readers.

---

Residual risk not covered above: behavior with providers whose context_window reports total (input+output) vs. input-only capacity, and the fixed IMAGE_TOKEN_ESTIMATE = 1024 under-counting high-resolution attachments — both are acknowledged in comments and bounded by the 20% reserve.

[github-actions]

Automated PR Review (Claude)

Model: claude-opus-5
Verdict: CHANGES_REQUESTED

Changed files:

CHANGELOG.md
notebook_intelligence/base_chat_participant.py
notebook_intelligence/chat_history_budget.py
notebook_intelligence/extension.py
tests/test_base_chat_participant_integration.py
tests/test_chat_history_budget.py
tests/test_websocket_handler_integration.py

Review result:

Review

Overall the refactor is well-structured (shared budgeting module, fail-open wrapper, lazy tokenizer). A few things look worth addressing before merge.

---

1. truncate_text is O(n·log n) tokenization on the event-loop path — Medium (performance)

File: notebook_intelligence/chat_history_budget.py (truncate_text), consumed by notebook_intelligence/extension.py as _truncate_context_content

The previous implementation encoded the content once, sliced the token list, and decoded once. The replacement:

• calls text_token_count(text) (full encode),
• then binary-searches prefix sizes, and for each probe decodes the prefix and re-encodes prefix + marker via text_token_count.

Worst case the probed prefix lengths sum to roughly n·log2(n) tokens of encode work, so this is ~1–2 orders of magnitude more tokenizer work than before for a large attachment. This matters because _truncate_context_content is called from WebsocketCopilotHandler.on_message, which runs on the Jupyter server's IOLoop thread — a multi-megabyte file attachment can now stall the whole server for a noticeable interval, where before it was a single pass.

Cheap fix that keeps the "strict" ≤ budget guarantee the new test asserts: encode once, compute marker_tokens = text_token_count(marker) once, slice encoded[: budget - marker_tokens], decode + rstrip(), and only fall back to a short correction loop (a handful of iterations) if the rstrip/re-encode pushes it over. That's O(n) with a bounded constant.

---

2. A transient tokenizer-load failure is cached for the process lifetime — Low/Medium (reliability)

File: notebook_intelligence/chat_history_budget.py (_get_encoding)

@lru_cache(maxsize=1) caches the None return as well. tiktoken.encoding_for_model fetches the BPE file over the network on first use in environments without a vendored/warm cache, so a single startup blip (which the new warm_tokenizer_encoding() thread makes more likely to be the first call, at server start, before any user activity) permanently degrades every subsequent request to the len(utf8)//3 estimator — with no further warnings after the first. Consider only caching successful loads (e.g., cache a sentinel and retry on failure, or bound the retry to N attempts), and/or not warming until the first real use.

---

3. Over-budget requests still go out unbudgeted in two paths — Low (correctness / fail-open)

File: notebook_intelligence/chat_history_budget.py (_budget_chat_messages)

• If the newest user message has non-string content (multimodal list — exactly the shape produced for pasted images in the non-Claude path), _truncate_text_message returns None, so the oversized request is sent as-is and fails at the provider. Handling the list case (truncate the text parts, drop image parts) would make the "mandatory context" guarantee real for multimodal turns.
• latest_user_index is None returns the original list unchanged even though we already know total_tokens > input_budget.

Both are documented as fail-open, but they're the cases where the user actually gets a hard provider error rather than a degraded answer.

---

4. Current-turn context selection is oldest-first with an early break — Low (behavior)

File: notebook_intelligence/chat_history_budget.py (selected_context loop)

The loop walks current_context forward and breaks on the first message that doesn't fit (including when _truncate_text_message returns None for a multimodal/image context). So one large early attachment can evict all later current-turn context, even small text items that would have fit. Given that the current-directory/current-cell markers and later @-mentions are appended after earlier attachments, this can drop the most task-relevant context. Worth at least a comment justifying the ordering, or continue-on-oversize instead of break.

---

5. Test coverage gaps — Low

• No test asserts initialize_handlers actually calls warm_tokenizer_encoding() (the only production wiring of it).
• No test pins the new invariant that agent mode / handle_chat_request_with_tools is not budgeted, which is the explicit claim in the CHANGELOG entry.
• test_builtin_generation_paths_budget_messages stubs budget_chat_messages out, so it verifies the call site but not that the generation paths still work end-to-end with real pruning applied.
• tests/test_chat_history_budget.py mutates process-global state (_get_encoding.cache_clear() plus monkeypatched tiktoken.encoding_for_model); the finally: cache_clear() leaves the cache empty for subsequent tests, which will then trigger a real tiktoken load. Not wrong, but it makes those tests order-sensitive if anything else asserts on load counts later.

---

Verification request (couldn't confirm from the diff)

extension.py deletes the module-level tiktoken_encoding global. The diff only shows the two helpers that used it being removed; please confirm there are no other references to tiktoken_encoding remaining in that file (a leftover would be a NameError at request time, not import time).

---

Residual risks acknowledged in the code comments and accepted here: the single gpt-4o tokenizer across all providers and the fixed IMAGE_TOKEN_ESTIMATE = 1024 will under-count for some providers/high-detail images; the 20% output reserve absorbs some but not all of that.

[github-actions]

Automated PR Review (Claude)

Model: claude-opus-5
Verdict: CHANGES_REQUESTED

Changed files:

CHANGELOG.md
notebook_intelligence/base_chat_participant.py
notebook_intelligence/chat_history_budget.py
notebook_intelligence/extension.py
tests/test_base_chat_participant_integration.py
tests/test_chat_history_budget.py
tests/test_websocket_handler_integration.py

Review result:

Review

Overall this is a well-structured, well-tested addition: the budgeter fails open, the tokenizer load moved off the import path, and the new test file covers most of the pruning matrix. A few things I'd like resolved or explicitly confirmed before merge.

---

1. PEP 604 annotations may break the supported Python floor — notebook_intelligence/chat_history_budget.py (Medium/High, needs verification)

_tokenizer_encoding: Any | None = None
...
def _truncate_text_message(message: dict, token_budget: int) -> dict | None:
def _truncate_mandatory_message(...) -> dict | None:

Both the module-level variable annotation and the function return annotations are evaluated at runtime (there is no from __future__ import annotations in this file). On Python 3.9, typing.Any | None raises TypeError: unsupported operand type(s) for |, which would make import notebook_intelligence.chat_history_budget — and therefore base_chat_participant and extension — fail outright.

The surrounding code (base_chat_participant.py uses Union[ToolPreInvokeResponse, None], extension.py imports Optional, Union) suggests the project may still target <3.10. Please confirm requires-python in pyproject.toml; if it allows 3.9, switch to Optional[...]/Union[...] or add from __future__ import annotations.

2. warm_tokenizer_encoding() can raise into the hot path — notebook_intelligence/chat_history_budget.py (Low/Medium)

    try:
        thread.start()
    except Exception:
        with _tokenizer_lock:
            _tokenizer_load_in_progress = False
        raise

_get_encoding() calls warm_tokenizer_encoding() lazily, so a RuntimeError("can't start new thread") (thread exhaustion) propagates out of text_token_count(). Inside budget_chat_messages that is caught by the fail-open wrapper, but extension.py re-exports this same function as _token_count and calls it directly in WebsocketCopilotHandler.on_message (the output-context sizing path) with no guard — an exception there aborts handling of the client message on the IOLoop thread. Given the module's stated fail-open philosophy, the lazy path should swallow the start failure and fall back to _fallback_text_token_count; re-raising is only reasonable for the explicit startup warm-up call.

3. Truncation budget ignores non-content message fields — notebook_intelligence/chat_history_budget.py (Low)

estimate_message_tokens() adds tokens for name, tool_call_id, and tool_calls, but _truncate_text_message() only subtracts MESSAGE_OVERHEAD_TOKENS:

content_budget = token_budget - MESSAGE_OVERHEAD_TOKENS

For a message carrying any of those keys, the truncated result's estimate exceeds token_budget, so remaining_budget -= message_tokens(truncated) can go negative and the final selection can overshoot the input budget. Not fatal (subsequent truncations return "" and the turn loop breaks), but it undermines the invariant the tests assert (sum(estimate_message_tokens(...)) <= input_budget) for tool-bearing messages — and there is no test covering that case.

4. Returning [] on "no user request" turns a degradation into a hard provider error — notebook_intelligence/chat_history_budget.py (Low)

    if latest_user_index is None:
        log.warning("Dropping over-budget chat history because it has no user request")
        return []

Every other failure mode in this module falls open (over-budget mandatory content is still sent). Here an empty message list is handed to chat_model.completions(), which every provider rejects. Practically unreachable for the current call sites, but returning list(messages) (or system + truncated tail) would be consistent with the module contract and cheaper to debug if it ever fires.

5. Full history is re-tokenized on every ask-mode turn (Low, performance)

_budget_chat_messages computes total_tokens = sum(message_tokens(m) for m in messages) before the under-budget short-circuit, so a session with a large accumulated history re-encodes everything each turn (O(n²) over a session). This runs on the request worker thread rather than the IOLoop, so it isn't a server-wide stall, but a cheap upper-bound pre-check (e.g. total UTF-8 byte length vs. input_budget * 4) before exact counting would avoid the work in the common case.

---

Minor / non-blocking

• _truncate_mandatory_message silently converts a multimodal newest-user message to text and replaces the image with [Image omitted to fit model context.]. Only the in-history notice tells the user; consider whether this warrants a MarkdownData/progress note in the ask-mode path.
• The CHANGELOG entry is the only one in the Unreleased section without an issue/PR reference.
• Test coverage is helper-level for extension.py; the new _token_count/_truncate_context_content aliases are exercised directly but there's no end-to-end assertion that on_message's additional-context budgeting still behaves as before the refactor (e.g. that truncate_text returning "" for tiny budgets — a behavior change from the old always-non-empty truncator — is handled by _build_additional_context_message).

Residual risks

• The tokenizer fallback (~3 bytes/token) is only exercised in tests via a synthetic character encoder; real-world behavior during the warm-up window (first requests after server start, before the background load completes) isn't covered.
• Silent behavior change if any code outside the shown hunks still references the removed module-level extension.tiktoken_encoding; I couldn't verify that from the diff context.

[github-actions]

Automated PR Review (Claude)

Model: claude-opus-5
Verdict: CHANGES_REQUESTED

Changed files:

CHANGELOG.md
notebook_intelligence/base_chat_participant.py
notebook_intelligence/chat_history_budget.py
notebook_intelligence/extension.py
pyproject.toml
tests/test_base_chat_participant_integration.py
tests/test_chat_history_budget.py
tests/test_websocket_handler_integration.py

Review result:

Summary

The new chat_history_budget.py is well structured and has good unit coverage. Most of my comments are on edge-path behavior in newly added code; one of them (silently discarding image parts from the mandatory newest user message) is worth fixing before merge because it can make the model answer confidently about content it never received.

---

1. Medium — Images are silently dropped from the mandatory user message with no signal to the model

notebook_intelligence/chat_history_budget.py — _truncate_mandatory_message

    text_parts = []
    for item in content:
        ... collect only {"type": "text"} / str parts ...
    if not text_parts:
        text_parts.append(_IMAGE_OMISSION_TEXT)

When the newest user message is multimodal and has at least one text part, every image_url / image part is discarded and the resulting message contains no indication that an image was removed (_IMAGE_OMISSION_TEXT is only added when there is no text at all). With IMAGE_TOKEN_ESTIMATE = 1024, this branch is reachable on ordinary small-window models with one or two pasted images, so the user asks "what's wrong in this screenshot?" and the model answers from the prompt text alone. This is the classic silently-degraded-context hallucination path.

Suggest always appending _IMAGE_OMISSION_TEXT (or the omission notice) when any non-text part is dropped.

Related, same branch: when base_mandatory_tokens > input_budget and the newest user prompt is truncated, no CONTEXT_OMISSION_NOTICE is added to the system prompt, unlike the non-mandatory path. That asymmetry means the most aggressive form of pruning is the one the model is least likely to be told about.

---

2. Low — A hung tokenizer download pins _tokenizer_load_in_progress and permanently degrades estimates

notebook_intelligence/chat_history_budget.py — _load_tokenizer_encoding / warm_tokenizer_encoding

_tokenizer_load_in_progress is cleared at the end of _load_tokenizer_encoding, inside the normal flow (the except catches Exception). tiktoken.encoding_for_model() fetches the BPE file over HTTP on a cold cache with no client-side timeout, so a hung/blackholed proxy leaves the flag True forever: the retry budget is never consumed, the thread leaks, and every subsequent call falls back to the coarse bytes/3 estimator for the process lifetime. Using try/finally for the flag (and, ideally, an explicit timeout or a check for a locally cached encoding) would make the retry path actually reachable.

Related, non-blocking: because warm-up is asynchronous, the first requests after server start use the byte-based fallback, which over-estimates typical English/code by ~30%. Combined with the 0.8 reserve, the same conversation can prune differently before and after the tokenizer lands — non-deterministic behavior that will be hard to explain in a bug report. Worth logging once when the fallback is actually used for budgeting.

---

3. Low — Current-turn context is consumed oldest-first, starving the newest attachment

notebook_intelligence/chat_history_budget.py — _budget_chat_messages

    for message in current_context:
        ...
        truncated = _truncate_text_message(message, remaining_budget)

The loop walks the current-turn context in original (oldest-first) order and lets the first oversized entry truncate itself into whatever budget remains. In practice the last context message is the one the user just attached/@-mentioned, so a stale large attachment earlier in the same turn can consume the entire remaining budget and push the freshest context out. Iterating newest-first (and re-sorting into original order for emission) would match user expectation better.

---

4. Low — Over-budget request is still sent when the system prompt alone exceeds the budget

notebook_intelligence/chat_history_budget.py

If input_budget - system_tokens <= 0, _truncate_mandatory_message returns None and the untruncated latest_user is kept, so the assembled list is knowingly over budget and the provider will 400. The ask-mode caller turns that into the generic "Oops! There was a problem handling chat request" message, which gives the user no way to understand that their model's window is too small for the configured system prompt/rules. Consider trimming the system message (or surfacing a specific message) in that case rather than emitting a request that is guaranteed to fail.

---

5. Low — Verify no remaining references to the removed extension.tiktoken_encoding

notebook_intelligence/extension.py

import tiktoken and the module-level tiktoken_encoding global were removed. The diff only shows the two helpers that used them; please confirm nothing else in the package or the test suite imports/patches notebook_intelligence.extension.tiktoken_encoding (a stale patch(...) target would now raise AttributeError at test time, and a stale runtime reference would be a NameError on a rarely hit path).

---

6. Nit — Unrelated dependency change bundled without a changelog note

pyproject.toml: "mcp>=1.27.0" → "mcp>=1.27.0,<2". Reasonable on its own, but it is unrelated to chat-history budgeting and is not mentioned in the CHANGELOG entry added by this PR.

---

Test coverage notes

Coverage of the new module is good. Gaps I'd consider adding:

• A case asserting the emitted messages remain a valid provider sequence when current_context truncation empties the budget and complete turns are also selected (ordering: turns → context → latest user).
• A case where the mandatory message is multimodal with text (the finding in §1) — the existing test only covers the image-only placeholder path.
• test_context_budget_helpers_use_strict_shared_token_policy runs against whatever tokenizer state the process happens to be in (real tiktoken vs. UTF-8 fallback, depending on test ordering and the warm-up thread). The assertions happen to hold in both branches today, but pinning the encoding via the same fixture used in test_chat_history_budget.py would remove the ordering dependence.

Residual risks

• Agent tool loops (handle_chat_request_with_tools) remain unbudgeted, which is where context blowups are most likely (large tool outputs); this is called out as intentional, but it is the largest remaining exposure.
• The gpt-4o tokenizer is used as a cross-provider proxy; for non-OpenAI providers the 20% output reserve absorbs, but does not guarantee, the mismatch.

[herikwebb]

Closing this fork staging PR because the finalized change was promoted upstream in plmbr#412 and merged on 2026-08-31. The upstream version includes additional integration fixes and tests, while this fork PR is now conflicted after the base diverged.