diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 3ba06cba4..0e1e49ba7 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -660,10 +660,7 @@ impl TurnContext { .unwrap_or(compact::SUMMARIZATION_PROMPT) } - pub(crate) fn to_turn_context_item( - &self, - collaboration_mode: CollaborationMode, - ) -> TurnContextItem { + pub(crate) fn to_turn_context_item(&self) -> TurnContextItem { TurnContextItem { turn_id: Some(self.sub_id.clone()), cwd: self.cwd.clone(), @@ -672,7 +669,7 @@ impl TurnContext { network: self.turn_context_network_item(), model: self.model_info.slug.clone(), personality: self.personality, - collaboration_mode: Some(collaboration_mode), + collaboration_mode: Some(self.collaboration_mode.clone()), effort: self.reasoning_effort, summary: self.reasoning_summary, user_instructions: self.user_instructions.clone(), @@ -1995,11 +1992,6 @@ impl Session { .await } - pub(crate) async fn current_collaboration_mode(&self) -> CollaborationMode { - let state = self.state.lock().await; - state.session_configuration.collaboration_mode.clone() - } - pub(crate) fn is_model_switch_developer_message(item: &ResponseItem) -> bool { let ResponseItem::Message { role, content, .. } = item else { return false; @@ -2014,7 +2006,7 @@ impl Session { } fn build_settings_update_items( &self, - previous_context: Option<&Arc>, + previous_context: Option<&TurnContextItem>, resumed_model: Option<&str>, current_context: &TurnContext, ) -> Vec { @@ -2025,7 +2017,7 @@ impl Session { let shell = self.user_shell(); let exec_policy = self.services.exec_policy.current(); crate::context_manager::updates::build_settings_update_items( - previous_context.map(Arc::as_ref), + previous_context, resumed_model, current_context, shell.as_ref(), @@ -2660,6 +2652,16 @@ impl Session { state.clone_history() } + pub(crate) async fn previous_context_item(&self) -> Option { + let state = self.state.lock().await; + state.previous_context_item() + } + + pub(crate) async fn set_previous_context_item(&self, item: Option) { + let mut state = self.state.lock().await; + state.set_previous_context_item(item); + } + pub(crate) async fn update_token_usage_info( &self, turn_context: &TurnContext, @@ -3123,7 +3125,9 @@ impl Session { async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiver) { // Seed with context in case there is an OverrideTurnContext first. - let mut previous_context: Option> = Some(sess.new_default_turn().await); + let initial_context = sess.new_default_turn().await; + sess.set_previous_context_item(Some(initial_context.to_turn_context_item())) + .await; // To break out of this loop, send Op::Shutdown. while let Ok(sub) = rx_sub.recv().await { @@ -3173,8 +3177,7 @@ async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiv .await; } Op::UserInput { .. } | Op::UserTurn { .. } => { - handlers::user_input_or_turn(&sess, sub.id.clone(), sub.op, &mut previous_context) - .await; + handlers::user_input_or_turn(&sess, sub.id.clone(), sub.op).await; } Op::ExecApproval { id: approval_id, @@ -3251,13 +3254,7 @@ async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiv handlers::set_thread_name(&sess, sub.id.clone(), name).await; } Op::RunUserShellCommand { command } => { - handlers::run_user_shell_command( - &sess, - sub.id.clone(), - command, - &mut previous_context, - ) - .await; + handlers::run_user_shell_command(&sess, sub.id.clone(), command).await; } Op::ResolveElicitation { server_name, @@ -3285,7 +3282,6 @@ mod handlers { use crate::codex::Session; use crate::codex::SessionSettingsUpdate; use crate::codex::SteerInputError; - use crate::codex::TurnContext; use crate::codex::spawn_review_thread; use crate::config::Config; @@ -3362,12 +3358,7 @@ mod handlers { } } - pub async fn user_input_or_turn( - sess: &Arc, - sub_id: String, - op: Op, - previous_context: &mut Option>, - ) { + pub async fn user_input_or_turn(sess: &Arc, sub_id: String, op: Op) { let (items, updates) = match op { Op::UserTurn { cwd, @@ -3430,8 +3421,9 @@ mod handlers { if let Err(SteerInputError::NoActiveTurn(items)) = sess.steer_input(items, None).await { sess.seed_initial_context_if_needed(¤t_context).await; let previous_model = sess.previous_model().await; + let previous_context_item = sess.previous_context_item().await; let update_items = sess.build_settings_update_items( - previous_context.as_ref(), + previous_context_item.as_ref(), previous_model.as_deref(), ¤t_context, ); @@ -3445,16 +3437,12 @@ mod handlers { let regular_task = sess.take_startup_regular_task().await.unwrap_or_default(); sess.spawn_task(Arc::clone(¤t_context), items, regular_task) .await; - *previous_context = Some(current_context); + sess.set_previous_context_item(Some(current_context.to_turn_context_item())) + .await; } } - pub async fn run_user_shell_command( - sess: &Arc, - sub_id: String, - command: String, - previous_context: &mut Option>, - ) { + pub async fn run_user_shell_command(sess: &Arc, sub_id: String, command: String) { if let Some((turn_context, cancellation_token)) = sess.active_turn_context_and_cancellation_token().await { @@ -3479,7 +3467,8 @@ mod handlers { UserShellCommandTask::new(command), ) .await; - *previous_context = Some(turn_context); + sess.set_previous_context_item(Some(turn_context.to_turn_context_item())) + .await; } pub async fn resolve_elicitation( @@ -5465,9 +5454,7 @@ async fn try_run_sampling_request( prompt: &Prompt, cancellation_token: CancellationToken, ) -> CodexResult { - let collaboration_mode = sess.current_collaboration_mode().await; - let rollout_item = - RolloutItem::TurnContext(turn_context.to_turn_context_item(collaboration_mode)); + let rollout_item = RolloutItem::TurnContext(turn_context.to_turn_context_item()); feedback_tags!( model = turn_context.model_info.slug.clone(), @@ -7738,8 +7725,12 @@ mod tests { .expect("rebuild config layer stack with network requirements"); current_context.config = Arc::new(config); - let update_items = - session.build_settings_update_items(Some(&previous_context), None, ¤t_context); + let previous_context_item = previous_context.to_turn_context_item(); + let update_items = session.build_settings_update_items( + Some(&previous_context_item), + None, + ¤t_context, + ); let environment_update = update_items .iter() diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index a40f1a0b8..406525142 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -118,13 +118,7 @@ async fn run_compact_task_inner( // Reuse one client session so turn-scoped state (sticky routing, websocket append tracking) // survives retries within this compact turn. - // TODO: If we need to guarantee the persisted mode always matches the prompt used for this - // turn, capture it in TurnContext at creation time. Using SessionConfiguration here avoids - // duplicating model settings on TurnContext, but an Op after turn start could update the - // session config before this write occurs. - let collaboration_mode = sess.current_collaboration_mode().await; - let rollout_item = - RolloutItem::TurnContext(turn_context.to_turn_context_item(collaboration_mode)); + let rollout_item = RolloutItem::TurnContext(turn_context.to_turn_context_item()); sess.persist_rollout_items(&[rollout_item]).await; loop { diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index b6b3ceae5..0cd935523 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -18,6 +18,7 @@ use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::InputModality; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; +use codex_protocol::protocol::TurnContextItem; use std::ops::Deref; /// Transcript of thread history @@ -26,6 +27,12 @@ pub(crate) struct ContextManager { /// The oldest items are at the beginning of the vector. items: Vec, token_info: Option, + /// Previous turn context snapshot used for diffing context and producing + /// model-visible settings update items. + /// + /// When this is `None`, settings diffing treats the next turn as having no + /// baseline and emits a full reinjection of context state. + previous_context_item: Option, } #[derive(Debug, Clone, Copy, Default)] @@ -41,6 +48,7 @@ impl ContextManager { Self { items: Vec::new(), token_info: TokenUsageInfo::new_or_append(&None, &None, None), + previous_context_item: None, } } @@ -52,6 +60,14 @@ impl ContextManager { self.token_info = info; } + pub(crate) fn set_previous_context_item(&mut self, item: Option) { + self.previous_context_item = item; + } + + pub(crate) fn previous_context_item(&self) -> Option { + self.previous_context_item.clone() + } + pub(crate) fn set_token_usage_full(&mut self, context_window: i64) { match &mut self.token_info { Some(info) => info.fill_to_context_window(context_window), diff --git a/codex-rs/core/src/context_manager/updates.rs b/codex-rs/core/src/context_manager/updates.rs index 4a3f0c15e..217889a72 100644 --- a/codex-rs/core/src/context_manager/updates.rs +++ b/codex-rs/core/src/context_manager/updates.rs @@ -6,26 +6,27 @@ use codex_protocol::config_types::Personality; use codex_protocol::models::DeveloperInstructions; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; +use codex_protocol::protocol::TurnContextItem; fn build_environment_update_item( - previous: Option<&TurnContext>, + previous: Option<&TurnContextItem>, next: &TurnContext, shell: &Shell, ) -> Option { let prev = previous?; - let prev_context = EnvironmentContext::from_turn_context(prev, shell); + let prev_context = EnvironmentContext::from_turn_context_item(prev, shell); let next_context = EnvironmentContext::from_turn_context(next, shell); if prev_context.equals_except_shell(&next_context) { return None; } - Some(ResponseItem::from(EnvironmentContext::diff( - prev, next, shell, - ))) + Some(ResponseItem::from( + EnvironmentContext::diff_from_turn_context_item(prev, next, shell), + )) } fn build_permissions_update_item( - previous: Option<&TurnContext>, + previous: Option<&TurnContextItem>, next: &TurnContext, exec_policy: &Policy, ) -> Option { @@ -46,11 +47,11 @@ fn build_permissions_update_item( } fn build_collaboration_mode_update_item( - previous: Option<&TurnContext>, + previous: Option<&TurnContextItem>, next: &TurnContext, ) -> Option { let prev = previous?; - if prev.collaboration_mode != next.collaboration_mode { + if prev.collaboration_mode.as_ref() != Some(&next.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(DeveloperInstructions::from_collaboration_mode(&next.collaboration_mode)?.into()) @@ -60,7 +61,7 @@ fn build_collaboration_mode_update_item( } fn build_personality_update_item( - previous: Option<&TurnContext>, + previous: Option<&TurnContextItem>, next: &TurnContext, personality_feature_enabled: bool, ) -> Option { @@ -68,7 +69,7 @@ fn build_personality_update_item( return None; } let previous = previous?; - if next.model_info.slug != previous.model_info.slug { + if next.model_info.slug != previous.model { return None; } @@ -96,12 +97,11 @@ pub(crate) fn personality_message_for( } pub(crate) fn build_model_instructions_update_item( - previous: Option<&TurnContext>, + previous: Option<&TurnContextItem>, resumed_model: Option<&str>, next: &TurnContext, ) -> Option { - let previous_model = - resumed_model.or_else(|| previous.map(|prev| prev.model_info.slug.as_str()))?; + let previous_model = resumed_model.or_else(|| previous.map(|prev| prev.model.as_str()))?; if previous_model == next.model_info.slug { return None; } @@ -115,7 +115,7 @@ pub(crate) fn build_model_instructions_update_item( } pub(crate) fn build_settings_update_items( - previous: Option<&TurnContext>, + previous: Option<&TurnContextItem>, resumed_model: Option<&str>, next: &TurnContext, shell: &Shell, diff --git a/codex-rs/core/src/environment_context.rs b/codex-rs/core/src/environment_context.rs index e6995b7a9..3c6df40ad 100644 --- a/codex-rs/core/src/environment_context.rs +++ b/codex-rs/core/src/environment_context.rs @@ -4,6 +4,8 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::ENVIRONMENT_CONTEXT_CLOSE_TAG; use codex_protocol::protocol::ENVIRONMENT_CONTEXT_OPEN_TAG; +use codex_protocol::protocol::TurnContextItem; +use codex_protocol::protocol::TurnContextNetworkItem; use serde::Deserialize; use serde::Serialize; use std::path::PathBuf; @@ -44,8 +46,12 @@ impl EnvironmentContext { self.cwd == *cwd && self.network == *network } - pub fn diff(before: &TurnContext, after: &TurnContext, shell: &Shell) -> Self { - let before_network = Self::network_from_turn_context(before); + pub fn diff_from_turn_context_item( + before: &TurnContextItem, + after: &TurnContext, + shell: &Shell, + ) -> Self { + let before_network = Self::network_from_turn_context_item(before); let after_network = Self::network_from_turn_context(after); let cwd = if before.cwd != after.cwd { Some(after.cwd.clone()) @@ -68,6 +74,14 @@ impl EnvironmentContext { ) } + pub fn from_turn_context_item(turn_context_item: &TurnContextItem, shell: &Shell) -> Self { + Self::new( + Some(turn_context_item.cwd.clone()), + shell.clone(), + Self::network_from_turn_context_item(turn_context_item), + ) + } + fn network_from_turn_context(turn_context: &TurnContext) -> Option { let network = turn_context .config @@ -81,6 +95,19 @@ impl EnvironmentContext { denied_domains: network.denied_domains.clone().unwrap_or_default(), }) } + + fn network_from_turn_context_item( + turn_context_item: &TurnContextItem, + ) -> Option { + let TurnContextNetworkItem { + allowed_domains, + denied_domains, + } = turn_context_item.network.as_ref()?; + Some(NetworkContext { + allowed_domains: allowed_domains.clone(), + denied_domains: denied_domains.clone(), + }) + } } impl EnvironmentContext { diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index c86f0cdfd..76c365ab5 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -11,6 +11,7 @@ use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; use crate::tasks::RegularTask; use crate::truncate::TruncationPolicy; +use codex_protocol::protocol::TurnContextItem; /// Persistent, session-scoped state previously stored directly on `Session`. pub(crate) struct SessionState { @@ -80,6 +81,14 @@ impl SessionState { self.history.set_token_info(info); } + pub(crate) fn set_previous_context_item(&mut self, item: Option) { + self.history.set_previous_context_item(item); + } + + pub(crate) fn previous_context_item(&self) -> Option { + self.history.previous_context_item() + } + // Token/rate limit helpers pub(crate) fn update_token_info_from_usage( &mut self,