diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index da2efb63c..6738d7672 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -114,6 +114,11 @@ impl BottomPane { self.status.as_ref() } + #[cfg(test)] + pub(crate) fn context_window_percent(&self) -> Option { + self.context_window_percent + } + fn active_view(&self) -> Option<&dyn BottomPaneView> { self.view_stack.last().map(std::convert::AsRef::as_ref) } diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index e832b5890..3a79330b7 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -290,6 +290,8 @@ pub(crate) struct ChatWidget { pending_notification: Option, // Simple review mode flag; used to adjust layout and banners. is_review_mode: bool, + // Snapshot of token usage to restore after review mode exits. + pre_review_token_info: Option>, // Whether to add a final message separator after the last message needs_final_message_separator: bool, @@ -489,16 +491,39 @@ impl ChatWidget { } pub(crate) fn set_token_info(&mut self, info: Option) { - if let Some(info) = info { - let context_window = info - .model_context_window - .or(self.config.model_context_window); - let percent = context_window.map(|window| { + match info { + Some(info) => self.apply_token_info(info), + None => { + self.bottom_pane.set_context_window_percent(None); + self.token_info = None; + } + } + } + + fn apply_token_info(&mut self, info: TokenUsageInfo) { + let percent = self.context_remaining_percent(&info); + self.bottom_pane.set_context_window_percent(percent); + self.token_info = Some(info); + } + + fn context_remaining_percent(&self, info: &TokenUsageInfo) -> Option { + info.model_context_window + .or(self.config.model_context_window) + .map(|window| { info.last_token_usage .percent_of_context_window_remaining(window) - }); - self.bottom_pane.set_context_window_percent(percent); - self.token_info = Some(info); + }) + } + + fn restore_pre_review_token_info(&mut self) { + if let Some(saved) = self.pre_review_token_info.take() { + match saved { + Some(info) => self.apply_token_info(info), + None => { + self.bottom_pane.set_context_window_percent(None); + self.token_info = None; + } + } } } @@ -1150,6 +1175,7 @@ impl ChatWidget { suppress_session_configured_redraw: false, pending_notification: None, is_review_mode: false, + pre_review_token_info: None, needs_final_message_separator: false, last_rendered_width: std::cell::Cell::new(None), feedback, @@ -1223,6 +1249,7 @@ impl ChatWidget { suppress_session_configured_redraw: true, pending_notification: None, is_review_mode: false, + pre_review_token_info: None, needs_final_message_separator: false, last_rendered_width: std::cell::Cell::new(None), feedback, @@ -1693,6 +1720,9 @@ impl ChatWidget { fn on_entered_review_mode(&mut self, review: ReviewRequest) { // Enter review mode and emit a concise banner + if self.pre_review_token_info.is_none() { + self.pre_review_token_info = Some(self.token_info.clone()); + } self.is_review_mode = true; let banner = format!(">> Code review started: {} <<", review.user_facing_hint); self.add_to_history(history_cell::new_review_status_line(banner)); @@ -1733,6 +1763,7 @@ impl ChatWidget { } self.is_review_mode = false; + self.restore_pre_review_token_info(); // Append a finishing banner at the end of this turn. self.add_to_history(history_cell::new_review_status_line( "<< Code review finished >>".to_string(), diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index a6ba54647..a4cef2241 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -38,6 +38,9 @@ use codex_core::protocol::ReviewRequest; use codex_core::protocol::StreamErrorEvent; use codex_core::protocol::TaskCompleteEvent; use codex_core::protocol::TaskStartedEvent; +use codex_core::protocol::TokenCountEvent; +use codex_core::protocol::TokenUsage; +use codex_core::protocol::TokenUsageInfo; use codex_core::protocol::UndoCompletedEvent; use codex_core::protocol::UndoStartedEvent; use codex_core::protocol::ViewImageToolCallEvent; @@ -215,6 +218,81 @@ fn exited_review_mode_emits_results_and_finishes() { assert!(!chat.is_review_mode); } +/// Exiting review restores the pre-review context window indicator. +#[test] +fn review_restores_context_window_indicator() { + let (mut chat, mut rx, _ops) = make_chatwidget_manual(); + + let context_window = 13_000; + let pre_review_tokens = 12_700; // ~30% remaining after subtracting baseline. + let review_tokens = 12_030; // ~97% remaining after subtracting baseline. + + chat.handle_codex_event(Event { + id: "token-before".into(), + msg: EventMsg::TokenCount(TokenCountEvent { + info: Some(make_token_info(pre_review_tokens, context_window)), + rate_limits: None, + }), + }); + assert_eq!(chat.bottom_pane.context_window_percent(), Some(30)); + + chat.handle_codex_event(Event { + id: "review-start".into(), + msg: EventMsg::EnteredReviewMode(ReviewRequest { + prompt: "Review the latest changes".to_string(), + user_facing_hint: "feature branch".to_string(), + append_to_original_thread: true, + }), + }); + + chat.handle_codex_event(Event { + id: "token-review".into(), + msg: EventMsg::TokenCount(TokenCountEvent { + info: Some(make_token_info(review_tokens, context_window)), + rate_limits: None, + }), + }); + assert_eq!(chat.bottom_pane.context_window_percent(), Some(97)); + + chat.handle_codex_event(Event { + id: "review-end".into(), + msg: EventMsg::ExitedReviewMode(ExitedReviewModeEvent { + review_output: None, + }), + }); + let _ = drain_insert_history(&mut rx); + + assert_eq!(chat.bottom_pane.context_window_percent(), Some(30)); + assert!(!chat.is_review_mode); +} + +/// Receiving a TokenCount event without usage clears the context indicator. +#[test] +fn token_count_none_resets_context_indicator() { + let (mut chat, _rx, _ops) = make_chatwidget_manual(); + + let context_window = 13_000; + let pre_compact_tokens = 12_700; + + chat.handle_codex_event(Event { + id: "token-before".into(), + msg: EventMsg::TokenCount(TokenCountEvent { + info: Some(make_token_info(pre_compact_tokens, context_window)), + rate_limits: None, + }), + }); + assert_eq!(chat.bottom_pane.context_window_percent(), Some(30)); + + chat.handle_codex_event(Event { + id: "token-cleared".into(), + msg: EventMsg::TokenCount(TokenCountEvent { + info: None, + rate_limits: None, + }), + }); + assert_eq!(chat.bottom_pane.context_window_percent(), None); +} + #[cfg_attr( target_os = "macos", ignore = "system configuration APIs are blocked under macOS seatbelt" @@ -292,6 +370,7 @@ fn make_chatwidget_manual() -> ( suppress_session_configured_redraw: false, pending_notification: None, is_review_mode: false, + pre_review_token_info: None, needs_final_message_separator: false, last_rendered_width: std::cell::Cell::new(None), feedback: codex_feedback::CodexFeedback::new(), @@ -338,6 +417,21 @@ fn lines_to_single_string(lines: &[ratatui::text::Line<'static>]) -> String { s } +fn make_token_info(total_tokens: i64, context_window: i64) -> TokenUsageInfo { + fn usage(total_tokens: i64) -> TokenUsage { + TokenUsage { + total_tokens, + ..TokenUsage::default() + } + } + + TokenUsageInfo { + total_token_usage: usage(total_tokens), + last_token_usage: usage(total_tokens), + model_context_window: Some(context_window), + } +} + #[test] fn rate_limit_warnings_emit_thresholds() { let mut state = RateLimitWarningState::default();