Calculate remaining context based on last token usage (#1940)

We should only take last request size (in tokens) into account

codex-rs/tui/src/bottom_pane/chat_composer.rs

@@ -45,7 +45,8 @@ pub enum InputResult {
 }
 
 struct TokenUsageInfo {
-    token_usage: TokenUsage,
+    total_token_usage: TokenUsage,
+    last_token_usage: TokenUsage,
     model_context_window: Option<u64>,
 }
 
@@ -129,11 +130,13 @@ impl ChatComposer {
     /// context when the composer is empty.
     pub(crate) fn set_token_usage(
         &mut self,
-        token_usage: TokenUsage,
+        total_token_usage: TokenUsage,
+        last_token_usage: TokenUsage,
         model_context_window: Option<u64>,
     ) {
         self.token_usage_info = Some(TokenUsageInfo {
-            token_usage,
+            total_token_usage,
+            last_token_usage,
             model_context_window,
         });
     }
@@ -694,16 +697,18 @@ impl WidgetRef for &ChatComposer {
 
                 // Append token/context usage info to the footer hints when available.
                 if let Some(token_usage_info) = &self.token_usage_info {
-                    let token_usage = &token_usage_info.token_usage;
+                    let token_usage = &token_usage_info.total_token_usage;
                     hint.push(Span::from("   "));
                     hint.push(
                         Span::from(format!("{} tokens used", token_usage.total_tokens))
                             .style(Style::default().add_modifier(Modifier::DIM)),
                     );
+                    let last_token_usage = &token_usage_info.last_token_usage;
                     if let Some(context_window) = token_usage_info.model_context_window {
                         let percent_remaining: u8 = if context_window > 0 {
                             let percent = 100.0
-                                - (token_usage.total_tokens as f32 / context_window as f32 * 100.0);
+                                - (last_token_usage.total_tokens as f32 / context_window as f32
+                                    * 100.0);
                             percent.clamp(0.0, 100.0) as u8
                         } else {
                             100

codex-rs/tui/src/bottom_pane/mod.rs

@@ -290,11 +290,12 @@ impl BottomPane<'_> {
     /// is forwarded directly to the underlying `ChatComposer`.
     pub(crate) fn set_token_usage(
         &mut self,
-        token_usage: TokenUsage,
+        total_token_usage: TokenUsage,
+        last_token_usage: TokenUsage,
         model_context_window: Option<u64>,
     ) {
         self.composer
-            .set_token_usage(token_usage, model_context_window);
+            .set_token_usage(total_token_usage, last_token_usage, model_context_window);
         self.request_redraw();
     }
 

codex-rs/tui/src/chatwidget.rs

@@ -66,7 +66,8 @@ pub(crate) struct ChatWidget<'a> {
     active_history_cell: Option<HistoryCell>,
     config: Config,
     initial_user_message: Option<UserMessage>,
-    token_usage: TokenUsage,
+    total_token_usage: TokenUsage,
+    last_token_usage: TokenUsage,
     reasoning_buffer: String,
     content_buffer: String,
     // Buffer for streaming assistant answer text; we do not surface partial
@@ -213,7 +214,8 @@ impl ChatWidget<'_> {
                 initial_prompt.unwrap_or_default(),
                 initial_images,
             ),
-            token_usage: TokenUsage::default(),
+            total_token_usage: TokenUsage::default(),
+            last_token_usage: TokenUsage::default(),
             reasoning_buffer: String::new(),
             content_buffer: String::new(),
             answer_buffer: String::new(),
@@ -365,9 +367,13 @@ impl ChatWidget<'_> {
                 self.request_redraw();
             }
             EventMsg::TokenCount(token_usage) => {
-                self.token_usage = add_token_usage(&self.token_usage, &token_usage);
-                self.bottom_pane
-                    .set_token_usage(self.token_usage.clone(), self.config.model_context_window);
+                self.total_token_usage = add_token_usage(&self.total_token_usage, &token_usage);
+                self.last_token_usage = token_usage;
+                self.bottom_pane.set_token_usage(
+                    self.total_token_usage.clone(),
+                    self.last_token_usage.clone(),
+                    self.config.model_context_window,
+                );
             }
             EventMsg::Error(ErrorEvent { message }) => {
                 self.add_to_history(HistoryCell::new_error_event(message.clone()));
@@ -552,7 +558,7 @@ impl ChatWidget<'_> {
     pub(crate) fn add_status_output(&mut self) {
         self.add_to_history(HistoryCell::new_status_output(
             &self.config,
-            &self.token_usage,
+            &self.total_token_usage,
         ));
     }
 
@@ -611,13 +617,16 @@ impl ChatWidget<'_> {
     }
 
     pub(crate) fn token_usage(&self) -> &TokenUsage {
-        &self.token_usage
+        &self.total_token_usage
     }
 
     pub(crate) fn clear_token_usage(&mut self) {
-        self.token_usage = TokenUsage::default();
-        self.bottom_pane
-            .set_token_usage(self.token_usage.clone(), self.config.model_context_window);
+        self.total_token_usage = TokenUsage::default();
+        self.bottom_pane.set_token_usage(
+            self.total_token_usage.clone(),
+            self.last_token_usage.clone(),
+            self.config.model_context_window,
+        );
     }
 
     pub fn cursor_pos(&self, area: Rect) -> Option<(u16, u16)> {