(MOT-4325) feat(context-manager): named-part counting and an assemble breakdown - #683
Merged
Conversation
…re in count-tokens context::count-tokens gains an optional parts map of named texts, counted individually with the request's estimator and returned as by_part without joining the total, plus a tools_tokens field that reports the tools share of the existing total. Callers dissecting a context window (system prompt segments vs tool schemas) no longer re-implement the estimator or make one call per part.
…m assemble context::assemble now returns breakdown alongside the totals it already maintains: system prompt tokens (summary section included), tools tokens, per-role message tokens, and the estimator that produced them. Callers rendering a context meter get the categories from the one call that already happens per generation instead of a second count round trip.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Contributor
skill-check — worker0 verified, 54 skipped (no docs/).
Four for four. Nicely done. |
📝 WalkthroughWalkthroughThe context manager now reports token-estimation details for assembled responses and token-count requests. It adds per-role, tool, system-prompt, and named-part totals, estimator names, wire schemas, and acceptance coverage. ChangesToken estimation telemetry
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@context-manager/src/functions/assemble.rs`:
- Around line 116-129: Expose request overhead in AssembleBreakdown by adding
request_overhead_tokens: u64 and populating it from request_overhead_tokens in
the assemble result construction; update
context-manager/src/functions/assemble.rs lines 116-129 and 360-368 accordingly.
Regenerate the required schema field in
context-manager/tests/golden/schemas/context.assemble.json lines 932-965, and
update context-manager/tests/features/assemble.feature lines 327-338 to send
nonzero overhead and assert the returned breakdown value.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: d52f6fb8-385d-4591-86e0-dda2d5dc14eb
📒 Files selected for processing (8)
context-manager/src/functions/assemble.rscontext-manager/src/functions/count_tokens.rscontext-manager/tests/features/assemble.featurecontext-manager/tests/features/count_tokens.featurecontext-manager/tests/golden/schemas/context.assemble.jsoncontext-manager/tests/golden/schemas/context.count-tokens.jsoncontext-manager/tests/steps/call_steps.rscontext-manager/tests/steps/common_steps.rs
Comment on lines
+116
to
+129
| /// Where the returned `token_count` sits, by category — the same sums the | ||
| /// pipeline already maintains, exposed so callers can render a context | ||
| /// breakdown without re-counting: `token_count` equals the by_role sum plus | ||
| /// `system_prompt_tokens`, `tools_tokens`, and the request overhead. | ||
| #[derive(Debug, Serialize, JsonSchema)] | ||
| pub struct AssembleBreakdown { | ||
| /// Estimated tokens of the returned system prompt (any compaction | ||
| /// summary section included). | ||
| pub system_prompt_tokens: u64, | ||
| /// Estimated tokens of the invocation schemas. | ||
| pub tools_tokens: u64, | ||
| /// The returned messages' tokens by role. | ||
| pub by_role: ByRoleTokens, | ||
| pub estimator: EstimatorName, |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Expose request overhead in AssembleBreakdown.
token_count includes request_overhead_tokens, but breakdown does not return it. A caller cannot reconcile the breakdown with token_count when it sends a nonzero overhead.
context-manager/src/functions/assemble.rs#L116-L129: Addrequest_overhead_tokens: u64toAssembleBreakdown.context-manager/src/functions/assemble.rs#L360-L368: Set the field fromrequest_overhead_tokens.context-manager/tests/golden/schemas/context.assemble.json#L932-L965: Regenerate the schema with the required field.context-manager/tests/features/assemble.feature#L327-L338: Send a nonzero overhead and assert the returned breakdown value.
Proposed fix
pub struct AssembleBreakdown {
pub system_prompt_tokens: u64,
pub tools_tokens: u64,
+ pub request_overhead_tokens: u64,
pub by_role: ByRoleTokens,
pub estimator: EstimatorName,
}
breakdown: AssembleBreakdown {
system_prompt_tokens: prompt_tokens,
tools_tokens: tool_tokens,
+ request_overhead_tokens,
by_role,
estimator: match estimator.kind() {📍 Affects 3 files
context-manager/src/functions/assemble.rs#L116-L129(this comment)context-manager/src/functions/assemble.rs#L360-L368context-manager/tests/golden/schemas/context.assemble.json#L932-L965context-manager/tests/features/assemble.feature#L327-L338
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@context-manager/src/functions/assemble.rs` around lines 116 - 129, Expose
request overhead in AssembleBreakdown by adding request_overhead_tokens: u64 and
populating it from request_overhead_tokens in the assemble result construction;
update context-manager/src/functions/assemble.rs lines 116-129 and 360-368
accordingly. Regenerate the required schema field in
context-manager/tests/golden/schemas/context.assemble.json lines 932-965, and
update context-manager/tests/features/assemble.feature lines 327-338 to send
nonzero overhead and assert the returned breakdown value.
… and estimator name Review cleanups, wire-identical (goldens unchanged). ByRoleTokens and EstimatorName move to types.rs with From impls from the core ByRole and EstimatorKind, deleting two hand-written conversions and the dead as_str. The by-role accumulation the breakdown added lives once in core::estimate::by_role_from_sizes over the memoized sizes, shared by assemble.
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
context-manager/src/core/estimate.rs (1)
213-239: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert exact role totals using independent sizes.
Lines 231-237 derive
sizesfromest.message, so a regression that re-estimates messages can still pass. The aggregate assertions also do not verify each role mapping. Use distinct fixed sizes and assert everyByRolefield.Proposed test strengthening
- let est = HeuristicEstimator; - let sizes: Vec<u64> = messages.iter().map(|m| est.message(m)).collect(); + let sizes = vec![11_u64, 22, 33, 44]; let by_role = by_role_from_sizes(&messages, &sizes); - assert!(by_role.user > 0 && by_role.assistant > 0); - assert!(by_role.function_result > 0 && by_role.custom > 0); - assert_eq!( - by_role.user + by_role.assistant + by_role.function_result + by_role.custom, - sizes.iter().sum::<u64>() - ); + assert_eq!(by_role.user, 11); + assert_eq!(by_role.assistant, 22); + assert_eq!(by_role.function_result, 33); + assert_eq!(by_role.custom, 44);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@context-manager/src/core/estimate.rs` around lines 213 - 239, Strengthen by_role_from_sizes_partitions_every_role by replacing estimator-derived sizes with distinct fixed values, then assert each ByRole field exactly matches the corresponding role’s size and that the total remains correct. Keep the existing four-role message fixture and call by_role_from_sizes directly so the test detects incorrect role mapping or message re-estimation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@context-manager/src/core/estimate.rs`:
- Around line 108-122: The by_role_from_sizes function currently zips messages
and sizes, silently truncating mismatched inputs. Enforce that messages.len()
equals sizes.len() before calculating the breakdown, using the function’s
existing return contract to assert or return an error, and preserve the per-role
accumulation for valid inputs.
---
Nitpick comments:
In `@context-manager/src/core/estimate.rs`:
- Around line 213-239: Strengthen by_role_from_sizes_partitions_every_role by
replacing estimator-derived sizes with distinct fixed values, then assert each
ByRole field exactly matches the corresponding role’s size and that the total
remains correct. Keep the existing four-role message fixture and call
by_role_from_sizes directly so the test detects incorrect role mapping or
message re-estimation.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: e54244ae-38eb-48e3-a729-f1774eacf5c5
📒 Files selected for processing (10)
context-manager/src/core/estimate.rscontext-manager/src/functions/assemble.rscontext-manager/src/functions/count_tokens.rscontext-manager/src/types.rscontext-manager/tests/features/assemble.featurecontext-manager/tests/features/count_tokens.featurecontext-manager/tests/golden/schemas/context.assemble.jsoncontext-manager/tests/golden/schemas/context.count-tokens.jsoncontext-manager/tests/steps/call_steps.rscontext-manager/tests/steps/common_steps.rs
🚧 Files skipped from review as they are similar to previous changes (9)
- context-manager/tests/features/count_tokens.feature
- context-manager/tests/steps/common_steps.rs
- context-manager/tests/features/assemble.feature
- context-manager/src/types.rs
- context-manager/tests/golden/schemas/context.assemble.json
- context-manager/src/functions/assemble.rs
- context-manager/tests/steps/call_steps.rs
- context-manager/tests/golden/schemas/context.count-tokens.json
- context-manager/src/functions/count_tokens.rs
Comment on lines
+108
to
+122
| /// Per-role breakdown from already-computed message sizes (`assemble`'s | ||
| /// memoized `sizes`), so callers with a size memo don't re-estimate. | ||
| pub fn by_role_from_sizes(messages: &[AgentMessage], sizes: &[u64]) -> ByRole { | ||
| let mut by_role = ByRole::default(); | ||
| for (message, size) in messages.iter().zip(sizes) { | ||
| match message.role() { | ||
| Role::User => by_role.user += size, | ||
| Role::Assistant => by_role.assistant += size, | ||
| Role::FunctionResult => by_role.function_result += size, | ||
| Role::Custom => by_role.custom += size, | ||
| } | ||
| } | ||
| by_role | ||
| } | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject mismatched message and size slices before zipping.
Line 112 truncates to the shorter slice. If sizes drifts from messages, the breakdown silently omits entries and can no longer reconcile with token_count. Enforce equal lengths, or return an error when callers control these inputs.
Proposed invariant check
pub fn by_role_from_sizes(messages: &[AgentMessage], sizes: &[u64]) -> ByRole {
+ assert_eq!(
+ messages.len(),
+ sizes.len(),
+ "messages and sizes must have equal lengths"
+ );
let mut by_role = ByRole::default();
- for (message, size) in messages.iter().zip(sizes) {
+ for (message, size) in messages.iter().zip(sizes.iter()) {📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| /// Per-role breakdown from already-computed message sizes (`assemble`'s | |
| /// memoized `sizes`), so callers with a size memo don't re-estimate. | |
| pub fn by_role_from_sizes(messages: &[AgentMessage], sizes: &[u64]) -> ByRole { | |
| let mut by_role = ByRole::default(); | |
| for (message, size) in messages.iter().zip(sizes) { | |
| match message.role() { | |
| Role::User => by_role.user += size, | |
| Role::Assistant => by_role.assistant += size, | |
| Role::FunctionResult => by_role.function_result += size, | |
| Role::Custom => by_role.custom += size, | |
| } | |
| } | |
| by_role | |
| } | |
| /// Per-role breakdown from already-computed message sizes (`assemble`'s | |
| /// memoized `sizes`), so callers with a size memo don't re-estimate. | |
| pub fn by_role_from_sizes(messages: &[AgentMessage], sizes: &[u64]) -> ByRole { | |
| assert_eq!( | |
| messages.len(), | |
| sizes.len(), | |
| "messages and sizes must have equal lengths" | |
| ); | |
| let mut by_role = ByRole::default(); | |
| for (message, size) in messages.iter().zip(sizes.iter()) { | |
| match message.role() { | |
| Role::User => by_role.user += size, | |
| Role::Assistant => by_role.assistant += size, | |
| Role::FunctionResult => by_role.function_result += size, | |
| Role::Custom => by_role.custom += size, | |
| } | |
| } | |
| by_role | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@context-manager/src/core/estimate.rs` around lines 108 - 122, The
by_role_from_sizes function currently zips messages and sizes, silently
truncating mismatched inputs. Enforce that messages.len() equals sizes.len()
before calculating the breakdown, using the function’s existing return contract
to assert or return an error, and preserve the per-role accumulation for valid
inputs.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Two additive surfaces on context-manager so callers can render a context breakdown without re-implementing the estimator or making extra round trips.
context::count-tokensgains an optionalpartsmap of named texts, counted individually and returned asby_partwithout joining the total, plus atools_tokensfield reporting the tools share of the existing total.context::assemblegains abreakdownresponse object:system_prompt_tokens(summary section included),tools_tokens, per-role message tokens, and the estimator that produced them. Computed from the size memo the pipeline already maintains, so the numbers reconcile withtoken_countby construction.Why
The harness assembles a context every generation and previously discarded every per-category number. The console context meter (follow-up PRs on MOT-4326 and MOT-4327) needs those categories from the one call that already happens per step.
Testing
Part of MOT-4324. No behavior change for existing callers; both responses are additive.
Summary by CodeRabbit
New Features
Tests