diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index b65b151579..c967abf9c8 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -1280,6 +1280,12 @@ pub(crate) enum ReplayKind { ThreadSnapshot, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum TurnCompletionStateCleanup { + ClearPendingMarkers, + PreservePendingMarkers, +} + #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum ThreadItemRenderSource { Live, @@ -5469,7 +5475,7 @@ impl ChatWidget { ))); return; } - let starts_standalone_user_shell_turn = !self.bottom_pane.is_task_running() + let starts_standalone_user_shell_turn = !self.agent_turn_running && !self.pending_turn_start_after_submit && !self.should_hold_user_input_for_standalone_shell_turn(); if self.submit_op(AppCommand::run_user_shell_command(cmd.to_string())) @@ -5790,6 +5796,7 @@ impl ChatWidget { }, }, Some(replay_kind), + TurnCompletionStateCleanup::PreservePendingMarkers, ); } } @@ -6204,7 +6211,7 @@ impl ChatWidget { } ServerNotification::TurnStarted(notification) => { let turn_id = notification.turn.id; - if replay_kind.is_none() { + if !matches!(replay_kind, Some(ReplayKind::ResumeInitialMessages)) { self.note_turn_started(&turn_id); } self.last_turn_id = Some(turn_id); @@ -6214,7 +6221,15 @@ impl ChatWidget { } } ServerNotification::TurnCompleted(notification) => { - self.handle_turn_completed_notification(notification, replay_kind); + let state_cleanup = match replay_kind { + Some(ReplayKind::ResumeInitialMessages) => { + TurnCompletionStateCleanup::PreservePendingMarkers + } + Some(ReplayKind::ThreadSnapshot) | None => { + TurnCompletionStateCleanup::ClearPendingMarkers + } + }; + self.handle_turn_completed_notification(notification, replay_kind, state_cleanup); } ServerNotification::ItemStarted(notification) => { self.handle_item_started_notification(notification, replay_kind.is_some()); @@ -6464,9 +6479,15 @@ impl ChatWidget { &mut self, notification: TurnCompletedNotification, replay_kind: Option, + state_cleanup: TurnCompletionStateCleanup, ) { - self.pending_turn_start_after_submit = false; - self.clear_standalone_user_shell_turn(Some(¬ification.turn.id)); + if matches!( + state_cleanup, + TurnCompletionStateCleanup::ClearPendingMarkers + ) { + self.pending_turn_start_after_submit = false; + self.clear_standalone_user_shell_turn(Some(¬ification.turn.id)); + } match notification.turn.status { TurnStatus::Completed => { self.last_non_retry_error = None; @@ -6783,7 +6804,7 @@ impl ChatWidget { EventMsg::TurnStarted(event) => { let turn_id = event.turn_id; let model_context_window = event.model_context_window; - if !from_replay { + if !is_resume_initial_replay { self.note_turn_started(&turn_id); } self.last_turn_id = Some(turn_id); @@ -6797,8 +6818,10 @@ impl ChatWidget { last_agent_message, .. }) => { - self.pending_turn_start_after_submit = false; - self.clear_standalone_user_shell_turn(Some(&turn_id)); + if !is_resume_initial_replay { + self.pending_turn_start_after_submit = false; + self.clear_standalone_user_shell_turn(Some(&turn_id)); + } self.on_task_complete(last_agent_message, from_replay); } EventMsg::TokenCount(ev) => { @@ -6835,8 +6858,10 @@ impl ChatWidget { EventMsg::McpStartupUpdate(ev) => self.on_mcp_startup_update(ev), EventMsg::McpStartupComplete(ev) => self.on_mcp_startup_complete(ev), EventMsg::TurnAborted(ev) => { - self.pending_turn_start_after_submit = false; - self.clear_standalone_user_shell_turn(ev.turn_id.as_deref()); + if !is_resume_initial_replay { + self.pending_turn_start_after_submit = false; + self.clear_standalone_user_shell_turn(ev.turn_id.as_deref()); + } match ev.reason { TurnAbortReason::Interrupted => { self.on_interrupted_turn(ev.reason); diff --git a/codex-rs/tui/src/chatwidget/tests/exec_flow.rs b/codex-rs/tui/src/chatwidget/tests/exec_flow.rs index cf6ca84d10..de970ef92f 100644 --- a/codex-rs/tui/src/chatwidget/tests/exec_flow.rs +++ b/codex-rs/tui/src/chatwidget/tests/exec_flow.rs @@ -1198,6 +1198,28 @@ async fn bang_after_user_turn_submit_before_turn_started_does_not_mark_standalon assert!(!chat.pending_standalone_user_shell_command); } +#[tokio::test] +async fn bang_during_mcp_startup_marks_standalone_shell_command() { + let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.thread_id = Some(ThreadId::new()); + chat.mcp_startup_status = Some(HashMap::from([( + "server".to_string(), + McpStartupStatus::Starting, + )])); + chat.update_task_running_state(); + + chat.bottom_pane + .set_composer_text("!sleep 10".to_string(), Vec::new(), Vec::new()); + chat.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + match op_rx.try_recv() { + Ok(Op::RunUserShellCommand { command }) => assert_eq!(command, "sleep 10"), + other => panic!("expected RunUserShellCommand op, got {other:?}"), + } + assert!(chat.pending_turn_start_after_submit); + assert!(chat.pending_standalone_user_shell_command); +} + #[tokio::test] async fn completion_without_seen_shell_start_clears_pending_marker() { let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/tests/history_replay.rs b/codex-rs/tui/src/chatwidget/tests/history_replay.rs index d79865d060..9b7da74f8b 100644 --- a/codex-rs/tui/src/chatwidget/tests/history_replay.rs +++ b/codex-rs/tui/src/chatwidget/tests/history_replay.rs @@ -825,6 +825,79 @@ async fn replayed_in_progress_turn_consumes_pending_standalone_shell_marker() { assert!(chat.bottom_pane.is_task_running()); } +#[tokio::test] +async fn replayed_completed_turn_preserves_pending_shell_submit_markers() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.pending_turn_start_after_submit = true; + chat.pending_standalone_user_shell_command = true; + + chat.replay_thread_turns( + vec![AppServerTurn { + id: "old-turn".to_string(), + items: Vec::new(), + status: AppServerTurnStatus::Completed, + error: None, + started_at: None, + completed_at: Some(0), + duration_ms: None, + }], + ReplayKind::ThreadSnapshot, + ); + + assert!(chat.pending_turn_start_after_submit); + assert!(chat.pending_standalone_user_shell_command); + assert!(chat.standalone_user_shell_turn_id.is_none()); +} + +#[tokio::test] +async fn replayed_buffered_shell_turn_completion_clears_pending_shell_marker() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.pending_turn_start_after_submit = true; + chat.pending_standalone_user_shell_command = true; + + chat.handle_server_notification( + ServerNotification::TurnStarted(TurnStartedNotification { + thread_id: "thread-1".to_string(), + turn: AppServerTurn { + id: "shell-turn".to_string(), + items: Vec::new(), + status: AppServerTurnStatus::InProgress, + error: None, + started_at: Some(0), + completed_at: None, + duration_ms: None, + }, + }), + Some(ReplayKind::ThreadSnapshot), + ); + + assert!(!chat.pending_turn_start_after_submit); + assert!(!chat.pending_standalone_user_shell_command); + assert_eq!( + chat.standalone_user_shell_turn_id.as_deref(), + Some("shell-turn") + ); + + chat.handle_server_notification( + ServerNotification::TurnCompleted(TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: AppServerTurn { + id: "shell-turn".to_string(), + items: Vec::new(), + status: AppServerTurnStatus::Completed, + error: None, + started_at: Some(0), + completed_at: Some(1), + duration_ms: Some(1), + }, + }), + Some(ReplayKind::ThreadSnapshot), + ); + + assert!(chat.standalone_user_shell_turn_id.is_none()); + assert!(!chat.pending_standalone_user_shell_command); +} + #[tokio::test] async fn replayed_stream_error_does_not_set_retry_status_or_status_indicator() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await;