From b6ca728faf621c265355c0cdad722950d274a82c Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Wed, 8 Jul 2026 22:54:01 -0700 Subject: [PATCH] Flatten turn collaboration mode state --- codex-rs/core/src/compact.rs | 2 +- codex-rs/core/src/compact_remote.rs | 2 +- codex-rs/core/src/compact_remote_v2.rs | 2 +- codex-rs/core/src/compact_token_budget.rs | 2 +- codex-rs/core/src/context_manager/updates.rs | 8 ++--- codex-rs/core/src/session/inject.rs | 2 +- codex-rs/core/src/session/review.rs | 5 ++- .../session/rollout_reconstruction_tests.rs | 16 ++++----- codex-rs/core/src/session/tests.rs | 8 ++--- codex-rs/core/src/session/turn.rs | 4 +-- codex-rs/core/src/session/turn_context.rs | 35 +++++++++++++------ codex-rs/core/src/stream_events_utils.rs | 4 +-- codex-rs/core/src/tasks/lifecycle.rs | 3 +- codex-rs/core/src/tasks/regular.rs | 2 +- codex-rs/core/src/tasks/user_shell.rs | 2 +- codex-rs/core/src/tools/handlers/plan.rs | 2 +- 16 files changed, 58 insertions(+), 41 deletions(-) diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index d2a9fd5222..3a32441518 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -130,7 +130,7 @@ pub(crate) async fn run_compact_task( trace_id: turn_context.trace_id.clone(), started_at: turn_context.turn_timing_state.started_at_unix_secs().await, model_context_window: turn_context.model_context_window(), - collaboration_mode_kind: turn_context.collaboration_mode.mode, + collaboration_mode_kind: turn_context.mode, }); sess.send_event(&turn_context, start_event).await; run_compact_task_inner( diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 4b1db93555..12233197f2 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -81,7 +81,7 @@ pub(crate) async fn run_remote_compact_task( trace_id: turn_context.trace_id.clone(), started_at: turn_context.turn_timing_state.started_at_unix_secs().await, model_context_window: turn_context.model_context_window(), - collaboration_mode_kind: turn_context.collaboration_mode.mode, + collaboration_mode_kind: turn_context.mode, }); sess.send_event(&turn_context, start_event).await; diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 1571602cc1..bf9527f34e 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -92,7 +92,7 @@ pub(crate) async fn run_remote_compact_task( trace_id: turn_context.trace_id.clone(), started_at: turn_context.turn_timing_state.started_at_unix_secs().await, model_context_window: turn_context.model_context_window(), - collaboration_mode_kind: turn_context.collaboration_mode.mode, + collaboration_mode_kind: turn_context.mode, }); sess.send_event(&turn_context, start_event).await; diff --git a/codex-rs/core/src/compact_token_budget.rs b/codex-rs/core/src/compact_token_budget.rs index 710e7cefe5..0d71e7ab2f 100644 --- a/codex-rs/core/src/compact_token_budget.rs +++ b/codex-rs/core/src/compact_token_budget.rs @@ -31,7 +31,7 @@ pub(crate) async fn run_manual_compact_task( trace_id: turn_context.trace_id.clone(), started_at: turn_context.turn_timing_state.started_at_unix_secs().await, model_context_window: turn_context.model_context_window(), - collaboration_mode_kind: turn_context.collaboration_mode.mode, + collaboration_mode_kind: turn_context.mode, }); sess.send_event(&turn_context, start_event).await; diff --git a/codex-rs/core/src/context_manager/updates.rs b/codex-rs/core/src/context_manager/updates.rs index d4cf65e86d..ff392a127d 100644 --- a/codex-rs/core/src/context_manager/updates.rs +++ b/codex-rs/core/src/context_manager/updates.rs @@ -70,13 +70,11 @@ fn build_collaboration_mode_update_item( } let prev = previous?; - if prev.collaboration_mode.as_ref() != Some(&next.collaboration_mode) { + let collaboration_mode = next.collaboration_mode(); + if prev.collaboration_mode.as_ref() != Some(&collaboration_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( - CollaborationModeInstructions::from_collaboration_mode(&next.collaboration_mode)? - .render(), - ) + Some(CollaborationModeInstructions::from_collaboration_mode(&collaboration_mode)?.render()) } else { None } diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index 5c5b7d1feb..b0b9ef7926 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -86,7 +86,7 @@ impl Session { let turn_context = self .new_default_turn_with_sub_id(uuid::Uuid::new_v4().to_string()) .await; - if turn_context.collaboration_mode.mode == ModeKind::Plan { + if turn_context.mode == ModeKind::Plan { self.clear_reserved_idle_turn(&turn_state).await; self.maybe_start_turn_for_pending_work().await; return Err(TryStartTurnIfIdleError::new( diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 6b68447b92..9a89c1beeb 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -130,7 +130,10 @@ pub(super) async fn spawn_review_thread( timezone: parent_turn_context.timezone.clone(), app_server_client_name: parent_turn_context.app_server_client_name.clone(), developer_instructions: None, - collaboration_mode: parent_turn_context.collaboration_mode.clone(), + mode: parent_turn_context.mode, + collaboration_mode_developer_instructions: parent_turn_context + .collaboration_mode_developer_instructions + .clone(), multi_agent_version: MultiAgentVersion::Disabled, personality: parent_turn_context.personality, approval_policy: parent_turn_context.approval_policy.clone(), diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 460ae21d1c..d24c80dafc 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -180,7 +180,7 @@ async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previ model: previous_model.to_string(), comp_hash: None, personality: turn_context.personality, - collaboration_mode: Some(turn_context.collaboration_mode.clone()), + collaboration_mode: Some(turn_context.collaboration_mode()), multi_agent_version: None, multi_agent_mode: None, realtime_active: Some(turn_context.realtime_active), @@ -227,7 +227,7 @@ async fn record_initial_history_resumed_hydrates_previous_turn_settings_from_lif model: previous_model.to_string(), comp_hash: Some("comp-hash-a".to_string()), personality: turn_context.personality, - collaboration_mode: Some(turn_context.collaboration_mode.clone()), + collaboration_mode: Some(turn_context.collaboration_mode()), multi_agent_version: None, multi_agent_mode: None, realtime_active: Some(turn_context.realtime_active), @@ -1262,7 +1262,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis model: previous_model.to_string(), comp_hash: None, personality: turn_context.personality, - collaboration_mode: Some(turn_context.collaboration_mode.clone()), + collaboration_mode: Some(turn_context.collaboration_mode()), multi_agent_version: None, multi_agent_mode: None, realtime_active: Some(turn_context.realtime_active), @@ -1349,7 +1349,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis model: previous_model.to_string(), comp_hash: None, personality: turn_context.personality, - collaboration_mode: Some(turn_context.collaboration_mode.clone()), + collaboration_mode: Some(turn_context.collaboration_mode()), multi_agent_version: None, multi_agent_mode: None, realtime_active: Some(turn_context.realtime_active), @@ -1381,7 +1381,7 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu model: previous_model.to_string(), comp_hash: None, personality: turn_context.personality, - collaboration_mode: Some(turn_context.collaboration_mode.clone()), + collaboration_mode: Some(turn_context.collaboration_mode()), multi_agent_version: None, multi_agent_mode: None, realtime_active: Some(turn_context.realtime_active), @@ -1508,7 +1508,7 @@ async fn record_initial_history_resumed_unmatched_abort_preserves_active_turn_fo model: current_model.to_string(), comp_hash: None, personality: turn_context.personality, - collaboration_mode: Some(turn_context.collaboration_mode.clone()), + collaboration_mode: Some(turn_context.collaboration_mode()), multi_agent_version: None, multi_agent_mode: None, realtime_active: Some(turn_context.realtime_active), @@ -1630,7 +1630,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea model: previous_model.to_string(), comp_hash: None, personality: turn_context.personality, - collaboration_mode: Some(turn_context.collaboration_mode.clone()), + collaboration_mode: Some(turn_context.collaboration_mode()), multi_agent_version: None, multi_agent_mode: None, realtime_active: Some(turn_context.realtime_active), @@ -1798,7 +1798,7 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear model: previous_model.to_string(), comp_hash: None, personality: turn_context.personality, - collaboration_mode: Some(turn_context.collaboration_mode.clone()), + collaboration_mode: Some(turn_context.collaboration_mode()), multi_agent_version: None, multi_agent_mode: None, realtime_active: Some(turn_context.realtime_active), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 696a8adbd8..c1f36ad2b6 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2530,7 +2530,7 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() { thread_level_id: session.thread_id.to_string(), turn_level_id: turn_context.sub_id.clone(), turn_id: turn_context.sub_id.clone(), - collaboration_mode: turn_context.collaboration_mode.clone(), + collaboration_mode: turn_context.collaboration_mode(), token_usage_at_turn_start, saw_session_store: true, saw_thread_store: true, @@ -3097,7 +3097,7 @@ async fn record_initial_history_forked_hydrates_previous_turn_settings() { model: previous_model.to_string(), comp_hash: None, personality: turn_context.personality, - collaboration_mode: Some(turn_context.collaboration_mode.clone()), + collaboration_mode: Some(turn_context.collaboration_mode()), multi_agent_version: None, multi_agent_mode: None, realtime_active: Some(turn_context.realtime_active), @@ -3991,14 +3991,14 @@ async fn turn_context_with_model_updates_model_fields() { .await; assert_eq!(updated.config.model.as_deref(), Some("gpt-5.4")); - assert_eq!(updated.collaboration_mode.model(), "gpt-5.4"); + assert_eq!(updated.collaboration_mode().model(), "gpt-5.4"); assert_eq!(updated.model_info, expected_model_info); assert_eq!( updated.reasoning_effort, Some(ReasoningEffortConfig::Medium) ); assert_eq!( - updated.collaboration_mode.reasoning_effort(), + updated.collaboration_mode().reasoning_effort(), Some(ReasoningEffortConfig::Medium) ); assert_eq!( diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index bafb368b32..8811e2f0aa 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -787,7 +787,7 @@ async fn track_turn_resolved_config_analytics( approval_policy: turn_context.approval_policy.value(), approvals_reviewer: turn_context.config.approvals_reviewer, sandbox_network_access: turn_context.network_sandbox_policy().is_enabled(), - collaboration_mode: turn_context.collaboration_mode.mode, + collaboration_mode: turn_context.mode, personality: turn_context.personality, workspace_kind: turn_context.turn_metadata_state.workspace_kind(), is_first_turn, @@ -1994,7 +1994,7 @@ async fn try_run_sampling_request( let mut should_emit_turn_diff = false; let mut should_emit_token_count = false; let reasoning_effort = turn_context.effective_reasoning_effort_for_tracing(); - let plan_mode = turn_context.collaboration_mode.mode == ModeKind::Plan; + let plan_mode = turn_context.mode == ModeKind::Plan; let mut assistant_message_stream_parsers = AssistantMessageStreamParsers::new(plan_mode); let mut plan_mode_state = plan_mode.then(|| PlanModeStreamState::new(&turn_context.sub_id)); let defer_streamed_turn_items_for_contributors = diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index e200413d3d..690e7096cf 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -125,7 +125,8 @@ pub struct TurnContext { pub(crate) timezone: Option, pub(crate) app_server_client_name: Option, pub(crate) developer_instructions: Option, - pub(crate) collaboration_mode: CollaborationMode, + pub(crate) mode: ModeKind, + pub(crate) collaboration_mode_developer_instructions: Option, pub(crate) multi_agent_version: MultiAgentVersion, pub(crate) personality: Option, pub(crate) approval_policy: Constrained, @@ -156,6 +157,17 @@ impl TurnContext { || matches!(self.history_mode, ThreadHistoryMode::Paginated) } + pub(crate) fn collaboration_mode(&self) -> CollaborationMode { + CollaborationMode { + mode: self.mode, + settings: Settings { + model: self.model_info.slug.clone(), + reasoning_effort: self.reasoning_effort.clone(), + developer_instructions: self.collaboration_mode_developer_instructions.clone(), + }, + } + } + pub(crate) fn permission_profile(&self) -> PermissionProfile { self.permission_profile.clone() } @@ -245,11 +257,6 @@ impl TurnContext { }; config.model_reasoning_effort = reasoning_effort.clone(); - let collaboration_mode = self.collaboration_mode.with_updates( - Some(model.clone()), - Some(reasoning_effort.clone()), - /*developer_instructions*/ None, - ); let available_models = models_manager .list_models( RefreshStrategy::OnlineIfUncached, @@ -282,7 +289,10 @@ impl TurnContext { timezone: self.timezone.clone(), app_server_client_name: self.app_server_client_name.clone(), developer_instructions: self.developer_instructions.clone(), - collaboration_mode, + mode: self.mode, + collaboration_mode_developer_instructions: self + .collaboration_mode_developer_instructions + .clone(), multi_agent_version: self.multi_agent_version, personality: self.personality, approval_policy: self.approval_policy.clone(), @@ -380,7 +390,7 @@ impl TurnContext { model: self.model_info.slug.clone(), comp_hash: self.model_info.comp_hash.clone(), personality: self.personality, - collaboration_mode: Some(self.collaboration_mode.clone()), + collaboration_mode: Some(self.collaboration_mode()), multi_agent_version: Some(self.multi_agent_version), multi_agent_mode: super::multi_agents::effective_multi_agent_mode(self), realtime_active: Some(self.realtime_active), @@ -497,7 +507,8 @@ impl Session { sub_id: String, skills_snapshot: HostSkillsSnapshot, ) -> TurnContext { - let reasoning_effort = session_configuration.collaboration_mode.reasoning_effort(); + let collaboration_mode = &session_configuration.collaboration_mode; + let reasoning_effort = collaboration_mode.reasoning_effort(); let reasoning_summary = session_configuration .model_reasoning_summary .unwrap_or(model_info.default_reasoning_summary); @@ -562,7 +573,11 @@ impl Session { timezone: Some(timezone), app_server_client_name: session_configuration.app_server_client_name.clone(), developer_instructions: session_configuration.developer_instructions.clone(), - collaboration_mode: session_configuration.collaboration_mode.clone(), + mode: collaboration_mode.mode, + collaboration_mode_developer_instructions: collaboration_mode + .settings + .developer_instructions + .clone(), multi_agent_version, personality: session_configuration.personality, approval_policy: session_configuration.approval_policy.clone(), diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 7d03a0d271..c67bf7ef29 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -130,7 +130,7 @@ pub(crate) async fn record_completed_response_item_with_finalized_facts( || { completed_item_defers_mailbox_delivery_to_next_turn( item, - turn_context.collaboration_mode.mode == ModeKind::Plan, + turn_context.mode == ModeKind::Plan, ) }, |facts| facts.defers_mailbox_delivery_to_next_turn, @@ -321,7 +321,7 @@ pub(crate) async fn handle_output_item_done( previously_active_item: Option, ) -> Result { let mut output = OutputItemResult::default(); - let plan_mode = ctx.turn_context.collaboration_mode.mode == ModeKind::Plan; + let plan_mode = ctx.turn_context.mode == ModeKind::Plan; match ToolRouter::build_tool_call(item.clone()) { // The model emitted a tool call; log it, persist the item immediately, and queue the tool execution. diff --git a/codex-rs/core/src/tasks/lifecycle.rs b/codex-rs/core/src/tasks/lifecycle.rs index 782cc8c7de..789955df13 100644 --- a/codex-rs/core/src/tasks/lifecycle.rs +++ b/codex-rs/core/src/tasks/lifecycle.rs @@ -12,11 +12,12 @@ impl Session { turn_context: &TurnContext, token_usage_at_turn_start: &TokenUsage, ) { + let collaboration_mode = turn_context.collaboration_mode(); for contributor in self.services.extensions.turn_lifecycle_contributors() { contributor .on_turn_start(codex_extension_api::TurnStartInput { turn_id: turn_context.sub_id.as_str(), - collaboration_mode: &turn_context.collaboration_mode, + collaboration_mode: &collaboration_mode, token_usage_at_turn_start, session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 40837728a2..66c4d8d5c7 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -52,7 +52,7 @@ impl SessionTask for RegularTask { trace_id: ctx.trace_id.clone(), started_at: ctx.turn_timing_state.started_at_unix_secs().await, model_context_window: ctx.model_context_window(), - collaboration_mode_kind: ctx.collaboration_mode.mode, + collaboration_mode_kind: ctx.mode, }); sess.send_event(ctx.as_ref(), event).await; sess.set_server_reasoning_included(/*included*/ false).await; diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index a15c53e16b..ddf664c2c4 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -120,7 +120,7 @@ pub(crate) async fn execute_user_shell_command( trace_id: turn_context.trace_id.clone(), started_at: turn_context.turn_timing_state.started_at_unix_secs().await, model_context_window: turn_context.model_context_window(), - collaboration_mode_kind: turn_context.collaboration_mode.mode, + collaboration_mode_kind: turn_context.mode, }); session.send_event(turn_context.as_ref(), event).await; } diff --git a/codex-rs/core/src/tools/handlers/plan.rs b/codex-rs/core/src/tools/handlers/plan.rs index 1b30c281ce..33b91aa512 100644 --- a/codex-rs/core/src/tools/handlers/plan.rs +++ b/codex-rs/core/src/tools/handlers/plan.rs @@ -81,7 +81,7 @@ impl PlanHandler { } }; - if turn.collaboration_mode.mode == ModeKind::Plan { + if turn.mode == ModeKind::Plan { return Err(FunctionCallError::RespondToModel( "update_plan is a TODO/checklist tool and is not allowed in Plan mode".to_string(), ));