diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index 8dc612ecbb..8d340f0cef 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -215,6 +215,7 @@ impl App { scrollback_reflow, deferred_history_cell, )?; + self.chat_widget.note_stream_consolidation_completed(); self.insert_completed_token_activity_output_after_stream_shutdown(tui); } AppEvent::ConsolidateProposedPlan(source) => { @@ -222,6 +223,7 @@ impl App { if !self.transcript_reflow.history_cell_refresh_requested() { self.transcript_reflow.clear(); } + self.chat_widget.note_stream_consolidation_completed(); self.insert_completed_token_activity_output_after_stream_shutdown(tui); return Ok(AppRunControl::Continue); } @@ -257,6 +259,7 @@ impl App { self.maybe_finish_stream_reflow(tui)?; } + self.chat_widget.note_stream_consolidation_completed(); self.insert_completed_token_activity_output_after_stream_shutdown(tui); } AppEvent::ApplyThreadRollback { num_turns } => { diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 7db4e94644..d58ed95f34 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -564,6 +564,7 @@ pub(crate) struct ChatWidget { stream_controller: Option, // Stream lifecycle controller for proposed plan output. plan_stream_controller: Option, + pending_stream_consolidations: usize, /// Holds the platform clipboard lease so copied text remains available while supported. clipboard_lease: Option, copy_last_response_binding: Vec, diff --git a/codex-rs/tui/src/chatwidget/constructor.rs b/codex-rs/tui/src/chatwidget/constructor.rs index 46b0eb098c..d5ff119996 100644 --- a/codex-rs/tui/src/chatwidget/constructor.rs +++ b/codex-rs/tui/src/chatwidget/constructor.rs @@ -136,6 +136,7 @@ impl ChatWidget { adaptive_chunking: AdaptiveChunkingPolicy::default(), stream_controller: None, plan_stream_controller: None, + pending_stream_consolidations: 0, clipboard_lease: None, copy_last_response_binding, running_commands: HashMap::new(), diff --git a/codex-rs/tui/src/chatwidget/streaming.rs b/codex-rs/tui/src/chatwidget/streaming.rs index 3715de2a82..633d9fa6aa 100644 --- a/codex-rs/tui/src/chatwidget/streaming.rs +++ b/codex-rs/tui/src/chatwidget/streaming.rs @@ -39,6 +39,7 @@ impl ChatWidget { // that can re-render from source on resize. if let Some(source) = source { let source = parse_assistant_markdown(&source).visible_markdown; + self.note_stream_consolidation_queued(); self.app_event_tx.send(AppEvent::ConsolidateAgentMessage { source, cwd: self.config.cwd.to_path_buf(), @@ -174,12 +175,14 @@ impl ChatWidget { // TODO: Replace streamed output with the final plan item text if plan streaming is // removed or if we need to reconcile mismatches between streamed and final content. if let Some(source) = consolidated_plan_source { + self.note_stream_consolidation_queued(); self.app_event_tx .send(AppEvent::ConsolidateProposedPlan(source)); } } else if !plan_text.is_empty() { self.add_to_history(history_cell::new_proposed_plan(plan_text, &self.config.cwd)); } else if let Some(source) = consolidated_plan_source { + self.note_stream_consolidation_queued(); self.app_event_tx .send(AppEvent::ConsolidateProposedPlan(source)); } diff --git a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs index 577b65f1a3..af069e5613 100644 --- a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs +++ b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs @@ -1298,6 +1298,33 @@ async fn completed_token_activity_refresh_waits_for_active_stream() { assert!(chat.take_completed_token_activity_output().is_some()); } +#[tokio::test] +async fn completed_token_activity_refresh_waits_for_queued_stream_consolidation() { + 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()); + chat.finalize_completed_assistant_message(/*message*/ None); + assert!(chat.pending_stream_consolidations > 0); + + assert!( + chat.finish_token_activity_refresh( + request_id, + Err("token activity unavailable".to_string()), + ) + ); + assert!(chat.token_activity_history_insertion_blocked()); + + chat.note_stream_consolidation_completed(); + + assert!(!chat.token_activity_history_insertion_blocked()); +} + #[tokio::test] async fn completed_token_activity_refresh_waits_for_active_history_cell() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/tokens.rs b/codex-rs/tui/src/chatwidget/tokens.rs index ca0412ac38..92fc06bb75 100644 --- a/codex-rs/tui/src/chatwidget/tokens.rs +++ b/codex-rs/tui/src/chatwidget/tokens.rs @@ -687,9 +687,20 @@ impl ChatWidget { pub(crate) fn token_activity_history_insertion_blocked(&self) -> bool { self.stream_controller.is_some() || self.plan_stream_controller.is_some() + || self.pending_stream_consolidations > 0 || self.transcript.active_cell.is_some() } + pub(crate) fn note_stream_consolidation_queued(&mut self) { + self.pending_stream_consolidations = + self.pending_stream_consolidations.saturating_add(/*rhs*/ 1); + } + + pub(crate) fn note_stream_consolidation_completed(&mut self) { + self.pending_stream_consolidations = + self.pending_stream_consolidations.saturating_sub(/*rhs*/ 1); + } + pub(crate) fn take_completed_token_activity_output(&mut self) -> Option { let output = self.completed_token_activity_output.take()?; self.bump_active_cell_revision(); diff --git a/codex-rs/tui/src/chatwidget/turn_runtime.rs b/codex-rs/tui/src/chatwidget/turn_runtime.rs index e07b19c4ee..dd72daac86 100644 --- a/codex-rs/tui/src/chatwidget/turn_runtime.rs +++ b/codex-rs/tui/src/chatwidget/turn_runtime.rs @@ -126,6 +126,7 @@ impl ChatWidget { self.add_boxed_history(cell); } if let Some(source) = source { + self.note_stream_consolidation_queued(); self.app_event_tx .send(AppEvent::ConsolidateProposedPlan(source)); }