diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index ca32e0a061..7b2adaf155 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -215,12 +215,14 @@ impl App { scrollback_reflow, deferred_history_cell, )?; + self.insert_completed_token_activity_output_if_ready(tui); } AppEvent::ConsolidateProposedPlan(source) => { if !self.terminal_resize_reflow_enabled() { if !self.transcript_reflow.history_cell_refresh_requested() { self.transcript_reflow.clear(); } + self.insert_completed_token_activity_output(tui); return Ok(AppRunControl::Continue); } let end = self.transcript_cells.len(); @@ -255,6 +257,7 @@ impl App { self.maybe_finish_stream_reflow(tui)?; } + self.insert_completed_token_activity_output_if_ready(tui); } AppEvent::ApplyThreadRollback { num_turns } => { if self.apply_non_pending_thread_rollback(num_turns) { @@ -728,14 +731,16 @@ impl App { if let Err(err) = &result { tracing::warn!("account/tokenUsage/read failed during TUI refresh: {err}"); } - if let Some(cell) = self + if self .chat_widget .finish_token_activity_refresh(request_id, result) { // Commit synchronously so an already queued /clear cannot overtake this card. // Do not route through ChatWidget::add_to_history: /tokens may complete during // active work, and flushing an in-progress tool cell would corrupt its lifecycle. - self.insert_history_cell(tui, Box::new(cell)); + // If an answer stream is active, keep the settled card transient until its + // provisional transcript cells have been consolidated. + self.insert_completed_token_activity_output_if_ready(tui); } } AppEvent::ConnectorsLoaded { result, is_final } => { diff --git a/codex-rs/tui/src/app/history_ui.rs b/codex-rs/tui/src/app/history_ui.rs index 2f6c0e31f6..f07ee48f15 100644 --- a/codex-rs/tui/src/app/history_ui.rs +++ b/codex-rs/tui/src/app/history_ui.rs @@ -30,6 +30,24 @@ impl App { } } + pub(super) fn insert_completed_token_activity_output_if_ready(&mut self, tui: &mut tui::Tui) { + if self.chat_widget.token_activity_history_insertion_blocked() + || self.transcript_cells.last().is_some_and(|cell| { + cell.as_any().is::() + || cell.as_any().is::() + }) + { + return; + } + self.insert_completed_token_activity_output(tui); + } + + pub(super) fn insert_completed_token_activity_output(&mut self, tui: &mut tui::Tui) { + if let Some(cell) = self.chat_widget.take_completed_token_activity_output() { + self.insert_history_cell(tui, Box::new(cell)); + } + } + pub(super) fn open_url_in_browser(&mut self, url: String) { if let Err(err) = webbrowser::open(&url) { self.chat_widget diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index b1e1074261..111972d3ac 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -551,6 +551,7 @@ pub(crate) struct ChatWidget { refreshing_status_outputs: Vec<(u64, StatusHistoryHandle)>, next_status_refresh_request_id: u64, refreshing_token_activity_output: Option, + completed_token_activity_output: Option, next_token_activity_request_id: u64, plan_type: Option, codex_rate_limit_reached_type: Option, diff --git a/codex-rs/tui/src/chatwidget/constructor.rs b/codex-rs/tui/src/chatwidget/constructor.rs index aac6f1235e..46b0eb098c 100644 --- a/codex-rs/tui/src/chatwidget/constructor.rs +++ b/codex-rs/tui/src/chatwidget/constructor.rs @@ -125,6 +125,7 @@ impl ChatWidget { refreshing_status_outputs: Vec::new(), next_status_refresh_request_id: 0, refreshing_token_activity_output: None, + completed_token_activity_output: None, next_token_activity_request_id: 0, plan_type: initial_plan_type, codex_rate_limit_reached_type: None, diff --git a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs index 6f4fa12a86..2a5515648e 100644 --- a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs +++ b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs @@ -1199,11 +1199,10 @@ async fn clearing_pending_token_activity_refreshes_discards_late_result() { chat.clear_pending_token_activity_refreshes(); assert!( - chat.finish_token_activity_refresh( + !chat.finish_token_activity_refresh( request_id, Err("stale token activity result".to_string()), ) - .is_none() ); } @@ -1247,8 +1246,15 @@ async fn completed_token_activity_refresh_returns_one_history_cell() { }; assert!(chat.pending_token_activity_output().is_some()); + assert!( + chat.finish_token_activity_refresh( + request_id, + Err("token activity unavailable".to_string()), + ) + ); + assert!(chat.pending_token_activity_output().is_some()); let cell = chat - .finish_token_activity_refresh(request_id, Err("token activity unavailable".to_string())) + .take_completed_token_activity_output() .expect("completed token activity cell"); assert!(chat.pending_token_activity_output().is_none()); @@ -1258,6 +1264,36 @@ async fn completed_token_activity_refresh_returns_one_history_cell() { ); } +#[tokio::test] +async fn completed_token_activity_refresh_waits_for_active_stream() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + set_chatgpt_auth(&mut chat); + + chat.dispatch_command(SlashCommand::Tokens); + let request_id = match rx.try_recv() { + Ok(AppEvent::RefreshTokenActivity { request_id }) => request_id, + other => panic!("expected token activity refresh request, got {other:?}"), + }; + chat.on_agent_message_delta("partial response".to_string()); + assert!(chat.token_activity_history_insertion_blocked()); + + assert!( + chat.finish_token_activity_refresh( + request_id, + Err("token activity unavailable".to_string()), + ) + ); + assert_eq!( + chat.pending_token_activity_output() + .map(|cell| lines_to_single_string(&cell.display_lines(u16::MAX))), + Some("/tokens daily\n\n Token activity\n Token activity unavailable\n".to_string()), + ); + + chat.stream_controller = None; + assert!(!chat.token_activity_history_insertion_blocked()); + assert!(chat.take_completed_token_activity_output().is_some()); +} + #[tokio::test] async fn repeated_token_activity_refreshes_keep_only_latest_card() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; @@ -1280,20 +1316,14 @@ async fn repeated_token_activity_refreshes_keep_only_latest_card() { .map(|cell| lines_to_single_string(&cell.display_lines(u16::MAX))), Some("/tokens weekly\n\n Token activity\n Loading...\n".to_string()), ); - assert!( - chat.finish_token_activity_refresh( - first_request_id, - Err("stale token activity result".to_string()), - ) - .is_none() - ); - assert!( - chat.finish_token_activity_refresh( - second_request_id, - Err("token activity unavailable".to_string()), - ) - .is_some() - ); + assert!(!chat.finish_token_activity_refresh( + first_request_id, + Err("stale token activity result".to_string()), + )); + assert!(chat.finish_token_activity_refresh( + second_request_id, + Err("token activity unavailable".to_string()), + )); } #[tokio::test] diff --git a/codex-rs/tui/src/chatwidget/tokens.rs b/codex-rs/tui/src/chatwidget/tokens.rs index 9a304813dc..50e426d528 100644 --- a/codex-rs/tui/src/chatwidget/tokens.rs +++ b/codex-rs/tui/src/chatwidget/tokens.rs @@ -642,6 +642,7 @@ impl ChatWidget { self.next_token_activity_request_id = self.next_token_activity_request_id.wrapping_add(/*rhs*/ 1); let (cell, handle) = new_token_activity_output(view); + self.completed_token_activity_output = None; self.refreshing_token_activity_output = Some(PendingTokenActivityOutput { request_id, cell, @@ -657,26 +658,46 @@ impl ChatWidget { self.refreshing_token_activity_output .as_ref() .map(|output| &output.cell as &dyn HistoryCell) + .or_else(|| { + self.completed_token_activity_output + .as_ref() + .map(|cell| cell as &dyn HistoryCell) + }) } pub(crate) fn finish_token_activity_refresh( &mut self, request_id: u64, result: Result, - ) -> Option { - let output = self.refreshing_token_activity_output.take()?; + ) -> bool { + let Some(output) = self.refreshing_token_activity_output.take() else { + return false; + }; if output.request_id != request_id { self.refreshing_token_activity_output = Some(output); - return None; + return false; } output.handle.finish(result); + self.completed_token_activity_output = Some(output.cell); self.bump_active_cell_revision(); self.request_redraw(); - Some(output.cell) + true + } + + pub(crate) fn token_activity_history_insertion_blocked(&self) -> bool { + self.stream_controller.is_some() || self.plan_stream_controller.is_some() + } + + pub(crate) fn take_completed_token_activity_output(&mut self) -> Option { + let output = self.completed_token_activity_output.take()?; + self.bump_active_cell_revision(); + Some(output) } pub(crate) fn clear_pending_token_activity_refreshes(&mut self) { - if self.refreshing_token_activity_output.take().is_some() { + let cleared_refresh = self.refreshing_token_activity_output.take().is_some(); + let cleared_completed = self.completed_token_activity_output.take().is_some(); + if cleared_refresh || cleared_completed { self.bump_active_cell_revision(); self.request_redraw(); }