(MOT-4325) feat(context-manager): named-part counting and an assemble breakdown

context-manager/src/core/estimate.rs

@@ -11,15 +11,6 @@ pub enum EstimatorKind {
     Heuristic,
 }
 
-impl EstimatorKind {
-    pub fn as_str(&self) -> &'static str {
-        match self {
-            EstimatorKind::Tokenizer => "tokenizer",
-            EstimatorKind::Heuristic => "heuristic",
-        }
-    }
-}
-
 pub trait Estimator: Send + Sync {
     fn kind(&self) -> EstimatorKind;
 
@@ -114,6 +105,21 @@ pub fn estimate_by_role(est: &dyn Estimator, messages: &[AgentMessage]) -> ByRol
     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 {
+    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
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -204,4 +210,31 @@ mod tests {
             estimate_messages(&est, &messages)
         );
     }
+
+    #[test]
+    fn by_role_from_sizes_partitions_every_role() {
+        let messages = vec![
+            msg(
+                json!({ "role": "user", "content": [{ "type": "text", "text": "q" }], "timestamp": 1 }),
+            ),
+            msg(
+                json!({ "role": "assistant", "content": [], "stop_reason": "end",
+                        "model": "m", "provider": "p", "timestamp": 2 }),
+            ),
+            msg(
+                json!({ "role": "function_result", "function_call_id": "c", "function_id": "f",
+                        "content": [], "timestamp": 3 }),
+            ),
+            msg(json!({ "role": "custom", "custom_type": "t", "content": [], "timestamp": 4 })),
+        ];
+        let est = HeuristicEstimator;
+        let sizes: Vec<u64> = messages.iter().map(|m| est.message(m)).collect();
+        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>()
+        );
+    }
 }

context-manager/src/functions/assemble.rs

@@ -14,7 +14,7 @@ use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
 use crate::core::budget::{default_reserved, preserve_recent_budget, usable};
-use crate::core::estimate::{estimator_for_model, Estimator};
+use crate::core::estimate::{by_role_from_sizes, estimator_for_model, Estimator};
 use crate::core::lease;
 use crate::core::prune::{emergency_reduce_with_sizes, prune_with_sizes, PruneParams};
 use crate::core::selection::select;
@@ -24,7 +24,9 @@ use crate::core::summary::{
 use crate::error::ContextError;
 use crate::functions::resolve_model;
 use crate::ports::{Deps, SummarizeRequest};
-use crate::types::{AgentFunction, AgentMessage, ModelInput, Role, ThinkingLevel};
+use crate::types::{
+    AgentFunction, AgentMessage, ByRoleTokens, EstimatorName, ModelInput, Role, ThinkingLevel,
+};
 
 #[derive(Debug, Default, Deserialize, JsonSchema)]
 pub struct AssembleOptions {
@@ -112,6 +114,22 @@ pub struct Applied {
     pub summarized_head_tokens: Option<u64>,
 }
 
+/// 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,
+}
+
 #[derive(Debug, Serialize, JsonSchema)]
 pub struct AssembleResponse {
     pub system_prompt: String,
@@ -127,6 +145,7 @@ pub struct AssembleResponse {
     pub effective_max_output_tokens: u64,
     pub model_resolved: ModelResolvedWire,
     pub applied: Applied,
+    pub breakdown: AssembleBreakdown,
 }
 
 /// Test-only re-export of [`count_context`] so sibling function tests
@@ -312,6 +331,8 @@ pub async fn handle(deps: &Deps, req: AssembleRequest) -> Result<AssembleRespons
         });
     }
 
+    let by_role = by_role_from_sizes(&working, &sizes);
+
     Ok(AssembleResponse {
         system_prompt,
         messages: working,
@@ -324,6 +345,12 @@ pub async fn handle(deps: &Deps, req: AssembleRequest) -> Result<AssembleRespons
             crate::core::budget::ModelResolved::Fallback => ModelResolvedWire::Fallback,
         },
         applied,
+        breakdown: AssembleBreakdown {
+            system_prompt_tokens: prompt_tokens,
+            tools_tokens: tool_tokens,
+            by_role: by_role.into(),
+            estimator: estimator.kind().into(),
+        },
     })
 }
 

context-manager/src/functions/count_tokens.rs

@@ -6,13 +6,15 @@
 //! falls back to the generic heuristic, reported in `estimator`), so
 //! cost-sensitive callers can run this with no `llm-router` installed.
 
+use std::collections::BTreeMap;
+
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
 use crate::core::estimate::{estimate_by_role, estimate_messages, estimator_for_model};
 use crate::error::ContextError;
 use crate::ports::Deps;
-use crate::types::{AgentFunction, AgentMessage, ModelInput};
+use crate::types::{AgentFunction, AgentMessage, ByRoleTokens, EstimatorName, ModelInput};
 
 #[derive(Debug, Deserialize, JsonSchema)]
 pub struct CountTokensRequest {
@@ -25,34 +27,30 @@ pub struct CountTokensRequest {
     /// single `agent_trigger` entry).
     #[serde(default)]
     pub tools: Option<Vec<AgentFunction>>,
+    /// Named auxiliary texts to count individually with the same
+    /// estimator (for example the segments a system prompt is built
+    /// from). Counted in `by_part` only; never added to `tokens`.
+    #[serde(default)]
+    pub parts: Option<BTreeMap<String, String>>,
     /// Tokenizer selection; falls back to a generic estimator.
     pub model: ModelInput,
 }
 
-/// Per-role token breakdown of the `messages` array.
-#[derive(Debug, Serialize, JsonSchema)]
-pub struct ByRoleTokens {
-    pub user: u64,
-    pub assistant: u64,
-    pub function_result: u64,
-    pub custom: u64,
-}
-
-/// Which estimator produced the count.
-#[derive(Debug, Serialize, JsonSchema)]
-#[serde(rename_all = "snake_case")]
-pub enum EstimatorName {
-    Tokenizer,
-    Heuristic,
-}
-
 #[derive(Debug, Serialize, JsonSchema)]
 pub struct CountTokensResponse {
     /// Total estimate: messages + system prompt + tools.
     pub tokens: u64,
     /// Breakdown of the message tokens by role (system prompt and
     /// tools are not part of any role bucket).
     pub by_role: Option<ByRoleTokens>,
+    /// The `tools` share of `tokens`; present when the request carried
+    /// `tools`.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tools_tokens: Option<u64>,
+    /// Per-part estimates for the request's `parts`, keyed by the
+    /// caller's names; present when the request carried `parts`.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub by_part: Option<BTreeMap<String, u64>>,
     pub estimator: EstimatorName,
 }
 
@@ -70,24 +68,29 @@ pub async fn handle(
     if let Some(system_prompt) = &req.system_prompt {
         tokens += estimator.text(system_prompt);
     }
-    for tool in req.tools.iter().flatten() {
-        tokens += estimator.function(tool);
-    }
+    let tools_tokens = req.tools.as_ref().map(|tools| {
+        tools
+            .iter()
+            .map(|tool| estimator.function(tool))
+            .sum::<u64>()
+    });
+    tokens += tools_tokens.unwrap_or(0);
+
+    let by_part = req.parts.as_ref().map(|parts| {
+        parts
+            .iter()
+            .map(|(name, text)| (name.clone(), estimator.text(text)))
+            .collect::<BTreeMap<String, u64>>()
+    });
 
     let by_role = estimate_by_role(estimator, &messages);
 
     Ok(CountTokensResponse {
         tokens,
-        by_role: Some(ByRoleTokens {
-            user: by_role.user,
-            assistant: by_role.assistant,
-            function_result: by_role.function_result,
-            custom: by_role.custom,
-        }),
-        estimator: match estimator.kind() {
-            crate::core::estimate::EstimatorKind::Tokenizer => EstimatorName::Tokenizer,
-            crate::core::estimate::EstimatorKind::Heuristic => EstimatorName::Heuristic,
-        },
+        by_role: Some(by_role.into()),
+        tools_tokens,
+        by_part,
+        estimator: estimator.kind().into(),
     })
 }
 

context-manager/src/types.rs

@@ -13,6 +13,8 @@ use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
+use crate::core::estimate::{ByRole, EstimatorKind};
+
 /// Message role discriminator.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
 #[serde(rename_all = "snake_case")]
@@ -287,6 +289,43 @@ pub enum ExecutionMode {
     Sequential,
 }
 
+/// Per-role token breakdown of the `messages` array.
+#[derive(Debug, Serialize, JsonSchema)]
+pub struct ByRoleTokens {
+    pub user: u64,
+    pub assistant: u64,
+    pub function_result: u64,
+    pub custom: u64,
+}
+
+impl From<ByRole> for ByRoleTokens {
+    fn from(by_role: ByRole) -> Self {
+        ByRoleTokens {
+            user: by_role.user,
+            assistant: by_role.assistant,
+            function_result: by_role.function_result,
+            custom: by_role.custom,
+        }
+    }
+}
+
+/// Which estimator produced the count.
+#[derive(Debug, Serialize, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum EstimatorName {
+    Tokenizer,
+    Heuristic,
+}
+
+impl From<EstimatorKind> for EstimatorName {
+    fn from(kind: EstimatorKind) -> Self {
+        match kind {
+            EstimatorKind::Tokenizer => EstimatorName::Tokenizer,
+            EstimatorKind::Heuristic => EstimatorName::Heuristic,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

context-manager/tests/features/assemble.feature

@@ -323,3 +323,16 @@ Feature: context::assemble — the model-ready context pipeline
       { "allow_compaction": false }
       """
     Then the call fails with code "context/overflow"
+
+  # Prevents: the breakdown drifting from the total. Budget UIs render
+  # these categories as segments of token_count, so they must reconcile.
+  Scenario: the response reports a per-category breakdown
+    Given a user message "what is the weather"
+    And an assistant message "sunny with a chance of tokens"
+    When I assemble the history with model "any-model" and system prompt "xxxxxxxxxxxxxxxx"
+    Then the call succeeds
+    And the response field "breakdown.system_prompt_tokens" is 4
+    And the response field "breakdown.tools_tokens" is 0
+    And the response field "breakdown.by_role.user" exceeds 0
+    And the response field "breakdown.by_role.assistant" exceeds 0
+    And the response field "breakdown.estimator" is "heuristic"

context-manager/tests/features/count_tokens.feature

@@ -76,3 +76,30 @@ Feature: context::count-tokens — estimate token usage for a message set
     Then the call succeeds
     And the response field "tokens" is 0
     And the response field "by_role.user" is 0
+
+  # Prevents: named parts leaking into the total — parts are a side
+  # breakdown for callers dissecting a prompt they also count whole
+  # (e.g. system prompt segments), so double-billing must be impossible.
+  Scenario: parts are counted individually and never added to the total
+    Given an empty history
+    When I count tokens with model "any-model" and parts:
+      """
+      { "identity": "xxxxxxxxxxxxxxxx", "guidance": "xxxxxxxx" }
+      """
+    Then the call succeeds
+    And the response field "tokens" is 0
+    And the response field "by_part.identity" is 4
+    And the response field "by_part.guidance" is 2
+
+  # Prevents: the tools share being indistinguishable inside the total —
+  # budget UIs render schemas as their own category.
+  Scenario: the tools share of the total is reported separately
+    Given an empty history
+    When I count tokens with model "any-model" and tools:
+      """
+      [{ "name": "agent_trigger", "description": "Invoke any allowed iii function.",
+         "parameters": { "type": "object", "properties": { "function": { "type": "string" } } } }]
+      """
+    Then the call succeeds
+    And the response field "tokens" exceeds 0
+    And the response field "tools_tokens" equals the response field "tokens"