core: count post-usage tokens after last model item

codex-rs/core/src/context_manager/history.rs

@@ -236,15 +236,24 @@ impl ContextManager {
             })
     }
 
-    fn get_trailing_codex_generated_items_tokens(&self) -> i64 {
-        let mut total = 0i64;
-        for item in self.items.iter().rev() {
-            if !is_codex_generated_item(item) {
-                break;
-            }
-            total = total.saturating_add(estimate_item_token_count(item));
-        }
-        total
+    // These are local items added after the most recent model-emitted item.
+    // They are not reflected in `last_token_usage.total_tokens`.
+    // If no model item has been emitted yet, treat the delta as empty.
+    fn items_after_last_model_generated_item(&self) -> &[ResponseItem] {
+        let start = self
+            .items
+            .iter()
+            .rposition(is_model_generated_item)
+            .map_or(self.items.len(), |index| index.saturating_add(1));
+        &self.items[start..]
+    }
+
+    fn get_items_after_last_model_generated_tokens(&self) -> i64 {
+        self.items_after_last_model_generated_item()
+            .iter()
+            .fold(0i64, |acc, item| {
+                acc.saturating_add(estimate_item_token_count(item))
+            })
     }
 
     /// When true, the server already accounted for past reasoning tokens and
@@ -255,13 +264,14 @@ impl ContextManager {
             .as_ref()
             .map(|info| info.last_token_usage.total_tokens)
             .unwrap_or(0);
-        let trailing_codex_generated_tokens = self.get_trailing_codex_generated_items_tokens();
+        let items_after_last_model_generated_tokens =
+            self.get_items_after_last_model_generated_tokens();
         if server_reasoning_included {
-            last_tokens.saturating_add(trailing_codex_generated_tokens)
+            last_tokens.saturating_add(items_after_last_model_generated_tokens)
         } else {
             last_tokens
                 .saturating_add(self.get_non_last_reasoning_items_tokens())
-                .saturating_add(trailing_codex_generated_tokens)
+                .saturating_add(items_after_last_model_generated_tokens)
         }
     }
 
@@ -367,6 +377,22 @@ fn estimate_item_token_count(item: &ResponseItem) -> i64 {
     }
 }
 
+fn is_model_generated_item(item: &ResponseItem) -> bool {
+    match item {
+        ResponseItem::Message { role, .. } => role == "assistant",
+        ResponseItem::Reasoning { .. }
+        | ResponseItem::FunctionCall { .. }
+        | ResponseItem::WebSearchCall { .. }
+        | ResponseItem::CustomToolCall { .. }
+        | ResponseItem::LocalShellCall { .. }
+        | ResponseItem::Compaction { .. } => true,
+        ResponseItem::FunctionCallOutput { .. }
+        | ResponseItem::CustomToolCallOutput { .. }
+        | ResponseItem::GhostSnapshot { .. }
+        | ResponseItem::Other => false,
+    }
+}
+
 pub(crate) fn is_codex_generated_item(item: &ResponseItem) -> bool {
     matches!(
         item,

codex-rs/core/src/context_manager/history_tests.rs

@@ -189,48 +189,32 @@ fn non_last_reasoning_tokens_ignore_entries_after_last_user() {
 }
 
 #[test]
-fn trailing_codex_generated_tokens_stop_at_first_non_generated_item() {
-    let earlier_output = function_call_output("call-earlier", "earlier output");
-    let trailing_function_output = function_call_output("call-tail-1", "tail function output");
-    let trailing_custom_output = custom_tool_call_output("call-tail-2", "tail custom output");
+fn items_after_last_model_generated_tokens_include_user_and_tool_output() {
     let history = create_history_with_items(vec![
-        earlier_output,
-        user_msg("boundary item"),
-        trailing_function_output.clone(),
-        trailing_custom_output.clone(),
+        assistant_msg("already counted by API"),
+        user_msg("new user message"),
+        custom_tool_call_output("call-tail", "new tool output"),
     ]);
-    let expected_tokens = estimate_item_token_count(&trailing_function_output)
-        .saturating_add(estimate_item_token_count(&trailing_custom_output));
+    let expected_tokens = estimate_item_token_count(&user_msg("new user message")).saturating_add(
+        estimate_item_token_count(&custom_tool_call_output("call-tail", "new tool output")),
+    );
 
     assert_eq!(
-        history.get_trailing_codex_generated_items_tokens(),
+        history.get_items_after_last_model_generated_tokens(),
         expected_tokens
     );
 }
 
 #[test]
-fn trailing_codex_generated_tokens_exclude_function_call_tail() {
-    let history = create_history_with_items(vec![ResponseItem::FunctionCall {
-        id: None,
-        name: "not-generated".to_string(),
-        arguments: "{}".to_string(),
-        call_id: "call-tail".to_string(),
-    }]);
+fn items_after_last_model_generated_tokens_are_zero_without_model_generated_items() {
+    let history = create_history_with_items(vec![user_msg("no model output yet")]);
 
-    assert_eq!(history.get_trailing_codex_generated_items_tokens(), 0);
+    assert_eq!(history.get_items_after_last_model_generated_tokens(), 0);
 }
 
 #[test]
-fn total_token_usage_includes_only_trailing_codex_generated_items() {
-    let non_trailing_output = function_call_output("call-before-message", "not trailing");
-    let trailing_assistant = assistant_msg("assistant boundary");
-    let trailing_output = custom_tool_call_output("tool-tail", "trailing output");
-    let mut history = create_history_with_items(vec![
-        non_trailing_output,
-        user_msg("boundary"),
-        trailing_assistant,
-        trailing_output.clone(),
-    ]);
+fn total_token_usage_includes_all_items_after_last_model_generated_item() {
+    let mut history = create_history_with_items(vec![assistant_msg("already counted by API")]);
     history.update_token_info(
         &TokenUsage {
             total_tokens: 100,
@@ -238,10 +222,17 @@ fn total_token_usage_includes_only_trailing_codex_generated_items() {
         },
         None,
     );
+    let added_user = user_msg("new user message");
+    let added_tool_output = custom_tool_call_output("tool-tail", "new tool output");
+    history.record_items(
+        [&added_user, &added_tool_output],
+        TruncationPolicy::Tokens(10_000),
+    );
 
     assert_eq!(
         history.get_total_token_usage(true),
-        100 + estimate_item_token_count(&trailing_output)
+        100 + estimate_item_token_count(&added_user)
+            + estimate_item_token_count(&added_tool_output)
     );
 }