Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions codex-rs/tui/src/bottom_pane/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,11 @@ impl BottomPane {
self.status.as_ref()
}

#[cfg(test)]
pub(crate) fn context_window_percent(&self) -> Option<i64> {
self.context_window_percent
}

fn active_view(&self) -> Option<&dyn BottomPaneView> {
self.view_stack.last().map(std::convert::AsRef::as_ref)
}
Expand Down
47 changes: 39 additions & 8 deletions codex-rs/tui/src/chatwidget.rs
Original file line number Diff line number Diff line change
Expand Up @@ -267,6 +267,8 @@ pub(crate) struct ChatWidget {
pending_notification: Option<Notification>,
// 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<Option<TokenUsageInfo>>,
// List of ghost commits corresponding to each turn.
ghost_snapshots: Vec<GhostCommit>,
ghost_snapshots_disabled: bool,
Expand Down Expand Up @@ -460,16 +462,39 @@ impl ChatWidget {
}

pub(crate) fn set_token_info(&mut self, info: Option<TokenUsageInfo>) {
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<i64> {
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;
}
}
}
}

Expand Down Expand Up @@ -989,6 +1014,7 @@ impl ChatWidget {
suppress_session_configured_redraw: false,
pending_notification: None,
is_review_mode: false,
pre_review_token_info: None,
ghost_snapshots: Vec::new(),
ghost_snapshots_disabled: true,
needs_final_message_separator: false,
Expand Down Expand Up @@ -1057,6 +1083,7 @@ impl ChatWidget {
suppress_session_configured_redraw: true,
pending_notification: None,
is_review_mode: false,
pre_review_token_info: None,
ghost_snapshots: Vec::new(),
ghost_snapshots_disabled: true,
needs_final_message_separator: false,
Expand Down Expand Up @@ -1532,6 +1559,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));
Expand Down Expand Up @@ -1572,6 +1602,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(),
Expand Down
93 changes: 93 additions & 0 deletions codex-rs/tui/src/chatwidget/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,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::ViewImageToolCallEvent;
use codex_protocol::ConversationId;
use codex_protocol::plan_tool::PlanItemArg;
Expand Down Expand Up @@ -220,6 +223,80 @@ 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(),
}),
});

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"
Expand Down Expand Up @@ -294,6 +371,7 @@ fn make_chatwidget_manual() -> (
suppress_session_configured_redraw: false,
pending_notification: None,
is_review_mode: false,
pre_review_token_info: None,
ghost_snapshots: Vec::new(),
ghost_snapshots_disabled: false,
needs_final_message_separator: false,
Expand Down Expand Up @@ -342,6 +420,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();
Expand Down