[Akshat977]

Pull Request

Issue

Fixes #1319

Description

ChatContext._as_list_token_budget decides which history items fit the model's context window. It estimates each item at len(rendered) // 4. That guess can be well off, because tokenizers differ by model and chat templates add tokens of their own.

Backends already record real token usage on each reply (generation.usage), and they record what was sent on _call (context and action). This PR uses that record to correct the estimate:

1. Find the latest reply with usage.
2. Check the usage belongs to this context:
   • The reply's generation.model matches the context's model. With a ModelIdentifier, any of its platform names counts as a match.
   • The item just before the reply is the reply's _call.action.
   • The items before that are the same objects as _call.context (what the backend sent).
3. Ratio = measured total tokens ÷ estimated tokens for those items (sent context + action + reply).
4. Multiply every item's estimate by the ratio, including older items and items added after the reply. The newest-to-oldest walk and the 0.75 headroom are unchanged.

If any check fails, the ratio is 1.0 and behaviour is exactly as before.

Other changes:

• _last_usage_tokens in compactor.py is split into _last_usage_thunk and _usage_total so both can be reused. Its behaviour is unchanged.
• Items are rendered only when their estimate is needed, so a long history isn't rendered past what the budget can hold (same as before).

Decisions worth a look

• Model only, not backend. The issue discussion asked for estimates to be tied to the backend and model that produced them. A ChatContext knows its model but not its backend, so the check compares the model name only. If the context has no model (only token_context_length_limit), the correction is skipped.
• Template overhead goes into the ratio. The measured prompt count includes the chat template and any system prompt, which aren't separate items. Spreading that over the items makes the estimate a little conservative, so it drops slightly more rather than overflowing.
• One ratio for the whole history. A per-item estimate (comparing the usage of two consecutive replies) could be a follow-up. I kept this PR to the approach suggested in the issue.

Testing

• Tests added to the respective file if code was changed
• New code has 100% coverage if code was added
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

New tests in test/stdlib/test_base_context.py:

• Measured usage higher than the estimate drops more items.
• Measured usage lower than the estimate keeps more items. This includes an item older than what was sent.
• Items added after the reply are scaled too.
• Falls back to the plain estimate when there's no usage, zero usage, a different model, a context without a model, no receipt, a receipt mismatch, a receipt longer than the history, or an action mismatch.
• _model_matches with string and ModelIdentifier models.

The new trimming tests fail on main and pass with this change. The existing context, compactor, stdlib, core and backend unit tests give the same results before and after (run with --disable-default-mellea-plugins). I haven't checked this against a live backend.

Attribution

• AI coding assistants used

---

Adding a new component, requirement, sampling strategy, or tool?

If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.

• Component
• Requirement
• Sampling Strategy
• Tool

[Akshat977]

@jakelorocco this is ready for review whenever you have time. Checks are passing, and I'm happy to change anything.

[markstur]

The idea of correcting the len // 4 estimate with measured usage is sound and the fallback handling for mismatched receipts is careful. Comments are correctness issues that would affect ordinary use, since a first turn and any tool-using conversation hits them. The main concerns are about guarding the extreme against extreme scaling which would likely be a problem.

AI-generated review comments and I cross-checked them. Feel free to call out mistakes or ask for more clarity.

[markstur]

Unguarded next() (chat.py:285) — idx = next(...) with no default crashes view_for_generation with StopIteration if the items argument ever diverges from self.as_list(), instead of degrading to scale 1.0 like every other mismatch.

The unguarded next() can't fail today, but _usage_scale only works because of an assumption its signature doesn't state.

Why it's safe now
Only caller: _as_list_token_budget is the only caller. It builds items = self.as_list() and passes it in.
Where the thunk comes from: thunk comes from _last_usage_thunk(self), which iterates ctx.as_list() on the same context. The thunk is therefore always an element of items.
Contexts that can't be linearized: If as_list() couldn't linearize a context, it would fail earlier, before this line.

Why it's still fragile
_usage_scale takes items as a parameter, so its contract says nothing about items and self having to agree. Called directly with an items list that lacks the thunk, and it raised StopIteration. A bare StopIteration is also a poor failure: a future caller inside a generator or coroutine would get a confusing error instead of the fallback.

Everything else in the method returns 1.0 when the data doesn't line up. This line is the only one that raises, which is inconsistent with that pattern.

[markstur]

Redundant O(n) work (chat.py:274) — _usage_scale triggers a second full as_list() walk plus a third identity scan per call, and re-checks usage is None that _last_usage_thunk already guarantees. view_for_generation runs 2–3× per generation on hot paths.

I haven't looked into how to refactor this. Author should at least consider it.

[markstur]

the scale blows up when the measured segment is small.

On the first turn the receipt is empty, so estimated covers only the action and the reply, a handful of tokens. measured, though, includes the fixed template and system-prompt overhead, which can be tens of tokens. The ratio is then very large. For example, a 3-token question and a 2-token reply with a measured total of 40 give a scale of about 8. Every history item is then costed at that multiple until a larger reply replaces the measurement, so the context is trimmed far more than needed. This is the normal case on turn one, not a corner case.

Please guard against small denominators, for example:

skip the correction (return 1.0) when estimated is below a minimum, such as a fraction of the budget or a fixed token count, and/or
clamp the result to a sane range such as 0.5–2.0.
Please also add a test for a first-turn reply (empty receipt, small action and reply, measured usage dominated by overhead). The current tests all use 100-character items, so they don't reach this case.

[markstur]

tool-call and reasoning thunks are costed at the 1-token floor.

A ModelOutputThunk with an empty value, or one carrying only tool calls, renders as ''. Its estimate is therefore 1 token, whatever it really costs. This causes two problems:

If such a thunk is inside the measured segment, estimated is too low and the scale is inflated, so history is over-trimmed.
If it is added after the reply, it costs about scale × 1 tokens when it may cost hundreds. The context is then over-filled, which is the failure this budget is meant to prevent.
Please either estimate these thunks from their tool calls and reasoning text, or skip the correction when the measured segment contains one. Please add a test with a tool-call thunk both inside the measured segment and after the reply.

[markstur]

Non-blocking, but needs a decision: overhead is counted twice.

measured is the full request cost, including the system prompt, tool schemas and template tokens, none of which are in items. Dividing by the item-only estimate folds that overhead into every item's cost. The 0.75 headroom factor already reserves capacity for the same overhead, so history can be trimmed when the real request would have fit.

This is the same cause as comment 1 and the clamp or minimum there reduces it. Please either subtract a known overhead or switch to prompt_tokens. If you keep the current approach, update the docstrings on _usage_scale and _as_list_token_budget to say that overhead is included in the scale. Then the 0.75 headroom can be revisited deliberately rather than stacked by accident.

[Akshat977]

Thanks @markstur for the detailed review. All five points are addressed in the latest commit:

1. Unguarded next(): _usage_scale now finds the measured reply with one backwards pass over the items it's given, so there's no next(). If items has no measured reply, it returns 1.0 like every other mismatch (test added).
2. Redundant work: that same pass replaces the second as_list() walk and the identity scan, and the repeated usage is None check is gone. One side effect: if a reply is re-added later as a validation target, the later copy isn't preceded by its action, so the correction falls back to the plain estimate.
3. Small denominators: the correction is skipped when the measured items are estimated at under 100 tokens, and the ratio is clamped to 0.5-2.0. There's a first-turn test with your example (3-token question, 2-token reply, 40 measured); without the guard it drops the whole history.
4. Tool-call and reasoning replies: estimates now include a reply's tool calls (name and JSON arguments) and reasoning text, so a tool-call-only reply is no longer costed at 1 token. Reasoning is counted even when the backend doesn't send it back, which errs toward trimming. The test has a tool-call reply both inside the measured items and after the reply.
5. Overhead counted twice: I kept total usage, since prompt_tokens includes the same overhead. The docstrings on _usage_scale and _as_list_token_budget now say the scale includes a share of fixed overhead that overlaps with the 0.75 headroom; the minimum and the clamp keep it small. Happy to revisit the headroom in a follow-up if you'd prefer.

The new tests fail on the previous commit and pass now.

[ajbozarth]

Some feedback from Claude: verified end-to-end against Ollama (granite4.1:3b, 3 turns via both chat and instruct). The receipt and model checks pass on every turn and the scale is steady at ~0.82, so the correction is active on the default path. I also confirmed Ollama's prompt_eval_count still counts KV-cached prefix tokens, so cached turns don't skew the ratio. One nit inline.

[Akshat977]

Thanks @ajbozarth for checking it end-to-end against Ollama! Went with inlining: _last_usage_thunk is folded back into _last_usage_tokens, and _usage_total stays as the shared helper. Pushed.