From 4b7b5cd7cc207668f60e2aabf99ef8c437c09060 Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Sun, 1 Feb 2026 15:40:04 -0800 Subject: [PATCH] Sync rollback model-visible mode state across core and tui --- .../src/protocol/thread_history.rs | 10 +- codex-rs/core/src/codex.rs | 154 ++++++++++++------ codex-rs/core/src/rollout/truncation.rs | 1 + codex-rs/core/src/state/mod.rs | 1 + codex-rs/core/src/state/session.rs | 38 ++++- codex-rs/protocol/src/protocol.rs | 12 ++ codex-rs/tui/src/app.rs | 98 ++++++++++- codex-rs/tui/src/app_backtrack.rs | 28 +++- codex-rs/tui/src/chatwidget.rs | 56 +++++++ 9 files changed, 333 insertions(+), 65 deletions(-) diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 39efe47655..751bc2663c 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -490,7 +490,10 @@ mod tests { EventMsg::AgentMessage(AgentMessageEvent { message: "A2".into(), }), - EventMsg::ThreadRolledBack(ThreadRolledBackEvent { num_turns: 1 }), + EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + model_visible_state: None, + }), EventMsg::UserMessage(UserMessageEvent { message: "Third".into(), images: None, @@ -565,7 +568,10 @@ mod tests { EventMsg::AgentMessage(AgentMessageEvent { message: "A2".into(), }), - EventMsg::ThreadRolledBack(ThreadRolledBackEvent { num_turns: 99 }), + EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 99, + model_visible_state: None, + }), ]; let turns = build_turns_from_event_msgs(&events); diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index db70d82cc9..61ae4a9972 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -185,6 +185,7 @@ use crate::skills::injection::app_id_from_path; use crate::skills::injection::tool_kind_for_path; use crate::skills::resolve_skill_dependencies_for_turn; use crate::state::ActiveTurn; +use crate::state::PendingModelVisibleStateSync; use crate::state::SessionServices; use crate::state::SessionState; use crate::state_db; @@ -1405,19 +1406,23 @@ impl Session { &self, previous_collaboration_mode: &CollaborationMode, next_collaboration_mode: Option<&CollaborationMode>, - force_inject_collaboration_instructions: bool, + pending_model_visible_state_sync: &PendingModelVisibleStateSync, ) -> Option { - if let Some(next_mode) = next_collaboration_mode { - if !force_inject_collaboration_instructions && previous_collaboration_mode == next_mode - { - return None; + let next_mode = next_collaboration_mode?; + let should_emit = match pending_model_visible_state_sync { + PendingModelVisibleStateSync::None => previous_collaboration_mode != next_mode, + PendingModelVisibleStateSync::Snapshot(snapshot) => { + snapshot.collaboration_mode.as_ref() != Some(next_mode) } - // If the next mode has empty developer instructions, this returns None and we emit no - // update, so prior collaboration instructions remain in the prompt history. - Some(DeveloperInstructions::from_collaboration_mode(&next.collaboration_mode)?.into()) - } else { - None + PendingModelVisibleStateSync::ForceEmitAll => true, + }; + if !should_emit { + return None; } + + // If the next mode has empty developer instructions, this returns None and we emit no + // update, so prior collaboration instructions remain in the prompt history. + Some(DeveloperInstructions::from_collaboration_mode(next_mode)?.into()) } fn build_settings_update_items( @@ -1426,7 +1431,7 @@ impl Session { current_context: &TurnContext, previous_collaboration_mode: &CollaborationMode, next_collaboration_mode: Option<&CollaborationMode>, - force_inject_collaboration_instructions: bool, + pending_model_visible_state_sync: &PendingModelVisibleStateSync, ) -> Vec { let mut update_items = Vec::new(); if let Some(env_item) = @@ -1442,7 +1447,7 @@ impl Session { if let Some(collaboration_mode_item) = self.build_collaboration_mode_update_item( previous_collaboration_mode, next_collaboration_mode, - force_inject_collaboration_instructions, + pending_model_visible_state_sync, ) { update_items.push(collaboration_mode_item); } @@ -2731,6 +2736,7 @@ mod handlers { use codex_protocol::protocol::ListCustomPromptsResponseEvent; use codex_protocol::protocol::ListSkillsResponseEvent; use codex_protocol::protocol::McpServerRefreshConfig; + use codex_protocol::protocol::ModelVisibleState; use codex_protocol::protocol::Op; use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::ReviewRequest; @@ -2756,6 +2762,7 @@ mod handlers { use tracing::info; use tracing::warn; + use crate::state::PendingModelVisibleStateSync; use crate::state::SessionState; pub async fn interrupt(sess: &Arc) { @@ -2785,20 +2792,16 @@ mod handlers { } let current_context = sess.new_default_turn_with_sub_id(sub_id).await; - let force_inject_collaboration_instructions = { + let pending_model_visible_state_sync = { let mut state = sess.state.lock().await; - let force = state.force_inject_collaboration_instructions; - if next_collaboration_mode.is_some() { - state.force_inject_collaboration_instructions = false; - } - force + state.take_pending_model_visible_state_sync(next_collaboration_mode.is_some()) }; let update_items = sess.build_settings_update_items( Some(&previous_context), ¤t_context, &previous_collaboration_mode, next_collaboration_mode.as_ref(), - force_inject_collaboration_instructions, + &pending_model_visible_state_sync, ); if !update_items.is_empty() { sess.record_conversation_items(¤t_context, &update_items) @@ -2874,18 +2877,16 @@ mod handlers { // Attempt to inject input into current task if let Err(items) = sess.inject_input(items).await { sess.seed_initial_context_if_needed(¤t_context).await; - let force_inject_collaboration_instructions = { + let pending_model_visible_state_sync = { let mut state = sess.state.lock().await; - let force = state.force_inject_collaboration_instructions; - state.force_inject_collaboration_instructions = false; - force + state.take_pending_model_visible_state_sync(next_collaboration_mode.is_some()) }; let update_items = sess.build_settings_update_items( previous_context.as_ref(), ¤t_context, &previous_collaboration_mode, next_collaboration_mode.as_ref(), - force_inject_collaboration_instructions, + &pending_model_visible_state_sync, ); if !update_items.is_empty() { sess.record_conversation_items(¤t_context, &update_items) @@ -3191,20 +3192,23 @@ mod handlers { // version of the history. let user_turns = Session::user_turn_count(history.raw_items()); sess.replace_history(history.raw_items().to_vec()).await; - { + let model_visible_state = { let mut state = sess.state.lock().await; apply_rollback_turn_context_history( &mut state, existing_turn_context_history, num_turns, user_turns, - ); - } + ) + }; sess.recompute_token_usage(turn_context.as_ref()).await; sess.send_event_raw_flushed(Event { id: turn_context.sub_id.clone(), - msg: EventMsg::ThreadRolledBack(ThreadRolledBackEvent { num_turns }), + msg: EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns, + model_visible_state: Some(model_visible_state), + }), }) .await; } @@ -3278,13 +3282,14 @@ mod handlers { /// Apply turn-context history updates after rollback. /// /// Truncates the stored history by the requested rollback depth, then pads to match the - /// surviving user-turn count. + /// surviving user-turn count. Also updates rollback sync state for model-visible settings + /// and returns the authoritative model-visible state after rollback. fn apply_rollback_turn_context_history( state: &mut SessionState, existing_turn_context_history: Vec>, num_turns: u32, user_turns: usize, - ) { + ) -> ModelVisibleState { let mut updated_turn_context_history = existing_turn_context_history; let truncated_len = updated_turn_context_history .len() @@ -3294,18 +3299,41 @@ mod handlers { updated_turn_context_history.resize_with(user_turns, || None); } state.set_turn_context_history(updated_turn_context_history); - if state.turn_context_history.len() == user_turns - && let Some(collaboration_mode) = state + + let rollback_snapshot = if state.turn_context_history.len() == user_turns { + state .turn_context_history .last() .and_then(Option::as_ref) - .and_then(|turn_context| turn_context.collaboration_mode.clone()) + .and_then(model_visible_state_from_turn_context) + } else { + None + }; + + if let Some(model_visible_state) = rollback_snapshot + && let Some(collaboration_mode) = model_visible_state.collaboration_mode.clone() { state.session_configuration.collaboration_mode = collaboration_mode; - state.force_inject_collaboration_instructions = false; + state.pending_model_visible_state_sync = + PendingModelVisibleStateSync::Snapshot(model_visible_state); } else { - state.force_inject_collaboration_instructions = true; + state.pending_model_visible_state_sync = PendingModelVisibleStateSync::ForceEmitAll; } + + ModelVisibleState { + collaboration_mode: Some(state.session_configuration.collaboration_mode.clone()), + } + } + + fn model_visible_state_from_turn_context( + turn_context: &TurnContextItem, + ) -> Option { + turn_context + .collaboration_mode + .clone() + .map(|collaboration_mode| ModelVisibleState { + collaboration_mode: Some(collaboration_mode), + }) } pub async fn shutdown(sess: &Arc, sub_id: String) -> bool { @@ -4769,12 +4797,14 @@ mod tests { use crate::protocol::CompactedItem; use crate::protocol::CreditsSnapshot; use crate::protocol::InitialHistory; + use crate::protocol::ModelVisibleState; use crate::protocol::RateLimitSnapshot; use crate::protocol::RateLimitWindow; use crate::protocol::ResumedHistory; use crate::protocol::TokenCountEvent; use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; + use crate::state::PendingModelVisibleStateSync; use crate::state::TaskKind; use crate::tasks::SessionTask; use crate::tasks::SessionTaskContext; @@ -5192,9 +5222,26 @@ mod tests { handlers::thread_rollback(&sess, "sub-1".to_string(), 1).await; let rollback_event = wait_for_thread_rolled_back(&rx).await; assert_eq!(rollback_event.num_turns, 1); + assert_eq!( + rollback_event.model_visible_state, + Some(ModelVisibleState { + collaboration_mode: Some(plan_mode.clone()), + }) + ); let current_mode = sess.current_collaboration_mode().await; assert_eq!(current_mode.mode, ModeKind::Plan); + + let pending_model_visible_state_sync = { + let mut state = sess.state.lock().await; + state.take_pending_model_visible_state_sync(true) + }; + assert_eq!( + pending_model_visible_state_sync, + PendingModelVisibleStateSync::Snapshot(ModelVisibleState { + collaboration_mode: Some(plan_mode), + }) + ); } #[tokio::test] @@ -5232,14 +5279,21 @@ mod tests { handlers::thread_rollback(&sess, "sub-1".to_string(), 1).await; let rollback_event = wait_for_thread_rolled_back(&rx).await; assert_eq!(rollback_event.num_turns, 1); + assert_eq!( + rollback_event.model_visible_state, + Some(ModelVisibleState { + collaboration_mode: Some(collaboration_mode.clone()), + }) + ); - let force_inject_collaboration_instructions = { + let pending_model_visible_state_sync = { let mut state = sess.state.lock().await; - let force = state.force_inject_collaboration_instructions; - state.force_inject_collaboration_instructions = false; - force + state.take_pending_model_visible_state_sync(true) }; - assert!(force_inject_collaboration_instructions); + assert_eq!( + pending_model_visible_state_sync, + PendingModelVisibleStateSync::ForceEmitAll + ); let current_context = sess.new_default_turn_with_sub_id("sub-2".to_string()).await; let previous_collaboration_mode = { @@ -5251,7 +5305,7 @@ mod tests { ¤t_context, &previous_collaboration_mode, Some(&previous_collaboration_mode), - force_inject_collaboration_instructions, + &pending_model_visible_state_sync, ); let expected_item: ResponseItem = DeveloperInstructions::from_collaboration_mode(&previous_collaboration_mode) @@ -5376,19 +5430,27 @@ mod tests { handlers::thread_rollback(&sess, "sub-1".to_string(), 1).await; let rollback_event = wait_for_thread_rolled_back(&rx).await; pretty_assertions::assert_eq!(rollback_event.num_turns, 1); + assert_eq!( + rollback_event.model_visible_state, + Some(ModelVisibleState { + collaboration_mode: Some(code_mode.clone()), + }) + ); - let (collaboration_mode, force_inject_collaboration_instructions) = { + let (collaboration_mode, pending_model_visible_state_sync) = { let mut state = sess.state.lock().await; - let force = state.force_inject_collaboration_instructions; - state.force_inject_collaboration_instructions = false; + let pending = state.take_pending_model_visible_state_sync(true); ( state.session_configuration.collaboration_mode.clone(), - force, + pending, ) }; pretty_assertions::assert_eq!(collaboration_mode, code_mode); - assert!(force_inject_collaboration_instructions); + assert_eq!( + pending_model_visible_state_sync, + PendingModelVisibleStateSync::ForceEmitAll + ); } #[tokio::test] diff --git a/codex-rs/core/src/rollout/truncation.rs b/codex-rs/core/src/rollout/truncation.rs index c50eacc48b..c456b88353 100644 --- a/codex-rs/core/src/rollout/truncation.rs +++ b/codex-rs/core/src/rollout/truncation.rs @@ -173,6 +173,7 @@ mod tests { RolloutItem::ResponseItem(assistant_msg("a2")), RolloutItem::EventMsg(EventMsg::ThreadRolledBack(ThreadRolledBackEvent { num_turns: 1, + model_visible_state: None, })), RolloutItem::ResponseItem(user_msg("u3")), RolloutItem::ResponseItem(assistant_msg("a3")), diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index 642433a786..2e3000a191 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -3,6 +3,7 @@ mod session; mod turn; pub(crate) use service::SessionServices; +pub(crate) use session::PendingModelVisibleStateSync; pub(crate) use session::SessionState; pub(crate) use turn::ActiveTurn; pub(crate) use turn::RunningTask; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index bed23779be..1d4725d2f3 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -1,6 +1,7 @@ //! Session-wide mutable state. use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::ModelVisibleState; use codex_protocol::protocol::TurnContextItem; use std::collections::HashMap; use std::collections::HashSet; @@ -12,17 +13,25 @@ use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; use crate::truncate::TruncationPolicy; +/// One-shot synchronization state for model-visible settings after rollback/backtrack. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub(crate) enum PendingModelVisibleStateSync { + /// No synchronization is pending. + #[default] + None, + /// Compare the next submitted model-visible state against this rollback snapshot. + Snapshot(ModelVisibleState), + /// Snapshot is missing/unreliable, so emit all tracked model-visible updates once. + ForceEmitAll, +} + /// Persistent, session-scoped state previously stored directly on `Session`. pub(crate) struct SessionState { pub(crate) session_configuration: SessionConfiguration, pub(crate) history: ContextManager, pub(crate) turn_context_history: Vec>, - /// Force a developer-instruction update for collaboration mode on the next turn. - /// - /// This is set when rollback/backtrack lacks a reliable turn context for the latest user turn, - /// so core should re-emit the collaboration developer instructions even if the mode is - /// unchanged. - pub(crate) force_inject_collaboration_instructions: bool, + /// Pending one-shot sync for model-visible state after rollback/backtrack. + pub(crate) pending_model_visible_state_sync: PendingModelVisibleStateSync, pub(crate) latest_rate_limits: Option, pub(crate) server_reasoning_included: bool, pub(crate) dependency_env: HashMap, @@ -42,7 +51,7 @@ impl SessionState { session_configuration, history, turn_context_history: Vec::new(), - force_inject_collaboration_instructions: false, + pending_model_visible_state_sync: PendingModelVisibleStateSync::None, latest_rate_limits: None, server_reasoning_included: false, dependency_env: HashMap::new(), @@ -108,6 +117,21 @@ impl SessionState { self.turn_context_history = turn_context_history; } + /// Consume pending model-visible state sync only when a collaboration update is being built. + /// + /// If the next operation does not carry a collaboration update (for example, model-only + /// overrides), keep the pending sync for a later turn that does. + pub(crate) fn take_pending_model_visible_state_sync( + &mut self, + has_collaboration_update: bool, + ) -> PendingModelVisibleStateSync { + if has_collaboration_update { + std::mem::take(&mut self.pending_model_visible_state_sync) + } else { + PendingModelVisibleStateSync::None + } + } + pub(crate) fn clone_history(&self) -> ContextManager { self.history.clone() } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index bf193cca9c..992299832a 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1979,6 +1979,18 @@ pub struct UndoCompletedEvent { pub struct ThreadRolledBackEvent { /// Number of user turns that were removed from context. pub num_turns: u32, + /// Model-visible session state after rollback. + /// + /// This lets clients synchronize UI mode/model indicators with core after history rewind. + /// Optional for backward compatibility with older persisted events. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model_visible_state: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +pub struct ModelVisibleState { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub collaboration_mode: Option, } #[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, TS)] diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 03d1b7e188..a4862d6fa2 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -2528,13 +2528,17 @@ mod tests { use codex_core::protocol::ErrorEvent; use codex_core::protocol::Event; use codex_core::protocol::EventMsg; + use codex_core::protocol::ModelVisibleState; use codex_core::protocol::SandboxPolicy; use codex_core::protocol::SessionConfiguredEvent; use codex_core::protocol::SessionSource; + use codex_core::protocol::ThreadRolledBackEvent; use codex_otel::OtelManager; use codex_protocol::ThreadId; + use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::config_types::ModeKind; + use codex_protocol::config_types::Settings; use codex_protocol::user_input::TextElement; use insta::assert_snapshot; use pretty_assertions::assert_eq; @@ -2742,6 +2746,17 @@ mod tests { } } + fn collaboration_mode(mode: ModeKind) -> CollaborationMode { + CollaborationMode { + mode, + settings: Settings { + model: "gpt-test".to_string(), + reasoning_effort: None, + developer_instructions: None, + }, + } + } + fn test_otel_manager(config: &Config, model: &str) -> OtelManager { let model_info = ModelsManager::construct_model_info_offline(model, config); OtelManager::new( @@ -3222,7 +3237,7 @@ mod tests { } #[tokio::test] - async fn backtrack_to_custom_clears_collaboration_mode() { + async fn backtrack_missing_mode_data_preserves_current_mode() { let (mut app, _app_event_rx, mut op_rx) = make_test_app_with_channels().await; app.chat_widget .set_feature_enabled(Feature::CollaborationModes, true); @@ -3280,7 +3295,7 @@ mod tests { assert_eq!( app.chat_widget.active_collaboration_mode_kind(), - ModeKind::Custom + ModeKind::Plan ); app.chat_widget @@ -3295,7 +3310,84 @@ mod tests { else { panic!("expected Op::UserTurn"); }; - assert_eq!(collaboration_mode, None); + let collaboration_mode = collaboration_mode.expect("expected collaboration mode"); + assert_eq!(collaboration_mode.mode, ModeKind::Plan); + } + + #[tokio::test] + async fn backtrack_success_applies_authoritative_mode_from_core() { + let (mut app, _app_event_rx, _op_rx) = make_test_app_with_channels().await; + app.chat_widget + .set_feature_enabled(Feature::CollaborationModes, true); + + let plan_mask = plan_mask(); + let code_mask = code_mask(); + + let thread_id = ThreadId::new(); + app.chat_widget.handle_codex_event(Event { + id: String::new(), + msg: EventMsg::SessionConfigured(SessionConfiguredEvent { + session_id: thread_id, + forked_from_id: None, + thread_name: None, + model: "gpt-test".to_string(), + model_provider_id: "test-provider".to_string(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::ReadOnly, + cwd: PathBuf::from("/home/user/project"), + reasoning_effort: None, + history_log_id: 0, + history_entry_count: 0, + initial_messages: None, + rollout_path: Some(PathBuf::new()), + }), + }); + + app.chat_widget.set_collaboration_mask(code_mask.clone()); + + app.transcript_cells = vec![ + Arc::new(UserHistoryCell { + message: "plan".to_string(), + text_elements: Vec::new(), + local_image_paths: Vec::new(), + collaboration_mode: Some(plan_mask), + }) as Arc, + Arc::new(AgentMessageCell::new( + vec![Line::from("plan response")], + true, + )) as Arc, + Arc::new(UserHistoryCell { + message: "code".to_string(), + text_elements: Vec::new(), + local_image_paths: Vec::new(), + collaboration_mode: Some(code_mask), + }) as Arc, + ]; + + app.backtrack.base_id = Some(thread_id); + app.backtrack.primed = true; + app.backtrack.nth_user_message = 0; + let selection = app + .confirm_backtrack_from_main() + .expect("backtrack selection"); + app.apply_backtrack_rollback(selection); + assert_eq!( + app.chat_widget.active_collaboration_mode_kind(), + ModeKind::Plan + ); + + app.handle_backtrack_event(&EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + model_visible_state: Some(ModelVisibleState { + collaboration_mode: Some(collaboration_mode(ModeKind::Code)), + }), + })); + + assert_eq!( + app.chat_widget.active_collaboration_mode_kind(), + ModeKind::Code + ); + assert!(app.backtrack.pending_rollback.is_none()); } #[tokio::test] diff --git a/codex-rs/tui/src/app_backtrack.rs b/codex-rs/tui/src/app_backtrack.rs index 5a3dda029d..c72628d298 100644 --- a/codex-rs/tui/src/app_backtrack.rs +++ b/codex-rs/tui/src/app_backtrack.rs @@ -37,6 +37,7 @@ use codex_core::protocol::CodexErrorInfo; use codex_core::protocol::ErrorEvent; use codex_core::protocol::EventMsg; use codex_core::protocol::Op; +use codex_core::protocol::ThreadRolledBackEvent; use codex_protocol::ThreadId; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::user_input::TextElement; @@ -97,6 +98,7 @@ pub(crate) struct BacktrackSelection { pub(crate) struct PendingBacktrackRollback { pub(crate) selection: BacktrackSelection, pub(crate) thread_id: Option, + /// Active mask before optimistic rollback preview mode changes. pub(crate) previous_collaboration_mask: Option, } @@ -207,11 +209,14 @@ impl App { return; } - // Optimistically apply the target mode for the rollback preview, but remember the - // prior mode so we can restore it if core rejects the rollback. - let previous_collaboration_mask = self - .chat_widget - .replace_collaboration_mask(selection.collaboration_mode.clone()); + // Optimistically apply the target mode for rollback preview when available. + // If the selected history item has no mode data (legacy/replayed entries), preserve the + // current mode until core returns authoritative rollback state. + let previous_collaboration_mask = self.chat_widget.active_collaboration_mask(); + if let Some(collaboration_mode) = selection.collaboration_mode.clone() { + self.chat_widget + .replace_collaboration_mask(Some(collaboration_mode)); + } let prefill = selection.prefill.clone(); let text_elements = selection.text_elements.clone(); let local_image_paths = selection.local_image_paths.clone(); @@ -463,7 +468,7 @@ impl App { pub(crate) fn handle_backtrack_event(&mut self, event: &EventMsg) { match event { - EventMsg::ThreadRolledBack(_) => self.finish_pending_backtrack(), + EventMsg::ThreadRolledBack(rollback) => self.finish_pending_backtrack(rollback), EventMsg::Error(ErrorEvent { codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), .. @@ -486,7 +491,7 @@ impl App { /// /// We ignore events that do not correspond to the currently active thread to avoid applying /// stale updates after a session switch. - fn finish_pending_backtrack(&mut self) { + fn finish_pending_backtrack(&mut self, rollback: &ThreadRolledBackEvent) { let Some(pending) = self.backtrack.pending_rollback.take() else { return; }; @@ -494,6 +499,15 @@ impl App { // Ignore rollbacks targeting a prior thread. return; } + + if let Some(model_visible_state) = rollback.model_visible_state.as_ref() { + self.chat_widget + .apply_model_visible_state(model_visible_state); + } else { + self.chat_widget + .replace_collaboration_mask(pending.previous_collaboration_mask.clone()); + } + self.trim_transcript_for_backtrack(pending.selection.nth_user_message); self.backtrack_render_pending = true; } diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index d28f4a47e1..35fa4d4242 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -69,6 +69,7 @@ use codex_core::protocol::McpStartupStatus; use codex_core::protocol::McpStartupUpdateEvent; use codex_core::protocol::McpToolCallBeginEvent; use codex_core::protocol::McpToolCallEndEvent; +use codex_core::protocol::ModelVisibleState; use codex_core::protocol::Op; use codex_core::protocol::PatchApplyBeginEvent; use codex_core::protocol::RateLimitSnapshot; @@ -5371,6 +5372,10 @@ impl ChatWidget { self.active_mode_kind() } + pub(crate) fn active_collaboration_mask(&self) -> Option { + self.active_collaboration_mask.clone() + } + fn is_session_configured(&self) -> bool { self.thread_id.is_some() } @@ -5530,6 +5535,57 @@ impl ChatWidget { previous } + /// Synchronize collaboration-mode UI state from core's model-visible snapshot. + /// + /// This is primarily used after rollback/backtrack so the composer mode and model display + /// reflect core's authoritative session state. + pub(crate) fn apply_model_visible_state(&mut self, state: &ModelVisibleState) { + let Some(collaboration_mode) = state.collaboration_mode.clone() else { + return; + }; + self.current_collaboration_mode = collaboration_mode.clone(); + if self.collaboration_modes_enabled() { + self.active_collaboration_mask = self.collaboration_mask_from_mode(&collaboration_mode); + self.update_collaboration_mode_indicator(); + } + self.refresh_model_display(); + self.request_redraw(); + } + + fn collaboration_mask_from_mode( + &self, + collaboration_mode: &CollaborationMode, + ) -> Option { + if collaboration_mode.mode == ModeKind::Custom { + return None; + } + + let mut mask = collaboration_modes::mask_for_kind( + self.models_manager.as_ref(), + collaboration_mode.mode, + ) + .unwrap_or_else(|| CollaborationModeMask { + name: match collaboration_mode.mode { + ModeKind::Plan => "Plan", + ModeKind::Code => "Code", + ModeKind::PairProgramming => "Pair Programming", + ModeKind::Execute => "Execute", + ModeKind::Custom => "Custom", + } + .to_string(), + mode: Some(collaboration_mode.mode), + model: None, + reasoning_effort: None, + developer_instructions: None, + }); + mask.mode = Some(collaboration_mode.mode); + mask.model = Some(collaboration_mode.settings.model.clone()); + mask.reasoning_effort = Some(collaboration_mode.settings.reasoning_effort); + mask.developer_instructions = + Some(collaboration_mode.settings.developer_instructions.clone()); + Some(mask) + } + fn connectors_enabled(&self) -> bool { self.config.features.enabled(Feature::Apps) }