diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 8a66fca0e1..fde223771e 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -155,6 +155,7 @@ use crate::skills::SkillMetadata; use crate::skills::SkillsManager; use crate::skills::build_skill_injections; use crate::state::ActiveTurn; +use crate::state::RunningTask; use crate::state::SessionServices; use crate::state::SessionState; use crate::tasks::GhostSnapshotTask; @@ -1902,12 +1903,12 @@ impl Session { pub async fn inject_input(&self, input: Vec) -> Result<(), Vec> { let mut active = self.active_turn.lock().await; match active.as_mut() { - Some(at) => { + Some(at) if at.tasks.values().any(RunningTask::accepts_pending_input) => { let mut ts = at.turn_state.lock().await; ts.push_pending_input(input.into()); Ok(()) } - None => Err(input), + Some(_) | None => Err(input), } } @@ -1918,14 +1919,14 @@ impl Session { ) -> Result<(), Vec> { let mut active = self.active_turn.lock().await; match active.as_mut() { - Some(at) => { + Some(at) if at.tasks.values().any(RunningTask::accepts_pending_input) => { let mut ts = at.turn_state.lock().await; for item in input { ts.push_pending_input(item); } Ok(()) } - None => Err(input), + Some(_) | None => Err(input), } } @@ -2256,6 +2257,7 @@ mod handlers { use crate::review_prompts::resolve_review_request; use crate::tasks::CompactTask; use crate::tasks::RegularTask; + use crate::tasks::SpawnAbortPolicy; use crate::tasks::UndoTask; use crate::tasks::UserShellCommandTask; use codex_protocol::custom_prompts::CustomPrompt; @@ -2435,10 +2437,11 @@ mod handlers { previous_context: &mut Option>, ) { let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; - sess.spawn_task( + sess.spawn_task_with_policy( Arc::clone(&turn_context), Vec::new(), UserShellCommandTask::new(command), + SpawnAbortPolicy::KeepActiveTurn, ) .await; *previous_context = Some(turn_context); diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index ccc50d066b..74054480c7 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -50,6 +50,12 @@ pub(crate) struct RunningTask { pub(crate) _timer: Option, } +impl RunningTask { + pub(crate) fn accepts_pending_input(&self) -> bool { + self.task.accepts_pending_input() + } +} + impl ActiveTurn { pub(crate) fn add_task(&mut self, task: RunningTask) { let sub_id = task.turn_context.sub_id.clone(); diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index a2f30aa522..6643b77882 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -43,6 +43,12 @@ pub(crate) use user_shell::UserShellCommandTask; const GRACEFULL_INTERRUPTION_TIMEOUT_MS: u64 = 100; const TURN_ABORTED_INTERRUPTED_GUIDANCE: &str = "The user interrupted the previous turn. Do not continue or repeat work from that turn unless the user explicitly asks. If any tools/commands were aborted, they may have partially executed; verify current state before retrying."; +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum SpawnAbortPolicy { + AbortActiveTurn, + KeepActiveTurn, +} + /// Thin wrapper that exposes the parts of [`Session`] task runners need. #[derive(Clone)] pub(crate) struct SessionTaskContext { @@ -81,6 +87,11 @@ pub(crate) trait SessionTask: Send + Sync + 'static { /// surface it in telemetry and UI. fn kind(&self) -> TaskKind; + /// Whether this task can consume pending user input injected mid-turn. + fn accepts_pending_input(&self) -> bool { + true + } + /// Executes the task until completion or cancellation. /// /// Implementations typically stream protocol events using `session` and @@ -114,7 +125,20 @@ impl Session { input: Vec, task: T, ) { - self.abort_all_tasks(TurnAbortReason::Replaced).await; + self.spawn_task_with_policy(turn_context, input, task, SpawnAbortPolicy::AbortActiveTurn) + .await; + } + + pub async fn spawn_task_with_policy( + self: &Arc, + turn_context: Arc, + input: Vec, + task: T, + policy: SpawnAbortPolicy, + ) { + if policy == SpawnAbortPolicy::AbortActiveTurn { + self.abort_all_tasks(TurnAbortReason::Replaced).await; + } let task: Arc = Arc::new(task); let task_kind = task.kind(); @@ -198,9 +222,14 @@ impl Session { async fn register_new_active_task(&self, task: RunningTask) { let mut active = self.active_turn.lock().await; - let mut turn = ActiveTurn::default(); - turn.add_task(task); - *active = Some(turn); + match active.as_mut() { + Some(turn) => turn.add_task(task), + None => { + let mut turn = ActiveTurn::default(); + turn.add_task(task); + *active = Some(turn); + } + } } async fn take_all_running_tasks(&self) -> Vec { diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index 1f1f1eb59e..4cb76212cb 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -52,6 +52,10 @@ impl SessionTask for UserShellCommandTask { TaskKind::Regular } + fn accepts_pending_input(&self) -> bool { + false + } + async fn run( self: Arc, session: Arc, diff --git a/codex-rs/core/tests/suite/user_shell_cmd.rs b/codex-rs/core/tests/suite/user_shell_cmd.rs index 74cb1c55c8..239dbbc793 100644 --- a/codex-rs/core/tests/suite/user_shell_cmd.rs +++ b/codex-rs/core/tests/suite/user_shell_cmd.rs @@ -2,6 +2,8 @@ use anyhow::Context; use codex_core::NewThread; use codex_core::ThreadManager; use codex_core::features::Feature; +use codex_core::protocol::AskForApproval; +use codex_core::protocol::Event; use codex_core::protocol::EventMsg; use codex_core::protocol::ExecCommandEndEvent; use codex_core::protocol::ExecCommandSource; @@ -9,6 +11,8 @@ use codex_core::protocol::ExecOutputStream; use codex_core::protocol::Op; use codex_core::protocol::SandboxPolicy; use codex_core::protocol::TurnAbortReason; +use codex_protocol::config_types::ReasoningSummary; +use codex_protocol::user_input::UserInput; use core_test_support::assert_regex_match; use core_test_support::load_default_config_for_test; use core_test_support::responses; @@ -23,9 +27,13 @@ use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; +use pretty_assertions::assert_eq; use regex_lite::escape; use std::path::PathBuf; +use std::time::Duration; +use std::time::Instant; use tempfile::TempDir; +use tokio::time::timeout; #[tokio::test] async fn user_shell_cmd_ls_and_cat_in_temp_dir() { @@ -127,6 +135,81 @@ async fn user_shell_cmd_can_be_interrupted() { assert_eq!(ev.reason, TurnAbortReason::Interrupted); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn user_shell_cmd_does_not_abort_active_turn() -> anyhow::Result<()> { + let server = start_mock_server().await; + let sse_body = responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]); + let response = responses::sse_response(sse_body).set_delay(Duration::from_millis(200)); + let _mock = responses::mount_response_once(&server, response).await; + + let test = test_codex().build(&server).await?; + let codex = test.codex.clone(); + let session_model = test.session_configured.model.clone(); + let cwd = test.cwd_path().to_path_buf(); + + codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "keep running".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: cwd.clone(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::DangerFullAccess, + model: session_model, + effort: None, + summary: ReasoningSummary::Auto, + collaboration_mode: None, + personality: None, + }) + .await?; + + let regular_turn_id = loop { + let event = timeout(Duration::from_secs(1), codex.next_event()).await??; + if matches!(event.msg, EventMsg::TurnStarted(_)) { + break event.id; + } + }; + + codex + .submit(Op::RunUserShellCommand { + command: "true".to_string(), + }) + .await?; + + let deadline = Instant::now() + Duration::from_secs(2); + let mut replaced_abort_for_regular = false; + let mut regular_completed = false; + while Instant::now() < deadline && !regular_completed { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + break; + } + let event = timeout(remaining, codex.next_event()).await??; + match event { + Event { + id, + msg: EventMsg::TurnAborted(ev), + } if id == regular_turn_id && ev.reason == TurnAbortReason::Replaced => { + replaced_abort_for_regular = true; + } + Event { + id, + msg: EventMsg::TurnComplete(_), + } if id == regular_turn_id => { + regular_completed = true; + } + _ => {} + } + } + + assert_eq!(regular_completed, true); + assert_eq!(replaced_abort_for_regular, false); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn user_shell_command_history_is_persisted_and_shared_with_model() -> anyhow::Result<()> { let server = responses::start_mock_server().await; diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index aa420a3ea9..2e5d0da86a 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -18,7 +18,7 @@ //! The bottom pane exposes a single "task running" indicator that drives the spinner and interrupt //! hints. This module treats that indicator as derived UI-busy state: it is set while an agent turn //! is in progress and while MCP server startup is in progress. Those lifecycles are tracked -//! independently (`agent_turn_running` and `mcp_startup_status`) and synchronized via +//! independently (`running_turn_ids` and `mcp_startup_status`) and synchronized via //! `update_task_running_state`. use std::collections::HashMap; use std::collections::HashSet; @@ -444,11 +444,11 @@ pub(crate) struct ChatWidget { unified_exec_wait_streak: Option, task_complete_pending: bool, unified_exec_processes: Vec, - /// Tracks whether codex-core currently considers an agent turn to be in progress. + /// Tracks the set of active turn ids currently running. /// /// This is kept separate from `mcp_startup_status` so that MCP startup progress (or completion) /// can update the status header without accidentally clearing the spinner for an active turn. - agent_turn_running: bool, + running_turn_ids: HashSet, /// Tracks per-server MCP startup state while startup is in progress. /// /// The map is `Some(_)` from the first `McpStartupUpdate` until `McpStartupComplete`, and the @@ -671,8 +671,9 @@ impl ChatWidget { /// The bottom pane only has one running flag, but this module treats it as a derived state of /// both the agent turn lifecycle and MCP startup lifecycle. fn update_task_running_state(&mut self) { + let turn_running = !self.running_turn_ids.is_empty(); self.bottom_pane - .set_task_running(self.agent_turn_running || self.mcp_startup_status.is_some()); + .set_task_running(turn_running || self.mcp_startup_status.is_some()); } fn restore_reasoning_status_header(&mut self) { @@ -860,8 +861,13 @@ impl ChatWidget { // Raw reasoning uses the same flow as summarized reasoning - fn on_task_started(&mut self) { - self.agent_turn_running = true; + fn on_task_started(&mut self, turn_id: Option<&str>, from_replay: bool) { + if turn_id.is_none() && from_replay { + return; + } + if let Some(turn_id) = turn_id { + self.running_turn_ids.insert(turn_id.to_string()); + } self.saw_plan_update_this_turn = false; self.bottom_pane.clear_quit_shortcut_hint(); self.quit_shortcut_expires_at = None; @@ -875,18 +881,29 @@ impl ChatWidget { self.request_redraw(); } - fn on_task_complete(&mut self, last_agent_message: Option, from_replay: bool) { + fn on_task_complete( + &mut self, + turn_id: Option<&str>, + last_agent_message: Option, + from_replay: bool, + ) { // If a stream is currently active, finalize it. self.flush_answer_stream_with_separator(); self.flush_unified_exec_wait_streak(); - // Mark task stopped and request redraw now that all content is in history. - self.agent_turn_running = false; + if let Some(turn_id) = turn_id { + self.running_turn_ids.remove(turn_id); + } else { + self.running_turn_ids.clear(); + } self.update_task_running_state(); - self.running_commands.clear(); - self.suppressed_exec_calls.clear(); - self.last_unified_wait = None; - self.unified_exec_wait_streak = None; - self.clear_unified_exec_processes(); + if self.running_turn_ids.is_empty() { + // Only clear shared command state when all turns are finished. + self.running_commands.clear(); + self.suppressed_exec_calls.clear(); + self.last_unified_wait = None; + self.unified_exec_wait_streak = None; + self.clear_unified_exec_processes(); + } self.request_redraw(); if !from_replay && self.queued_user_messages.is_empty() { @@ -1096,7 +1113,7 @@ impl ChatWidget { // Ensure any spinner is replaced by a red ✗ and flushed into history. self.finalize_active_cell_as_failed(); // Reset running state and clear streaming buffers. - self.agent_turn_running = false; + self.running_turn_ids.clear(); self.update_task_running_state(); self.running_commands.clear(); self.suppressed_exec_calls.clear(); @@ -1991,7 +2008,7 @@ impl ChatWidget { unified_exec_wait_streak: None, task_complete_pending: false, unified_exec_processes: Vec::new(), - agent_turn_running: false, + running_turn_ids: HashSet::new(), mcp_startup_status: None, interrupts: InterruptManager::new(), reasoning_buffer: String::new(), @@ -2116,7 +2133,7 @@ impl ChatWidget { unified_exec_wait_streak: None, task_complete_pending: false, unified_exec_processes: Vec::new(), - agent_turn_running: false, + running_turn_ids: HashSet::new(), mcp_startup_status: None, interrupts: InterruptManager::new(), reasoning_buffer: String::new(), @@ -2242,7 +2259,7 @@ impl ChatWidget { unified_exec_wait_streak: None, task_complete_pending: false, unified_exec_processes: Vec::new(), - agent_turn_running: false, + running_turn_ids: HashSet::new(), mcp_startup_status: None, interrupts: InterruptManager::new(), reasoning_buffer: String::new(), @@ -2946,9 +2963,9 @@ impl ChatWidget { self.on_agent_reasoning_final(); } EventMsg::AgentReasoningSectionBreak(_) => self.on_reasoning_section_break(), - EventMsg::TurnStarted(_) => self.on_task_started(), + EventMsg::TurnStarted(_) => self.on_task_started(id.as_deref(), from_replay), EventMsg::TurnComplete(TurnCompleteEvent { last_agent_message }) => { - self.on_task_complete(last_agent_message, from_replay) + self.on_task_complete(id.as_deref(), last_agent_message, from_replay) } EventMsg::TokenCount(ev) => { self.set_token_info(ev.info); diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index fab175edcf..c87798718a 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -818,7 +818,7 @@ async fn make_chatwidget_manual( unified_exec_wait_streak: None, task_complete_pending: false, unified_exec_processes: Vec::new(), - agent_turn_running: false, + running_turn_ids: HashSet::new(), mcp_startup_status: None, interrupts: InterruptManager::new(), reasoning_buffer: String::new(), @@ -1265,7 +1265,7 @@ async fn plan_implementation_popup_skips_when_messages_queued() { chat.bottom_pane.set_task_running(true); chat.queue_user_message("Queued message".into()); - chat.on_task_complete(Some("Plan details".to_string()), false); + chat.on_task_complete(None, Some("Plan details".to_string()), false); let popup = render_bottom_popup(&chat, 80); assert!( @@ -1283,7 +1283,7 @@ async fn plan_implementation_popup_shows_on_plan_update_without_message() { .expect("expected plan collaboration mask"); chat.set_collaboration_mask(plan_mask); - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); chat.on_plan_update(UpdatePlanArgs { explanation: None, plan: vec![PlanItemArg { @@ -1291,7 +1291,7 @@ async fn plan_implementation_popup_shows_on_plan_update_without_message() { status: StepStatus::Pending, }], }); - chat.on_task_complete(None, false); + chat.on_task_complete(Some("turn-1"), None, false); let popup = render_bottom_popup(&chat, 80); assert!( @@ -1311,7 +1311,7 @@ async fn plan_implementation_popup_skips_when_rate_limit_prompt_pending() { .expect("expected plan collaboration mask"); chat.set_collaboration_mask(plan_mask); - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); chat.on_plan_update(UpdatePlanArgs { explanation: None, plan: vec![PlanItemArg { @@ -1320,7 +1320,7 @@ async fn plan_implementation_popup_skips_when_rate_limit_prompt_pending() { }], }); chat.on_rate_limit_snapshot(Some(snapshot(92.0))); - chat.on_task_complete(None, false); + chat.on_task_complete(Some("turn-1"), None, false); let popup = render_bottom_popup(&chat, 80); assert!( @@ -1678,7 +1678,7 @@ async fn streaming_final_answer_keeps_task_running_state() { let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(None).await; chat.thread_id = Some(ThreadId::new()); - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); chat.on_agent_message_delta("Final answer line\n".to_string()); chat.on_commit_tick(); @@ -1704,6 +1704,41 @@ async fn streaming_final_answer_keeps_task_running_state() { assert!(!chat.bottom_pane.quit_shortcut_hint_visible()); } +#[tokio::test] +async fn task_running_persists_when_user_shell_turn_completes_first() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(None).await; + + chat.handle_codex_event(Event { + id: "turn-regular".into(), + msg: EventMsg::TurnStarted(TurnStartedEvent { + model_context_window: None, + }), + }); + chat.handle_codex_event(Event { + id: "turn-shell".into(), + msg: EventMsg::TurnStarted(TurnStartedEvent { + model_context_window: None, + }), + }); + assert_eq!(chat.bottom_pane.is_task_running(), true); + + chat.handle_codex_event(Event { + id: "turn-shell".into(), + msg: EventMsg::TurnComplete(TurnCompleteEvent { + last_agent_message: None, + }), + }); + assert_eq!(chat.bottom_pane.is_task_running(), true); + + chat.handle_codex_event(Event { + id: "turn-regular".into(), + msg: EventMsg::TurnComplete(TurnCompleteEvent { + last_agent_message: None, + }), + }); + assert_eq!(chat.bottom_pane.is_task_running(), false); +} + #[tokio::test] async fn ctrl_c_shutdown_works_with_caps_lock() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; @@ -1865,7 +1900,7 @@ async fn exec_end_without_begin_uses_event_command() { #[tokio::test] async fn exec_history_shows_unified_exec_startup_commands() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); let begin = begin_exec_with_source( &mut chat, @@ -1892,7 +1927,7 @@ async fn exec_history_shows_unified_exec_startup_commands() { #[tokio::test] async fn exec_history_shows_unified_exec_tool_calls() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(None).await; - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); let begin = begin_exec_with_source( &mut chat, @@ -1909,7 +1944,7 @@ async fn exec_history_shows_unified_exec_tool_calls() { #[tokio::test] async fn unified_exec_end_after_task_complete_is_suppressed() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); let begin = begin_exec_with_source( &mut chat, @@ -1919,7 +1954,7 @@ async fn unified_exec_end_after_task_complete_is_suppressed() { ); drain_insert_history(&mut rx); - chat.on_task_complete(None, false); + chat.on_task_complete(Some("turn-1"), None, false); end_exec(&mut chat, begin, "", "", 0); let cells = drain_insert_history(&mut rx); @@ -1932,8 +1967,8 @@ async fn unified_exec_end_after_task_complete_is_suppressed() { #[tokio::test] async fn unified_exec_interaction_after_task_complete_is_suppressed() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; - chat.on_task_started(); - chat.on_task_complete(None, false); + chat.on_task_started(Some("turn-1"), false); + chat.on_task_complete(Some("turn-1"), None, false); chat.handle_codex_event(Event { id: "call-1".to_string(), @@ -2027,7 +2062,7 @@ async fn unified_exec_wait_before_streamed_agent_message_snapshot() { #[tokio::test] async fn unified_exec_wait_status_header_updates_on_late_command_display() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(None).await; - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); chat.unified_exec_processes.push(UnifiedExecProcessSummary { key: "proc-1".to_string(), command_display: "sleep 5".to_string(), @@ -2049,7 +2084,7 @@ async fn unified_exec_wait_status_header_updates_on_late_command_display() { #[tokio::test] async fn unified_exec_waiting_multiple_empty_snapshots() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); begin_unified_exec_startup(&mut chat, "call-wait-1", "proc-1", "just fix"); terminal_interaction(&mut chat, "call-wait-1a", "proc-1", ""); @@ -2077,7 +2112,7 @@ async fn unified_exec_waiting_multiple_empty_snapshots() { #[tokio::test] async fn unified_exec_empty_then_non_empty_snapshot() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); begin_unified_exec_startup(&mut chat, "call-wait-2", "proc-2", "just fix"); terminal_interaction(&mut chat, "call-wait-2a", "proc-2", ""); @@ -2094,7 +2129,7 @@ async fn unified_exec_empty_then_non_empty_snapshot() { #[tokio::test] async fn unified_exec_non_empty_then_empty_snapshots() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); begin_unified_exec_startup(&mut chat, "call-wait-3", "proc-3", "just fix"); terminal_interaction(&mut chat, "call-wait-3a", "proc-3", "pwd\n"); @@ -2212,7 +2247,7 @@ async fn collab_mode_shift_tab_cycles_only_when_enabled_and_idle() { assert_eq!(chat.active_collaboration_mode_kind(), ModeKind::Code); assert_eq!(chat.current_collaboration_mode(), &initial); - chat.on_task_started(); + chat.on_task_started(Some("turn-1"), false); let before = chat.active_collaboration_mode_kind(); chat.handle_key_event(KeyEvent::from(KeyCode::BackTab)); assert_eq!(chat.active_collaboration_mode_kind(), before);