From 30db592c7abc4e3bd5a1ca810e0c48fdc220d91b Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 5 Jul 2026 18:41:42 -0400 Subject: [PATCH] Protect restored startup drafts at submission --- codex-rs/tui/src/app.rs | 17 +++-- codex-rs/tui/src/app/app_server_events.rs | 3 - codex-rs/tui/src/app/input.rs | 52 -------------- codex-rs/tui/src/app/test_support.rs | 1 - codex-rs/tui/src/app/tests.rs | 2 - codex-rs/tui/src/app/tests/startup.rs | 71 +++++++++++++------ codex-rs/tui/src/app/thread_routing.rs | 8 --- codex-rs/tui/src/bottom_pane/chat_composer.rs | 49 +++++++++++++ codex-rs/tui/src/bottom_pane/mod.rs | 32 +++++++++ codex-rs/tui/src/chatwidget.rs | 5 ++ codex-rs/tui/src/chatwidget/constructor.rs | 1 + codex-rs/tui/src/chatwidget/input_flow.rs | 24 +++++++ codex-rs/tui/src/chatwidget/mcp_startup.rs | 14 ++-- codex-rs/tui/src/chatwidget/session_flow.rs | 1 + 14 files changed, 180 insertions(+), 100 deletions(-) diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 2a6f2b7639..1c02539674 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -574,7 +574,6 @@ pub(crate) struct App { primary_thread_id: Option, last_subagent_backfill_attempt: Option, primary_session_configured: Option, - startup_draft_protected: bool, pending_primary_events: VecDeque, pending_app_server_requests: PendingAppServerRequests, pending_startup_thread_start: bool, @@ -1065,7 +1064,6 @@ See the Codex keymap documentation for supported actions and examples." primary_thread_id: None, last_subagent_backfill_attempt: None, primary_session_configured: None, - startup_draft_protected: false, pending_primary_events: VecDeque::new(), pending_app_server_requests: PendingAppServerRequests::default(), pending_startup_thread_start, @@ -1117,11 +1115,16 @@ See the Codex keymap documentation for supported actions and examples." } } - if allow_startup_text { - if let Some(startup_text) = tui.take_startup_text()? { - app.chat_widget.insert_str(&startup_text); - } - app.startup_draft_protected = true; + if allow_startup_text && let Some(startup_text) = tui.take_startup_text()? { + let expected_mcp_servers = app + .config + .mcp_servers + .get() + .iter() + .filter_map(|(name, server)| server.enabled.then_some(name.clone())) + .collect::>(); + app.chat_widget + .restore_startup_draft(&startup_text, expected_mcp_servers); } let event_stream_started_at = Instant::now(); diff --git a/codex-rs/tui/src/app/app_server_events.rs b/codex-rs/tui/src/app/app_server_events.rs index bba60afe83..06ed4320d8 100644 --- a/codex-rs/tui/src/app/app_server_events.rs +++ b/codex-rs/tui/src/app/app_server_events.rs @@ -41,9 +41,6 @@ impl App { ); self.refresh_mcp_startup_expected_servers_from_config(); self.chat_widget.finish_mcp_startup_after_lag(); - if self.startup_draft_protected { - self.clear_startup_draft_protection_if_mcp_settled(); - } } AppServerEvent::ServerNotification(notification) => { self.handle_server_notification_event(app_server_client, notification) diff --git a/codex-rs/tui/src/app/input.rs b/codex-rs/tui/src/app/input.rs index cd577df63d..5d891c128b 100644 --- a/codex-rs/tui/src/app/input.rs +++ b/codex-rs/tui/src/app/input.rs @@ -95,10 +95,6 @@ impl App { app_server: &mut AppServerSession, key_event: KeyEvent, ) { - if self.startup_draft_blocks_key(key_event) { - return; - } - // Some terminals, especially on macOS, encode Option+Left/Right as Option+b/f unless // enhanced keyboard reporting is available. We only treat those word-motion fallbacks as // agent-switch shortcuts when the composer is empty so we never steal the expected @@ -272,26 +268,6 @@ impl App { && !self.chat_widget.should_handle_vim_insert_escape(key_event) } - pub(super) fn startup_draft_blocks_key(&self, key_event: KeyEvent) -> bool { - self.startup_draft_protected - && (self.primary_session_configured.is_none() || self.chat_widget.mcp_startup_active()) - && matches!( - key_event, - KeyEvent { - code: KeyCode::Enter, - modifiers: KeyModifiers::NONE, - kind: KeyEventKind::Press | KeyEventKind::Repeat, - .. - } - ) - } - - pub(super) fn clear_startup_draft_protection_if_mcp_settled(&mut self) { - if !self.chat_widget.mcp_startup_active() { - self.startup_draft_protected = false; - } - } - pub(super) fn reject_side_backtrack_esc(&mut self) { self.reset_backtrack_state(); self.chat_widget @@ -310,11 +286,6 @@ impl App { #[cfg(test)] mod tests { use super::super::test_support::make_test_app; - use crossterm::event::KeyCode; - use crossterm::event::KeyEvent; - use crossterm::event::KeyEventKind; - use crossterm::event::KeyEventState; - use crossterm::event::KeyModifiers; #[tokio::test] async fn app_keymap_shortcuts_are_disabled_while_keymap_view_is_active() { @@ -326,27 +297,4 @@ mod tests { assert!(!app.app_keymap_shortcuts_available()); } - - #[tokio::test] - async fn startup_draft_only_blocks_unmodified_enter_presses() { - let mut app = make_test_app().await; - app.startup_draft_protected = true; - - assert!(app.startup_draft_blocks_key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE))); - assert!(app.startup_draft_blocks_key(KeyEvent::new_with_kind( - KeyCode::Enter, - KeyModifiers::NONE, - KeyEventKind::Repeat, - ))); - assert!(!app.startup_draft_blocks_key(KeyEvent::new(KeyCode::Enter, KeyModifiers::SHIFT))); - assert!( - !app.startup_draft_blocks_key(KeyEvent::new(KeyCode::Char('a'), KeyModifiers::NONE)) - ); - assert!(!app.startup_draft_blocks_key(KeyEvent { - code: KeyCode::Enter, - modifiers: KeyModifiers::NONE, - kind: KeyEventKind::Release, - state: KeyEventState::NONE, - })); - } } diff --git a/codex-rs/tui/src/app/test_support.rs b/codex-rs/tui/src/app/test_support.rs index 919c9118e4..e0eafaa23b 100644 --- a/codex-rs/tui/src/app/test_support.rs +++ b/codex-rs/tui/src/app/test_support.rs @@ -61,7 +61,6 @@ pub(super) async fn make_test_app() -> App { primary_thread_id: None, last_subagent_backfill_attempt: None, primary_session_configured: None, - startup_draft_protected: false, pending_primary_events: VecDeque::new(), pending_app_server_requests: PendingAppServerRequests::default(), pending_startup_thread_start: false, diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 000cc2d44e..d8d63575bb 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -4098,7 +4098,6 @@ async fn make_test_app() -> App { primary_thread_id: None, last_subagent_backfill_attempt: None, primary_session_configured: None, - startup_draft_protected: false, pending_primary_events: VecDeque::new(), pending_app_server_requests: PendingAppServerRequests::default(), pending_startup_thread_start: false, @@ -4164,7 +4163,6 @@ async fn make_test_app_with_channels() -> ( primary_thread_id: None, last_subagent_backfill_attempt: None, primary_session_configured: None, - startup_draft_protected: false, pending_primary_events: VecDeque::new(), pending_app_server_requests: PendingAppServerRequests::default(), pending_startup_thread_start: false, diff --git a/codex-rs/tui/src/app/tests/startup.rs b/codex-rs/tui/src/app/tests/startup.rs index 133a8d04fb..23ab975c62 100644 --- a/codex-rs/tui/src/app/tests/startup.rs +++ b/codex-rs/tui/src/app/tests/startup.rs @@ -138,20 +138,20 @@ fn startup_waiting_gate_not_applied_for_resume_or_fork_session_selection() { } #[tokio::test] -async fn startup_thread_started_submits_queued_startup_input() { +async fn restored_startup_draft_waits_for_initial_mcp_settlement() { let (mut app, _app_event_rx, mut op_rx) = make_test_app_with_channels().await; app.pending_startup_thread_start = true; - app.startup_draft_protected = true; app.chat_widget .set_queue_submissions_until_session_configured(/*queue*/ true); app.chat_widget - .apply_external_edit("queued before startup completes".to_string()); - app.chat_widget - .handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + .restore_startup_draft("captured during startup", ["test".to_string()]); + let enter = KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE); + app.chat_widget.handle_key_event(enter); + assert!(app.chat_widget.queued_user_message_texts().is_empty()); assert_eq!( - app.chat_widget.queued_user_message_texts(), - vec!["queued before startup completes".to_string()] + app.chat_widget.composer_text_with_pending(), + "captured during startup" ); let mut app_server = Box::pin(crate::start_embedded_app_server_for_picker( @@ -169,8 +169,12 @@ async fn startup_thread_started_submits_queued_startup_input() { ) .await .expect("startup thread should attach"); - let enter = KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE); - assert!(!app.startup_draft_blocks_key(enter)); + + app.chat_widget.handle_key_event(enter); + assert_eq!( + app.chat_widget.composer_text_with_pending(), + "captured during startup" + ); app.chat_widget .set_mcp_startup_expected_servers(["test".to_string()]); @@ -183,7 +187,11 @@ async fn startup_thread_started_submits_queued_startup_input() { failure_reason: None, }), )); - assert!(app.startup_draft_blocks_key(enter)); + app.chat_widget.handle_key_event(enter); + assert_eq!( + app.chat_widget.composer_text_with_pending(), + "captured during startup" + ); app.handle_thread_event_now(ThreadBufferedEvent::Notification( ServerNotification::McpServerStatusUpdated(McpServerStatusUpdatedNotification { @@ -194,8 +202,18 @@ async fn startup_thread_started_submits_queued_startup_input() { failure_reason: None, }), )); - assert!(!app.startup_draft_blocks_key(enter)); - assert!(!app.startup_draft_protected); + + app.chat_widget.handle_key_event(enter); + match next_user_turn_op(&mut op_rx) { + Op::UserTurn { items, .. } => assert_eq!( + items, + vec![UserInput::Text { + text: "captured during startup".to_string(), + text_elements: Vec::new(), + }] + ), + other => panic!("expected restored startup input submission, got {other:?}"), + } app.handle_thread_event_now(ThreadBufferedEvent::Notification( ServerNotification::McpServerStatusUpdated(McpServerStatusUpdatedNotification { @@ -206,18 +224,25 @@ async fn startup_thread_started_submits_queued_startup_input() { failure_reason: None, }), )); - assert!(!app.startup_draft_blocks_key(enter)); + app.chat_widget.insert_str("later round"); + app.chat_widget.handle_key_event(enter); + assert_eq!(app.chat_widget.composer_text_with_pending(), ""); +} - match next_user_turn_op(&mut op_rx) { - Op::UserTurn { items, .. } => assert_eq!( - items, - vec![UserInput::Text { - text: "queued before startup completes".to_string(), - text_elements: Vec::new(), - }] - ), - other => panic!("expected queued startup input submission, got {other:?}"), - } +#[tokio::test] +async fn startup_without_restored_draft_uses_normal_queued_submission() { + let (mut app, _app_event_rx, _op_rx) = make_test_app_with_channels().await; + app.chat_widget + .set_queue_submissions_until_session_configured(/*queue*/ true); + app.chat_widget.insert_str("typed after the first frame"); + + app.chat_widget + .handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert_eq!( + app.chat_widget.queued_user_message_texts(), + vec!["typed after the first frame".to_string()] + ); } #[tokio::test] diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 22956f9274..b149fd53f3 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -1465,17 +1465,9 @@ impl App { ); match event { ThreadBufferedEvent::Notification(notification) => { - let starts_turn = matches!(¬ification, ServerNotification::TurnStarted(_)); - let updates_mcp_startup = - matches!(¬ification, ServerNotification::McpServerStatusUpdated(_)); self.cache_collab_receiver_threads_for_notification(¬ification); self.chat_widget .handle_server_notification(notification, /*replay_kind*/ None); - if self.startup_draft_protected && starts_turn { - self.startup_draft_protected = false; - } else if self.startup_draft_protected && updates_mcp_startup { - self.clear_startup_draft_protection_if_mcp_settled(); - } } ThreadBufferedEvent::Request(request) => { if self diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index e568436b40..191a199c8e 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -361,6 +361,7 @@ pub(crate) struct ChatComposer { placeholder_text: String, is_task_running: bool, queue_submissions: bool, + startup_draft_submission_blocked: bool, /// Slash-command draft staged for local recall after application-level dispatch. /// /// This slot is intentionally separate from `ChatComposerHistory` so inline slash commands can @@ -529,6 +530,7 @@ impl ChatComposer { placeholder_text, is_task_running: false, queue_submissions: false, + startup_draft_submission_blocked: false, pending_slash_command_history: None, skills: None, plugins: None, @@ -2759,6 +2761,9 @@ impl ChatComposer { should_queue: bool, now: Instant, ) -> (InputResult, bool) { + if self.startup_draft_submission_blocked { + return (InputResult::None, false); + } if should_queue { if let Some(pasted) = self.draft.paste_burst.flush_before_modified_input() { self.handle_paste(pasted); @@ -3891,6 +3896,10 @@ impl ChatComposer { self.queue_submissions = queue_submissions; } + pub(crate) fn set_startup_draft_submission_blocked(&mut self, blocked: bool) { + self.startup_draft_submission_blocked = blocked; + } + pub(crate) fn set_context_window(&mut self, percent: Option, used_tokens: Option) { if self.footer.context_window_percent == percent && self.footer.context_window_used_tokens == used_tokens @@ -8470,6 +8479,46 @@ mod tests { assert_eq!("explain the change\n", composer.draft.textarea.text()); } + #[test] + fn startup_draft_protection_respects_remapped_submit() { + use crate::key_hint; + use crate::keymap::RuntimeKeymap; + use crossterm::event::KeyCode; + use crossterm::event::KeyEvent; + use crossterm::event::KeyModifiers; + + let (tx, _rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + let mut composer = ChatComposer::new( + /*has_input_focus*/ true, + sender, + /*enhanced_keys_supported*/ false, + "Ask Codex to do anything".to_string(), + /*disable_paste_burst*/ false, + ); + composer + .draft + .textarea + .set_text_clearing_elements("captured during startup"); + composer + .draft + .textarea + .set_cursor(composer.draft.textarea.text().len()); + let mut keymap = RuntimeKeymap::defaults(); + keymap.composer.submit = vec![key_hint::ctrl(KeyCode::Char('j'))]; + composer.set_keymap_bindings(&keymap); + composer.set_startup_draft_submission_blocked(/*blocked*/ true); + + let (enter_result, _) = + composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + let (submit_result, _) = + composer.handle_key_event(KeyEvent::new(KeyCode::Char('j'), KeyModifiers::CONTROL)); + + assert_eq!(InputResult::None, enter_result); + assert_eq!(InputResult::None, submit_result); + assert_eq!("captured during startup\n", composer.draft.textarea.text()); + } + #[test] fn remapped_queue_does_not_fall_back_to_tab() { use crate::key_hint; diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index ad30dedf5a..d17070a45c 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -1020,6 +1020,10 @@ impl BottomPane { self.composer.set_queue_submissions(queue_submissions); } + pub(crate) fn set_startup_draft_submission_blocked(&mut self, blocked: bool) { + self.composer.set_startup_draft_submission_blocked(blocked); + } + /// Hide the status indicator while leaving task-running state untouched. pub(crate) fn hide_status_indicator(&mut self) { if self.status.take().is_some() { @@ -2864,6 +2868,34 @@ mod tests { assert!(matches!(rx.try_recv(), Ok(AppEvent::OpenApprovalsPopup))); } + #[test] + fn startup_draft_protection_does_not_block_active_view_enter() { + let (tx_raw, mut rx) = unbounded_channel::(); + let tx = AppEventSender::new(tx_raw); + let mut pane = test_pane(tx); + pane.insert_str("captured during startup"); + pane.set_startup_draft_submission_blocked(/*blocked*/ true); + pane.show_selection_view(SelectionViewParams { + title: Some("Confirm".to_string()), + items: vec![SelectionItem { + name: "Continue".to_string(), + actions: vec![Box::new(|tx: &_| { + tx.send(AppEvent::OpenApprovalsPopup); + })], + dismiss_on_select: true, + ..Default::default() + }], + ..Default::default() + }); + + let result = pane.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + + assert_eq!(result, InputResult::None); + assert!(pane.no_modal_or_popup_active()); + assert!(matches!(rx.try_recv(), Ok(AppEvent::OpenApprovalsPopup))); + assert_eq!(pane.composer_text(), "captured during startup"); + } + #[test] fn esc_routes_to_handle_key_event_when_requested() { #[derive(Default)] diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 42903d2e7f..16115a8011 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -598,6 +598,11 @@ pub(crate) struct ChatWidget { mcp_startup_status: Option>, /// Expected MCP servers for the current startup round, seeded from enabled local config. mcp_startup_expected_servers: Option>, + /// Initial MCP servers that must settle before a restored startup draft can be submitted. + /// + /// `Some(_)` also records that startup text was actually restored. The set may become empty + /// before session configuration arrives, but protection remains until both conditions settle. + startup_draft_pending_mcp_servers: Option>, /// After startup settles, ignore stale updates until enough notifications confirm a new round. mcp_startup_ignore_updates_until_next_start: bool, /// A lag signal for the next round means terminal-only updates are enough to settle it. diff --git a/codex-rs/tui/src/chatwidget/constructor.rs b/codex-rs/tui/src/chatwidget/constructor.rs index 7fb6fe54e9..06f6c98593 100644 --- a/codex-rs/tui/src/chatwidget/constructor.rs +++ b/codex-rs/tui/src/chatwidget/constructor.rs @@ -159,6 +159,7 @@ impl ChatWidget { unified_exec_processes: Vec::new(), mcp_startup_status: None, mcp_startup_expected_servers: None, + startup_draft_pending_mcp_servers: None, mcp_startup_ignore_updates_until_next_start: false, mcp_startup_allow_terminal_only_next_round: false, mcp_startup_pending_next_round: HashMap::new(), diff --git a/codex-rs/tui/src/chatwidget/input_flow.rs b/codex-rs/tui/src/chatwidget/input_flow.rs index 916c54cd5e..48681227f9 100644 --- a/codex-rs/tui/src/chatwidget/input_flow.rs +++ b/codex-rs/tui/src/chatwidget/input_flow.rs @@ -7,6 +7,30 @@ use super::*; impl ChatWidget { + pub(crate) fn restore_startup_draft(&mut self, text: &str, expected_mcp_servers: I) + where + I: IntoIterator, + { + self.insert_str(text); + self.startup_draft_pending_mcp_servers = Some(expected_mcp_servers.into_iter().collect()); + self.bottom_pane + .set_startup_draft_submission_blocked(/*blocked*/ true); + self.finish_startup_draft_protection_if_ready(); + } + + pub(super) fn finish_startup_draft_protection_if_ready(&mut self) { + let startup_ready = self.is_session_configured() + && self + .startup_draft_pending_mcp_servers + .as_ref() + .is_some_and(HashSet::is_empty); + if startup_ready { + self.startup_draft_pending_mcp_servers = None; + self.bottom_pane + .set_startup_draft_submission_blocked(/*blocked*/ false); + } + } + pub(super) fn handle_composer_input_result( &mut self, input_result: InputResult, diff --git a/codex-rs/tui/src/chatwidget/mcp_startup.rs b/codex-rs/tui/src/chatwidget/mcp_startup.rs index cc429896f8..2ba54b832a 100644 --- a/codex-rs/tui/src/chatwidget/mcp_startup.rs +++ b/codex-rs/tui/src/chatwidget/mcp_startup.rs @@ -181,10 +181,6 @@ impl ChatWidget { self.mcp_startup_expected_servers = Some(server_names.into_iter().collect()); } - pub(crate) fn mcp_startup_active(&self) -> bool { - self.mcp_startup_status.is_some() - } - pub(super) fn finish_mcp_startup(&mut self, failed: Vec, cancelled: Vec) { if !cancelled.is_empty() { self.on_warning(format!( @@ -215,6 +211,10 @@ impl ChatWidget { } pub(crate) fn finish_mcp_startup_after_lag(&mut self) { + if let Some(pending) = &mut self.startup_draft_pending_mcp_servers { + pending.clear(); + } + self.finish_startup_draft_protection_if_ready(); if self.mcp_startup_ignore_updates_until_next_start { if self.mcp_startup_pending_next_round.is_empty() { self.mcp_startup_pending_next_round_saw_starting = false; @@ -267,6 +267,12 @@ impl ChatWidget { &mut self, notification: McpServerStatusUpdatedNotification, ) { + if notification.status != McpServerStartupState::Starting + && let Some(pending) = &mut self.startup_draft_pending_mcp_servers + { + pending.remove(¬ification.name); + } + self.finish_startup_draft_protection_if_ready(); let status = match notification.status { McpServerStartupState::Starting => McpStartupStatus::Starting, McpServerStartupState::Ready => McpStartupStatus::Ready, diff --git a/codex-rs/tui/src/chatwidget/session_flow.rs b/codex-rs/tui/src/chatwidget/session_flow.rs index 177700d4f4..b54dc7d88a 100644 --- a/codex-rs/tui/src/chatwidget/session_flow.rs +++ b/codex-rs/tui/src/chatwidget/session_flow.rs @@ -20,6 +20,7 @@ impl ChatWidget { self.session_network_proxy = session.network_proxy.clone(); let previous_thread_id = self.thread_id; self.thread_id = Some(session.thread_id); + self.finish_startup_draft_protection_if_ready(); self.bottom_pane .set_queue_submissions(/*queue_submissions*/ false); if previous_thread_id != self.thread_id {