diff --git a/codex-rs/core/src/context/contextual_user_message.rs b/codex-rs/core/src/context/contextual_user_message.rs index de735b6f1ee7..514deb8937ec 100644 --- a/codex-rs/core/src/context/contextual_user_message.rs +++ b/codex-rs/core/src/context/contextual_user_message.rs @@ -3,7 +3,7 @@ use codex_protocol::items::parse_hook_prompt_fragment; use codex_protocol::models::ContentItem; use super::AdditionalContextUserFragment; -use super::EnvironmentContext; +use super::EnvironmentsState; use super::FragmentRegistration; use super::FragmentRegistrationProxy; use super::InternalModelContextFragment; @@ -19,7 +19,7 @@ use super::UserShellCommand; static USER_INSTRUCTIONS_REGISTRATION: FragmentRegistrationProxy = FragmentRegistrationProxy::new(); -static ENVIRONMENT_CONTEXT_REGISTRATION: FragmentRegistrationProxy = +static ENVIRONMENT_CONTEXT_REGISTRATION: FragmentRegistrationProxy = FragmentRegistrationProxy::new(); static ADDITIONAL_CONTEXT_REGISTRATION: FragmentRegistrationProxy = FragmentRegistrationProxy::new(); diff --git a/codex-rs/core/src/context/environment_context.rs b/codex-rs/core/src/context/environment_context.rs index e93ffb88ffda..6cf665909958 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -1,102 +1,16 @@ -use crate::session::turn_context::TurnContext; -use crate::session::turn_context::TurnEnvironment; -use crate::shell::Shell; use codex_protocol::models::ManagedFileSystemPermissions; use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; use codex_protocol::permissions::FileSystemSandboxEntry; use codex_protocol::permissions::FileSystemSpecialPath; -use codex_protocol::protocol::TurnContextItem; -use codex_protocol::protocol::TurnContextNetworkItem; use codex_utils_absolute_path::AbsolutePathBuf; -use codex_utils_path_uri::PathUri; -use std::collections::HashSet; +use indexmap::IndexMap; use std::path::PathBuf; -use super::ContextualUserFragment; - -#[derive(Debug, Clone, PartialEq)] -pub(crate) struct EnvironmentContext { - pub(crate) environments: EnvironmentContextEnvironments, - pub(crate) current_date: Option, - pub(crate) timezone: Option, - pub(crate) network: Option, - pub(crate) filesystem: Option, - pub(crate) subagents: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct EnvironmentContextEnvironment { - pub(crate) id: String, - pub(crate) cwd: PathUri, - pub(crate) shell: String, -} - -impl EnvironmentContextEnvironment { - fn legacy(cwd: PathUri, shell: String) -> Self { - Self { - id: String::new(), - cwd, - shell, - } - } - - fn from_turn_environments(environments: &[TurnEnvironment], shell: &Shell) -> Vec { - environments - .iter() - .map(|environment| Self { - id: environment.environment_id.clone(), - cwd: environment.cwd().clone(), - shell: environment - .shell - .as_ref() - .map(|shell| shell.name().to_string()) - .unwrap_or_else(|| shell.name().to_string()), - }) - .collect() - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) enum EnvironmentContextEnvironments { - None, - Single(EnvironmentContextEnvironment), - Multiple(Vec), -} - -impl EnvironmentContextEnvironments { - fn from_vec(environments: Vec) -> Self { - let mut environments = environments; - match environments.pop() { - None => Self::None, - Some(environment) if environments.is_empty() => Self::Single(environment), - Some(environment) => { - environments.push(environment); - Self::Multiple(environments) - } - } - } - - fn equals_except_shell(&self, other: &Self) -> bool { - match (self, other) { - (Self::None, Self::None) => true, - (Self::Single(left), Self::Single(right)) => left.cwd == right.cwd, - (Self::Multiple(left), Self::Multiple(right)) => { - left.len() == right.len() - && left - .iter() - .zip(right.iter()) - .all(|(left, right)| left.id == right.id && left.cwd == right.cwd) - } - _ => false, - } - } -} - #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct FileSystemContext { - workspace_roots: Vec, + workspace_roots: IndexMap, permission_profile: FileSystemPermissionProfileContext, } @@ -110,14 +24,14 @@ enum FileSystemPermissionProfileContext { #[derive(Debug, Clone, PartialEq, Eq)] enum ManagedFileSystemContext { Restricted { - entries: Vec, + entries: IndexMap, glob_scan_max_depth: Option, }, Unrestricted, } impl FileSystemContext { - fn from_permission_profile( + pub(super) fn from_permission_profile( permission_profile: &PermissionProfile, workspace_roots: &[AbsolutePathBuf], ) -> Self { @@ -126,7 +40,7 @@ impl FileSystemContext { .materialize_project_roots_with_workspace_roots(workspace_roots); let workspace_roots = workspace_roots .iter() - .map(|root| root.to_string_lossy().into_owned()) + .map(|root| (root.to_string_lossy().into_owned(), true)) .collect(); let permission_profile = match permission_profile { PermissionProfile::Managed { file_system, .. } => { @@ -143,11 +57,11 @@ impl FileSystemContext { } } - fn render(&self) -> String { + pub(super) fn render(&self) -> String { let mut rendered = "".to_string(); if !self.workspace_roots.is_empty() { rendered.push_str(""); - for root in &self.workspace_roots { + for root in self.workspace_roots.keys() { push_text_element(&mut rendered, "root", root); } rendered.push_str(""); @@ -162,15 +76,15 @@ impl From for ManagedFileSystemContext { fn from(file_system: ManagedFileSystemPermissions) -> Self { match file_system { ManagedFileSystemPermissions::Restricted { - mut entries, + entries, glob_scan_max_depth, - } => { - dedupe_file_system_entries(&mut entries); - Self::Restricted { - entries, - glob_scan_max_depth: glob_scan_max_depth.map(usize::from), - } - } + } => Self::Restricted { + entries: entries + .into_iter() + .map(|entry| (file_system_entry_key(&entry), entry)) + .collect(), + glob_scan_max_depth: glob_scan_max_depth.map(usize::from), + }, ManagedFileSystemPermissions::Unrestricted => Self::Unrestricted, } } @@ -215,7 +129,7 @@ impl ManagedFileSystemContext { rendered.push_str(&format!(" glob_scan_max_depth=\"{glob_scan_max_depth}\"")); } rendered.push('>'); - for entry in entries { + for entry in entries.values() { render_file_system_entry(rendered, entry); } rendered.push_str(""); @@ -272,9 +186,13 @@ fn render_special_path_with_subpath(base: &str, subpath: &Option) -> St } } -fn dedupe_file_system_entries(entries: &mut Vec) { - let mut seen = HashSet::new(); - entries.retain(|entry| seen.insert(entry.clone())); +fn file_system_entry_key(entry: &FileSystemSandboxEntry) -> String { + let path = match &entry.path { + FileSystemPath::Path { path } => format!("path:{}", path.to_string_lossy()), + FileSystemPath::GlobPattern { pattern } => format!("glob:{pattern}"), + FileSystemPath::Special { value } => format!("special:{}", render_special_path(value)), + }; + format!("{}:{path}", entry.access) } fn push_text_element(rendered: &mut String, name: &str, value: &str) { @@ -283,7 +201,7 @@ fn push_text_element(rendered: &mut String, name: &str, value: &str) { rendered.push_str(&format!("")); } -fn push_xml_escaped_text(rendered: &mut String, value: &str) { +pub(crate) fn push_xml_escaped_text(rendered: &mut String, value: &str) { for ch in value.chars() { match ch { '&' => rendered.push_str("&"), @@ -298,288 +216,58 @@ fn push_xml_escaped_text(rendered: &mut String, value: &str) { #[derive(Debug, Clone, PartialEq, Eq, Default)] pub(crate) struct NetworkContext { - allowed_domains: Vec, - denied_domains: Vec, + domains: IndexMap, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum NetworkDomainAccess { + Allow, + Deny, } impl NetworkContext { pub(crate) fn new(allowed_domains: Vec, denied_domains: Vec) -> Self { Self { - allowed_domains, - denied_domains, + domains: allowed_domains + .into_iter() + .map(|domain| (domain, NetworkDomainAccess::Allow)) + .chain( + denied_domains + .into_iter() + .map(|domain| (domain, NetworkDomainAccess::Deny)), + ) + .collect(), } } - fn render(&self) -> String { + pub(super) fn render(&self) -> String { let mut rendered = "".to_string(); - Self::push_rendered_domain_element(&mut rendered, "allowed", &self.allowed_domains); - Self::push_rendered_domain_element(&mut rendered, "denied", &self.denied_domains); + self.push_rendered_domain_element(&mut rendered, "allowed", NetworkDomainAccess::Allow); + self.push_rendered_domain_element(&mut rendered, "denied", NetworkDomainAccess::Deny); rendered.push_str(""); rendered } - fn push_rendered_domain_element(rendered_network: &mut String, name: &str, domains: &[String]) { - if domains.is_empty() { + fn push_rendered_domain_element( + &self, + rendered_network: &mut String, + name: &str, + access: NetworkDomainAccess, + ) { + let mut domains = self + .domains + .iter() + .filter_map(|(domain, current_access)| (*current_access == access).then_some(domain)); + let Some(first) = domains.next() else { return; - } - - rendered_network.push_str(&format!("<{name}>")); - rendered_network.push_str(&domains.join(",")); - rendered_network.push_str(&format!("")); - } -} - -impl EnvironmentContext { - pub(crate) fn new( - environments: Vec, - current_date: Option, - timezone: Option, - network: Option, - subagents: Option, - ) -> Self { - Self { - environments: EnvironmentContextEnvironments::from_vec(environments), - current_date, - timezone, - network, - filesystem: None, - subagents, - } - } - - fn new_with_environments( - environments: EnvironmentContextEnvironments, - current_date: Option, - timezone: Option, - network: Option, - filesystem: Option, - subagents: Option, - ) -> Self { - Self { - environments, - current_date, - timezone, - network, - filesystem, - subagents, - } - } - - /// Compares two environment contexts, ignoring the shell. Useful when - /// comparing turn to turn, since the initial environment_context will - /// include the shell, and then it is not configurable from turn to turn. - pub(crate) fn equals_except_shell(&self, other: &EnvironmentContext) -> bool { - self.environments.equals_except_shell(&other.environments) - && self.current_date == other.current_date - && self.timezone == other.timezone - && self.network == other.network - && self.filesystem == other.filesystem - && self.subagents == other.subagents - } - - pub(crate) fn diff_from_turn_context_item( - before: &TurnContextItem, - after: &EnvironmentContext, - ) -> Self { - let before_network = Self::network_from_turn_context_item(before); - let before_filesystem = Self::filesystem_from_turn_context_item(before); - let environments = match &after.environments { - EnvironmentContextEnvironments::Single(environment) => { - if PathUri::from_abs_path(&before.cwd) != environment.cwd { - EnvironmentContextEnvironments::Single(EnvironmentContextEnvironment::legacy( - environment.cwd.clone(), - environment.shell.clone(), - )) - } else { - EnvironmentContextEnvironments::None - } - } - EnvironmentContextEnvironments::Multiple(environments) => { - EnvironmentContextEnvironments::Multiple(environments.clone()) - } - EnvironmentContextEnvironments::None => EnvironmentContextEnvironments::None, }; - let network = if before_network != after.network { - after.network.clone() - } else { - before_network - }; - let filesystem = if before_filesystem != after.filesystem { - after.filesystem.clone() - } else { - before_filesystem - }; - EnvironmentContext::new_with_environments( - environments, - after.current_date.clone(), - after.timezone.clone(), - network, - filesystem, - /*subagents*/ None, - ) - } - pub(crate) fn from_turn_context(turn_context: &TurnContext, shell: &Shell) -> Self { - let mut context = Self::new( - EnvironmentContextEnvironment::from_turn_environments( - &turn_context.environments.turn_environments, - shell, - ), - turn_context.current_date.clone(), - turn_context.timezone.clone(), - Self::network_from_turn_context(turn_context), - /*subagents*/ None, - ); - context.filesystem = Some(FileSystemContext::from_permission_profile( - &turn_context.permission_profile, - &turn_context.config.effective_workspace_roots(), - )); - context - } - - pub(crate) fn from_turn_context_item( - turn_context_item: &TurnContextItem, - shell: String, - ) -> Self { - Self::new_with_environments( - EnvironmentContextEnvironments::from_vec(vec![EnvironmentContextEnvironment::legacy( - PathUri::from_abs_path(&turn_context_item.cwd), - shell, - )]), - turn_context_item.current_date.clone(), - turn_context_item.timezone.clone(), - Self::network_from_turn_context_item(turn_context_item), - Self::filesystem_from_turn_context_item(turn_context_item), - /*subagents*/ None, - ) - } - - pub(crate) fn with_subagents(mut self, subagents: String) -> Self { - if !subagents.is_empty() { - self.subagents = Some(subagents); - } - self - } - - fn network_from_turn_context(turn_context: &TurnContext) -> Option { - let network = turn_context - .config - .config_layer_stack - .requirements() - .network - .as_ref()?; - - Some(NetworkContext::new( - network - .domains - .as_ref() - .and_then(codex_config::NetworkDomainPermissionsToml::allowed_domains) - .unwrap_or_default(), - network - .domains - .as_ref() - .and_then(codex_config::NetworkDomainPermissionsToml::denied_domains) - .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::new( - allowed_domains.clone(), - denied_domains.clone(), - )) - } - - fn filesystem_from_turn_context_item( - turn_context_item: &TurnContextItem, - ) -> Option { - Some(FileSystemContext::from_permission_profile( - &turn_context_item.permission_profile(), - &workspace_roots_from_turn_context_item(turn_context_item), - )) - } -} - -fn workspace_roots_from_turn_context_item( - turn_context_item: &TurnContextItem, -) -> Vec { - if let Some(workspace_roots) = turn_context_item.workspace_roots.as_ref() { - return workspace_roots.clone(); - } - - vec![turn_context_item.cwd.clone()] -} - -impl ContextualUserFragment for EnvironmentContext { - fn role(&self) -> &'static str { - "user" - } - - fn markers(&self) -> (&'static str, &'static str) { - Self::type_markers() - } - - fn type_markers() -> (&'static str, &'static str) { - ( - codex_protocol::protocol::ENVIRONMENT_CONTEXT_OPEN_TAG, - codex_protocol::protocol::ENVIRONMENT_CONTEXT_CLOSE_TAG, - ) - } - - fn body(&self) -> String { - let mut lines = Vec::new(); - match &self.environments { - EnvironmentContextEnvironments::Single(environment) => { - let cwd = environment.cwd.inferred_native_path_string(); - lines.push(format!(" {cwd}")); - lines.push(format!(" {}", environment.shell)); - } - EnvironmentContextEnvironments::Multiple(environments) => { - lines.push(" ".to_string()); - for environment in environments { - lines.push(format!(" ", environment.id)); - let cwd = environment.cwd.inferred_native_path_string(); - lines.push(format!(" {cwd}")); - lines.push(format!(" {}", environment.shell)); - lines.push(" ".to_string()); - } - lines.push(" ".to_string()); - } - EnvironmentContextEnvironments::None => {} - } - if let Some(current_date) = &self.current_date { - lines.push(format!(" {current_date}")); - } - if let Some(timezone) = &self.timezone { - lines.push(format!(" {timezone}")); - } - match &self.network { - Some(network) => { - lines.push(format!(" {}", network.render())); - } - None => { - // TODO(mbolin): Include this line if it helps the model. - // lines.push(" ".to_string()); - } - } - if let Some(filesystem) = &self.filesystem { - lines.push(format!(" {}", filesystem.render())); - } - if let Some(subagents) = &self.subagents { - lines.push(" ".to_string()); - lines.extend(subagents.lines().map(|line| format!(" {line}"))); - lines.push(" ".to_string()); + rendered_network.push_str(&format!("<{name}>")); + rendered_network.push_str(first); + for domain in domains { + rendered_network.push(','); + rendered_network.push_str(domain); } - format!("\n{}\n", lines.join("\n")) + rendered_network.push_str(&format!("")); } } - -#[cfg(test)] -#[path = "environment_context_tests.rs"] -mod tests; diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 3347358bb587..616d999627ef 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -32,6 +32,8 @@ mod token_budget_context; mod turn_aborted; mod user_instructions; mod user_shell_command; +#[allow(dead_code)] +pub(crate) mod world_state; pub(crate) use approved_command_prefix_saved::ApprovedCommandPrefixSaved; pub(crate) use apps_instructions::AppsInstructions; @@ -47,7 +49,6 @@ pub(crate) use collaboration_mode_instructions::CollaborationModeInstructions; pub(crate) use contextual_user_message::is_contextual_user_fragment; pub(crate) use contextual_user_message::parse_visible_hook_prompt_message; pub(crate) use current_time_reminder::CurrentTimeReminder; -pub(crate) use environment_context::EnvironmentContext; pub(crate) use guardian_followup_review_reminder::GuardianFollowupReviewReminder; pub(crate) use hook_additional_context::HookAdditionalContext; pub(crate) use image_generation_instructions::ImageGenerationInstructions; @@ -76,3 +77,4 @@ pub(crate) use token_budget_context::TokenBudgetRemainingContext; pub(crate) use turn_aborted::TurnAborted; pub(crate) use user_instructions::UserInstructions; pub(crate) use user_shell_command::UserShellCommand; +pub(crate) use world_state::EnvironmentsState; diff --git a/codex-rs/core/src/context/world_state/collaboration_mode.rs b/codex-rs/core/src/context/world_state/collaboration_mode.rs new file mode 100644 index 000000000000..3fe1f0618b56 --- /dev/null +++ b/codex-rs/core/src/context/world_state/collaboration_mode.rs @@ -0,0 +1,38 @@ +use super::WorldStateSection; +use crate::context::CollaborationModeInstructions; +use crate::context::ContextualUserFragment; +use crate::session::turn_context::TurnContext; +use codex_protocol::config_types::CollaborationMode; +use codex_protocol::protocol::TurnContextItem; + +#[derive(Debug)] +pub(crate) struct CollaborationModeState { + mode: Option, + enabled: bool, +} + +impl CollaborationModeState { + pub(crate) fn from_turn_context(turn_context: &TurnContext) -> Self { + Self { + mode: Some(turn_context.collaboration_mode.clone()), + enabled: turn_context.config.include_collaboration_mode_instructions, + } + } + + pub(crate) fn from_turn_context_item(turn_context_item: &TurnContextItem) -> Self { + Self { + mode: turn_context_item.collaboration_mode.clone(), + enabled: false, + } + } +} + +impl WorldStateSection for CollaborationModeState { + fn render_diff(&self, previous: Option<&Self>) -> Option> { + if !self.enabled || previous.is_some_and(|previous| self.mode == previous.mode) { + return None; + } + CollaborationModeInstructions::from_collaboration_mode(self.mode.as_ref()?) + .map(|instructions| Box::new(instructions) as Box) + } +} diff --git a/codex-rs/core/src/context/world_state/environment.rs b/codex-rs/core/src/context/world_state/environment.rs new file mode 100644 index 000000000000..8f227e907953 --- /dev/null +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -0,0 +1,397 @@ +use super::WorldStateSection; +use crate::context::ContextualUserFragment; +use crate::context::environment_context::FileSystemContext; +use crate::context::environment_context::NetworkContext; +use crate::context::environment_context::push_xml_escaped_text; +use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::session::turn_context::TurnContext; +use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_protocol::protocol::TurnContextItem; +use codex_protocol::protocol::TurnContextNetworkItem; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_path_uri::PathUri; +use std::collections::BTreeMap; + +/// Environment values visible to the model. +#[derive(Clone, Debug, Default)] +pub(crate) struct EnvironmentsState { + environments: BTreeMap, + current_date: Option, + timezone: Option, + network: Option, + filesystem: Option, + subagents: Option, +} + +impl PartialEq for EnvironmentsState { + fn eq(&self, other: &Self) -> bool { + self.environments == other.environments + && self.current_date == other.current_date + && self.timezone == other.timezone + && self.network == other.network + && self.filesystem == other.filesystem + } +} + +impl Eq for EnvironmentsState {} + +impl EnvironmentsState { + pub(crate) fn from_turn_context(turn_context: &TurnContext) -> Self { + Self::from_turn_context_with_environments(turn_context, &turn_context.environments) + } + + pub(crate) fn from_turn_context_with_environments( + turn_context: &TurnContext, + environments: &TurnEnvironmentSnapshot, + ) -> Self { + let mut state = Self { + environments: environments + .turn_environments + .iter() + .map(|environment| { + ( + environment.environment_id.clone(), + EnvironmentState { + cwd: environment.cwd().clone(), + status: Some(EnvironmentStatus::Available), + shell: environment + .shell + .as_ref() + .map(|shell| shell.name().to_string()), + }, + ) + }) + .collect(), + current_date: turn_context.current_date.clone(), + timezone: turn_context.timezone.clone(), + network: network_from_turn_context(turn_context), + filesystem: Some(FileSystemContext::from_permission_profile( + &turn_context.permission_profile, + &turn_context.config.effective_workspace_roots(), + )), + subagents: None, + }; + for environment in &environments.starting { + state + .environments + .entry(environment.selection.environment_id.clone()) + .or_insert_with(|| EnvironmentState { + cwd: environment.selection.cwd.clone(), + status: Some(EnvironmentStatus::Starting), + shell: None, + }); + } + state + } + + pub(crate) fn from_turn_context_item(turn_context_item: &TurnContextItem) -> Self { + Self { + environments: [( + LOCAL_ENVIRONMENT_ID.to_string(), + EnvironmentState { + cwd: PathUri::from_abs_path(&turn_context_item.cwd), + status: Some(EnvironmentStatus::Available), + shell: None, + }, + )] + .into_iter() + .collect(), + current_date: turn_context_item.current_date.clone(), + timezone: turn_context_item.timezone.clone(), + network: network_from_turn_context_item(turn_context_item), + filesystem: Some(FileSystemContext::from_permission_profile( + &turn_context_item.permission_profile(), + &workspace_roots_from_turn_context_item(turn_context_item), + )), + subagents: None, + } + } + + pub(crate) fn with_subagents(mut self, subagents: String) -> Self { + if !subagents.is_empty() { + self.subagents = Some(subagents); + } + self + } + + pub(crate) fn render_diff( + &self, + previous: Option<&Self>, + ) -> Option> { + WorldStateSection::render_diff(self, previous) + } + + fn rendered_full(&self) -> RenderedEnvironments { + RenderedEnvironments { + updates: self + .environments + .iter() + .map(|(id, environment)| { + (id.clone(), EnvironmentUpdate::Current(environment.clone())) + }) + .collect(), + legacy_single: is_legacy_single(&self.environments), + current_date: self.current_date.clone(), + timezone: self.timezone.clone(), + network: self.network.clone(), + filesystem: self.filesystem.clone(), + subagents: self.subagents.clone(), + } + } +} + +impl WorldStateSection for EnvironmentsState { + fn render_diff(&self, previous: Option<&Self>) -> Option> { + let empty = Self::default(); + let previous = previous.unwrap_or(&empty); + let turn_context_values_changed = self.current_date != previous.current_date + || self.timezone != previous.timezone + || self.network != previous.network + || self.filesystem != previous.filesystem; + let mut updates = self + .environments + .iter() + .filter(|(id, environment)| previous.environments.get(*id) != Some(*environment)) + .map(|(id, environment)| (id.clone(), EnvironmentUpdate::Current(environment.clone()))) + .collect::>(); + updates.extend( + previous + .environments + .keys() + .filter(|id| !self.environments.contains_key(*id)) + .map(|id| (id.clone(), EnvironmentUpdate::Unavailable)), + ); + let legacy_single = is_legacy_single(&self.environments) + && updates + .values() + .all(|update| matches!(update, EnvironmentUpdate::Current(_))); + (!updates.is_empty() || turn_context_values_changed).then(|| { + Box::new(RenderedEnvironments { + updates, + legacy_single, + current_date: self.current_date.clone(), + timezone: self.timezone.clone(), + network: self.network.clone(), + filesystem: self.filesystem.clone(), + subagents: self.subagents.clone(), + }) as Box + }) + } +} + +impl ContextualUserFragment for EnvironmentsState { + fn role(&self) -> &'static str { + "user" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + environment_context_markers() + } + + fn body(&self) -> String { + self.rendered_full().body() + } +} + +struct RenderedEnvironments { + updates: BTreeMap, + legacy_single: bool, + current_date: Option, + timezone: Option, + network: Option, + filesystem: Option, + subagents: Option, +} + +enum EnvironmentUpdate { + Current(EnvironmentState), + Unavailable, +} + +impl ContextualUserFragment for RenderedEnvironments { + fn role(&self) -> &'static str { + "user" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + environment_context_markers() + } + + fn body(&self) -> String { + let mut rendered = "\n".to_string(); + if self.legacy_single { + if let Some(EnvironmentUpdate::Current(environment)) = self.updates.values().next() { + push_environment_values(&mut rendered, environment, " "); + } + } else if !self.updates.is_empty() { + rendered.push_str(" \n"); + for (id, update) in &self.updates { + match update { + EnvironmentUpdate::Current(environment) => { + rendered.push_str(" \n"); + push_environment_values(&mut rendered, environment, " "); + rendered.push_str(" \n"); + } + EnvironmentUpdate::Unavailable => { + rendered.push_str(" \n"); + } + } + } + rendered.push_str(" \n"); + } + push_optional_element(&mut rendered, "current_date", self.current_date.as_deref()); + push_optional_element(&mut rendered, "timezone", self.timezone.as_deref()); + if let Some(network) = &self.network { + rendered.push_str(" "); + rendered.push_str(&network.render()); + rendered.push('\n'); + } + if let Some(filesystem) = &self.filesystem { + rendered.push_str(" "); + rendered.push_str(&filesystem.render()); + rendered.push('\n'); + } + if let Some(subagents) = &self.subagents { + rendered.push_str(" \n"); + for line in subagents.lines() { + rendered.push_str(" "); + rendered.push_str(line); + rendered.push('\n'); + } + rendered.push_str(" \n"); + } + rendered + } +} + +fn push_environment_values(rendered: &mut String, environment: &EnvironmentState, indent: &str) { + rendered.push_str(indent); + rendered.push_str(""); + push_xml_escaped_text(rendered, &environment.cwd.inferred_native_path_string()); + rendered.push_str("\n"); + if let Some(shell) = &environment.shell { + rendered.push_str(indent); + rendered.push_str(""); + push_xml_escaped_text(rendered, shell); + rendered.push_str("\n"); + } +} + +fn push_optional_element(rendered: &mut String, name: &str, value: Option<&str>) { + let Some(value) = value else { + return; + }; + rendered.push_str(" <"); + rendered.push_str(name); + rendered.push('>'); + push_xml_escaped_text(rendered, value); + rendered.push_str("\n"); +} + +#[derive(Clone, Debug)] +struct EnvironmentState { + cwd: PathUri, + status: Option, + shell: Option, +} + +impl PartialEq for EnvironmentState { + fn eq(&self, other: &Self) -> bool { + self.cwd == other.cwd + } +} + +impl Eq for EnvironmentState {} + +#[derive(Clone, Copy, Debug)] +enum EnvironmentStatus { + Starting, + Available, +} + +fn is_legacy_single(environments: &BTreeMap) -> bool { + environments.len() == 1 + && matches!( + environments + .values() + .next() + .and_then(|environment| environment.status), + Some(EnvironmentStatus::Available) + ) +} + +fn environment_context_markers() -> (&'static str, &'static str) { + ( + codex_protocol::protocol::ENVIRONMENT_CONTEXT_OPEN_TAG, + codex_protocol::protocol::ENVIRONMENT_CONTEXT_CLOSE_TAG, + ) +} + +fn network_from_turn_context(turn_context: &TurnContext) -> Option { + let network = turn_context + .config + .config_layer_stack + .requirements() + .network + .as_ref()?; + + Some(NetworkContext::new( + network + .domains + .as_ref() + .and_then(codex_config::NetworkDomainPermissionsToml::allowed_domains) + .unwrap_or_default(), + network + .domains + .as_ref() + .and_then(codex_config::NetworkDomainPermissionsToml::denied_domains) + .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::new( + allowed_domains.clone(), + denied_domains.clone(), + )) +} + +fn workspace_roots_from_turn_context_item( + turn_context_item: &TurnContextItem, +) -> Vec { + if let Some(workspace_roots) = turn_context_item.workspace_roots.as_ref() { + return workspace_roots.clone(); + } + + vec![turn_context_item.cwd.clone()] +} + +#[cfg(test)] +#[path = "environment_tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "environment_render_tests.rs"] +mod render_tests; diff --git a/codex-rs/core/src/context/environment_context_tests.rs b/codex-rs/core/src/context/world_state/environment_render_tests.rs similarity index 67% rename from codex-rs/core/src/context/environment_context_tests.rs rename to codex-rs/core/src/context/world_state/environment_render_tests.rs index d40ce4e671d6..b8715023551a 100644 --- a/codex-rs/core/src/context/environment_context_tests.rs +++ b/codex-rs/core/src/context/world_state/environment_render_tests.rs @@ -30,15 +30,43 @@ fn test_abs_path(unix_path: &str) -> AbsolutePathBuf { test_path_buf(unix_path).abs() } +fn environment(id: &str, cwd: PathUri, shell: impl Into) -> (String, EnvironmentState) { + ( + id.to_string(), + EnvironmentState { + cwd, + status: Some(EnvironmentStatus::Available), + shell: Some(shell.into()), + }, + ) +} + +fn environment_state( + environments: impl IntoIterator, + current_date: Option, + timezone: Option, + network: Option, + subagents: Option, +) -> EnvironmentsState { + EnvironmentsState { + environments: environments.into_iter().collect(), + current_date, + timezone, + network, + filesystem: None, + subagents, + } +} + #[test] fn serialize_workspace_write_environment_context() { let cwd = test_path_buf("/repo"); - let context = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&cwd.abs()), - shell: fake_shell_name(), - }], + let context = environment_state( + [environment( + "local", + PathUri::from_abs_path(&cwd.abs()), + fake_shell_name(), + )], Some("2026-02-26".to_string()), Some("America/Los_Angeles".to_string()), /*network*/ None, @@ -60,12 +88,12 @@ fn serialize_workspace_write_environment_context() { #[test] fn serialize_environment_context_with_foreign_windows_cwd() { - let context = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "remote".to_string(), - cwd: PathUri::parse("file:///C:/windows").expect("Windows cwd URI"), - shell: "powershell".to_string(), - }], + let context = environment_state( + [environment( + "remote", + PathUri::parse("file:///C:/windows").expect("Windows cwd URI"), + "powershell", + )], /*current_date*/ None, /*timezone*/ None, /*network*/ None, @@ -87,12 +115,12 @@ fn serialize_environment_context_with_network() { vec!["api.example.com".to_string(), "*.openai.com".to_string()], vec!["blocked.example.com".to_string()], ); - let context = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&test_abs_path("/repo")), - shell: fake_shell_name(), - }], + let context = environment_state( + [environment( + "local", + PathUri::from_abs_path(&test_abs_path("/repo")), + fake_shell_name(), + )], Some("2026-02-26".to_string()), Some("America/Los_Angeles".to_string()), Some(network), @@ -149,12 +177,12 @@ fn serialize_environment_context_with_full_filesystem_profile() { AbsolutePathBuf::resolve_path_against_base(Path::new("private/**"), repo.as_path()); let other_repo_private_glob = AbsolutePathBuf::resolve_path_against_base(Path::new("private/**"), other_repo.as_path()); - let mut context = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&test_abs_path("/repo")), - shell: fake_shell_name(), - }], + let mut context = environment_state( + [environment( + "local", + PathUri::from_abs_path(&test_abs_path("/repo")), + fake_shell_name(), + )], /*current_date*/ None, /*timezone*/ None, /*network*/ None, @@ -210,7 +238,7 @@ fn turn_context_item_filesystem_uses_workspace_roots_instead_of_cwd() { summary: codex_protocol::config_types::ReasoningSummary::Auto, }; - let context = EnvironmentContext::from_turn_context_item(&item, fake_shell_name()).render(); + let context = EnvironmentsState::from_turn_context_item(&item).render(); assert!( context.contains(&format!( @@ -237,7 +265,7 @@ fn turn_context_item_filesystem_uses_workspace_roots_instead_of_cwd() { #[test] fn serialize_read_only_environment_context() { - let context = EnvironmentContext::new( + let context = environment_state( Vec::new(), Some("2026-02-26".to_string()), Some("America/Los_Angeles".to_string()), @@ -253,97 +281,14 @@ fn serialize_read_only_environment_context() { assert_eq!(context.render(), expected); } -#[test] -fn equals_except_shell_compares_cwd() { - let context1 = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&test_abs_path("/repo")), - shell: fake_shell_name(), - }], - /*current_date*/ None, - /*timezone*/ None, - /*network*/ None, - /*subagents*/ None, - ); - let context2 = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&test_abs_path("/repo")), - shell: fake_shell_name(), - }], - /*current_date*/ None, - /*timezone*/ None, - /*network*/ None, - /*subagents*/ None, - ); - assert!(context1.equals_except_shell(&context2)); -} - -#[test] -fn equals_except_shell_compares_cwd_differences() { - let context1 = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&test_abs_path("/repo1")), - shell: fake_shell_name(), - }], - /*current_date*/ None, - /*timezone*/ None, - /*network*/ None, - /*subagents*/ None, - ); - let context2 = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&test_abs_path("/repo2")), - shell: fake_shell_name(), - }], - /*current_date*/ None, - /*timezone*/ None, - /*network*/ None, - /*subagents*/ None, - ); - - assert!(!context1.equals_except_shell(&context2)); -} - -#[test] -fn equals_except_shell_ignores_shell() { - let context1 = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&test_abs_path("/repo")), - shell: "bash".to_string(), - }], - /*current_date*/ None, - /*timezone*/ None, - /*network*/ None, - /*subagents*/ None, - ); - let context2 = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "other".to_string(), - cwd: PathUri::from_abs_path(&test_abs_path("/repo")), - shell: "zsh".to_string(), - }], - /*current_date*/ None, - /*timezone*/ None, - /*network*/ None, - /*subagents*/ None, - ); - - assert!(context1.equals_except_shell(&context2)); -} - #[test] fn serialize_environment_context_with_subagents() { - let context = EnvironmentContext::new( - vec![EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&test_abs_path("/repo")), - shell: fake_shell_name(), - }], + let context = environment_state( + [environment( + "local", + PathUri::from_abs_path(&test_abs_path("/repo")), + fake_shell_name(), + )], Some("2026-02-26".to_string()), Some("America/Los_Angeles".to_string()), /*network*/ None, @@ -371,18 +316,10 @@ fn serialize_environment_context_with_subagents() { fn serialize_environment_context_with_multiple_selected_environments() { let local_cwd = test_path_buf("/repo/local"); let remote_cwd = test_path_buf("/repo/remote"); - let context = EnvironmentContext::new( - vec![ - EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&local_cwd.abs()), - shell: "bash".to_string(), - }, - EnvironmentContextEnvironment { - id: "remote".to_string(), - cwd: PathUri::from_abs_path(&remote_cwd.abs()), - shell: "bash".to_string(), - }, + let context = environment_state( + [ + environment("local", PathUri::from_abs_path(&local_cwd.abs()), "bash"), + environment("remote", PathUri::from_abs_path(&remote_cwd.abs()), "bash"), ], Some("2026-02-26".to_string()), Some("America/Los_Angeles".to_string()), @@ -416,18 +353,14 @@ fn serialize_environment_context_with_multiple_selected_environments() { fn serialize_environment_context_prefers_environment_shell_when_present() { let local_cwd = test_path_buf("/repo/local"); let remote_cwd = test_path_buf("/repo/remote"); - let context = EnvironmentContext::new( - vec![ - EnvironmentContextEnvironment { - id: "local".to_string(), - cwd: PathUri::from_abs_path(&local_cwd.abs()), - shell: "powershell".to_string(), - }, - EnvironmentContextEnvironment { - id: "remote".to_string(), - cwd: PathUri::from_abs_path(&remote_cwd.abs()), - shell: "cmd".to_string(), - }, + let context = environment_state( + [ + environment( + "local", + PathUri::from_abs_path(&local_cwd.abs()), + "powershell", + ), + environment("remote", PathUri::from_abs_path(&remote_cwd.abs()), "cmd"), ], /*current_date*/ None, /*timezone*/ None, diff --git a/codex-rs/core/src/context/world_state/environment_tests.rs b/codex-rs/core/src/context/world_state/environment_tests.rs new file mode 100644 index 000000000000..bc54ae4bb6c1 --- /dev/null +++ b/codex-rs/core/src/context/world_state/environment_tests.rs @@ -0,0 +1,253 @@ +use super::*; +use crate::context::ContextualUserFragment; +use crate::context::world_state::WorldState; +use crate::context_manager::updates::merge_contextual_fragments; +use anyhow::Result; +use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_protocol::models::ContentItem; +use codex_protocol::models::PermissionProfile; +use codex_protocol::models::ResponseItem; +use codex_protocol::permissions::NetworkSandboxPolicy; +use pretty_assertions::assert_eq; + +#[test] +fn renders_full_environment_state() -> Result<()> { + let context = EnvironmentsState { + environments: [ + ("laptop".to_string(), available("file:///repo", "zsh")?), + ( + "devbox".to_string(), + available("file:///workspace", "bash")?, + ), + ] + .into_iter() + .collect(), + current_date: Some("2026-06-20".to_string()), + timezone: Some("America/Los_Angeles".to_string()), + network: Some(NetworkContext::new( + vec!["api.example.com".to_string()], + vec!["blocked.example.com".to_string()], + )), + filesystem: Some(FileSystemContext::from_permission_profile( + &PermissionProfile::Disabled, + &[], + )), + subagents: Some("task_1: running\ntask_2: completed".to_string()), + }; + + let mut world_state = WorldState::default(); + world_state.add_section(context); + + assert_eq!( + vec![user_message( + r#" + + + /workspace + bash + + + /repo + zsh + + + 2026-06-20 + America/Los_Angeles + api.example.comblocked.example.com + + + task_1: running + task_2: completed + +"#, + )], + merge_contextual_fragments(world_state.render_full()), + ); + Ok(()) +} + +#[test] +fn renders_only_changed_environments() -> Result<()> { + let mut previous = WorldState::default(); + previous.add_section(EnvironmentsState { + environments: [ + ("laptop".to_string(), available("file:///repo", "bash")?), + ("devbox".to_string(), starting("file:///workspace")?), + ("old".to_string(), available("file:///old", "sh")?), + ] + .into_iter() + .collect(), + ..Default::default() + }); + let mut current = WorldState::default(); + current.add_section(EnvironmentsState { + environments: [ + ("laptop".to_string(), available("file:///new-repo", "zsh")?), + ( + "devbox".to_string(), + available("file:///workspace", "powershell")?, + ), + ("remote".to_string(), starting("file:///remote")?), + ] + .into_iter() + .collect(), + ..Default::default() + }); + + assert_eq!( + vec![user_message( + r#" + + + /new-repo + zsh + + + + /remote + + +"#, + )], + merge_contextual_fragments(current.render_diff(&previous)), + ); + Ok(()) +} + +#[test] +fn persisted_turn_context_values_render_a_diff() -> Result<()> { + let environments = EnvironmentsState { + environments: [( + LOCAL_ENVIRONMENT_ID.to_string(), + available("file:///repo", "zsh")?, + )] + .into_iter() + .collect(), + ..Default::default() + }; + let mut previous = WorldState::default(); + previous.add_section(EnvironmentsState { + current_date: Some("2026-06-19".to_string()), + timezone: Some("UTC".to_string()), + network: Some(NetworkContext::new( + vec!["old.example.com".to_string()], + vec![], + )), + filesystem: Some(FileSystemContext::from_permission_profile( + &PermissionProfile::Disabled, + &[], + )), + ..environments.clone() + }); + let mut current = WorldState::default(); + current.add_section(EnvironmentsState { + current_date: Some("2026-06-20".to_string()), + timezone: Some("America/Los_Angeles".to_string()), + network: Some(NetworkContext::new( + vec!["new.example.com".to_string()], + vec!["blocked.example.com".to_string()], + )), + filesystem: Some(FileSystemContext::from_permission_profile( + &PermissionProfile::External { + network: NetworkSandboxPolicy::Restricted, + }, + &[], + )), + ..environments + }); + + assert_eq!( + vec![user_message( + r#" + 2026-06-20 + America/Los_Angeles + new.example.comblocked.example.com + +"#, + )], + merge_contextual_fragments(current.render_diff(&previous)), + ); + Ok(()) +} + +#[test] +fn single_environment_diff_ignores_shell() -> Result<()> { + let previous = EnvironmentsState { + environments: [( + LOCAL_ENVIRONMENT_ID.to_string(), + available("file:///repo", "bash")?, + )] + .into_iter() + .collect(), + ..Default::default() + }; + let current = EnvironmentsState { + environments: [( + LOCAL_ENVIRONMENT_ID.to_string(), + available("file:///repo", "zsh")?, + )] + .into_iter() + .collect(), + ..Default::default() + }; + + assert_eq!(None, fragment_item(current.render_diff(Some(&previous)))); + Ok(()) +} + +#[test] +fn removed_legacy_environment_renders_unavailable() -> Result<()> { + let previous = EnvironmentsState { + environments: [( + LOCAL_ENVIRONMENT_ID.to_string(), + available("file:///repo", "bash")?, + )] + .into_iter() + .collect(), + ..Default::default() + }; + + assert_eq!( + Some(user_message( + r#" + + + +"#, + )), + fragment_item(EnvironmentsState::default().render_diff(Some(&previous))), + ); + Ok(()) +} + +fn available(cwd: &str, shell: &str) -> Result { + Ok(EnvironmentState { + cwd: PathUri::parse(cwd)?, + status: Some(EnvironmentStatus::Available), + shell: Some(shell.to_string()), + }) +} + +fn starting(cwd: &str) -> Result { + Ok(EnvironmentState { + cwd: PathUri::parse(cwd)?, + status: Some(EnvironmentStatus::Starting), + shell: None, + }) +} + +fn user_message(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + phase: None, + metadata: None, + } +} + +fn fragment_item(fragment: Option>) -> Option { + fragment.map(codex_context_fragments::ContextualUserFragment::into_boxed_response_item) +} diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs new file mode 100644 index 000000000000..de1b52536c5e --- /dev/null +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -0,0 +1,94 @@ +mod collaboration_mode; +mod environment; +mod model; +mod multi_agent_mode; +mod permissions; +mod personality; +mod realtime; + +use crate::context::ContextualUserFragment; +use indexmap::IndexMap; +use std::any::Any; +use std::any::TypeId; +use std::fmt; + +pub(crate) use collaboration_mode::CollaborationModeState; +pub(crate) use environment::EnvironmentsState; +pub(crate) use model::ModelState; +pub(crate) use multi_agent_mode::MultiAgentModeState; +pub(crate) use permissions::PermissionsState; +pub(crate) use personality::PersonalityState; +pub(crate) use realtime::RealtimeState; + +trait ErasedWorldStateSection: Send + Sync { + fn as_any(&self) -> &dyn Any; + + fn render_diff(&self, previous: Option<&dyn Any>) -> Option>; +} + +impl ErasedWorldStateSection for S { + fn as_any(&self) -> &dyn Any { + self + } + + fn render_diff(&self, previous: Option<&dyn Any>) -> Option> { + let previous = match previous { + Some(previous) => { + let Some(previous) = previous.downcast_ref::() else { + unreachable!("world-state section type must match its type ID"); + }; + Some(previous) + } + None => None, + }; + WorldStateSection::render_diff(self, previous) + } +} + +/// A typed portion of the state visible to the model. +/// +/// Implementations own how their current state is rendered relative to an +/// earlier value of the same section type. A missing previous value requests +/// the section's complete current representation; sections without a +/// standalone full representation may return no fragment. +pub(crate) trait WorldStateSection: Any + Send + Sync { + fn render_diff(&self, previous: Option<&Self>) -> Option>; +} + +/// A snapshot of the model-visible world with one section per concrete type. +#[derive(Default)] +pub(crate) struct WorldState { + sections: IndexMap>, +} + +impl fmt::Debug for WorldState { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("WorldState") + .field("section_count", &self.sections.len()) + .finish() + } +} + +impl WorldState { + pub(crate) fn add_section(&mut self, section: S) { + self.sections.insert(TypeId::of::(), Box::new(section)); + } + + pub(crate) fn render_full(&self) -> Vec> { + self.render_diff(&Self::default()) + } + + pub(crate) fn render_diff(&self, previous: &Self) -> Vec> { + self.sections + .iter() + .filter_map(|(type_id, section)| { + let previous = previous + .sections + .get(type_id) + .map(|section| section.as_any()); + section.render_diff(previous) + }) + .collect() + } +} diff --git a/codex-rs/core/src/context/world_state/model.rs b/codex-rs/core/src/context/world_state/model.rs new file mode 100644 index 000000000000..628cbe3fac12 --- /dev/null +++ b/codex-rs/core/src/context/world_state/model.rs @@ -0,0 +1,38 @@ +use super::WorldStateSection; +use crate::context::ContextualUserFragment; +use crate::context::ModelSwitchInstructions; +use crate::session::turn_context::TurnContext; + +#[derive(Debug)] +pub(crate) struct ModelState { + model: String, + instructions: String, +} + +impl ModelState { + pub(crate) fn from_turn_context(turn_context: &TurnContext) -> Self { + Self { + model: turn_context.model_info.slug.clone(), + instructions: turn_context + .model_info + .get_model_instructions(turn_context.personality), + } + } + + pub(crate) fn from_previous_model(model: &str) -> Self { + Self { + model: model.to_string(), + instructions: String::new(), + } + } +} + +impl WorldStateSection for ModelState { + fn render_diff(&self, previous: Option<&Self>) -> Option> { + let previous = previous?; + if self.model == previous.model || self.instructions.is_empty() { + return None; + } + Some(Box::new(ModelSwitchInstructions::new(&self.instructions))) + } +} diff --git a/codex-rs/core/src/context/world_state/multi_agent_mode.rs b/codex-rs/core/src/context/world_state/multi_agent_mode.rs new file mode 100644 index 000000000000..76b51d85fb17 --- /dev/null +++ b/codex-rs/core/src/context/world_state/multi_agent_mode.rs @@ -0,0 +1,48 @@ +use super::WorldStateSection; +use crate::context::ContextualUserFragment; +use crate::context::MultiAgentModeInstructions; +use crate::session::multi_agents; +use crate::session::turn_context::TurnContext; +use codex_features::Feature; +use codex_protocol::config_types::MultiAgentMode; +use codex_protocol::protocol::TurnContextItem; + +#[derive(Debug)] +pub(crate) struct MultiAgentModeState(Option); + +impl MultiAgentModeState { + pub(crate) fn from_turn_context(turn_context: &TurnContext) -> Self { + Self(multi_agents::effective_multi_agent_mode( + turn_context.multi_agent_version, + &turn_context.config.multi_agent_v2, + &turn_context.session_source, + turn_context.multi_agent_mode, + turn_context + .config + .features + .enabled(Feature::MultiAgentMode), + )) + } + + pub(crate) fn from_turn_context_item(turn_context_item: &TurnContextItem) -> Self { + Self(turn_context_item.multi_agent_mode) + } +} + +impl WorldStateSection for MultiAgentModeState { + fn render_diff(&self, previous: Option<&Self>) -> Option> { + if previous.is_some_and(|previous| self.0 == previous.0) { + return None; + } + let mode = match self.0 { + Some(mode) => mode, + None if previous + .is_some_and(|previous| previous.0 == Some(MultiAgentMode::Proactive)) => + { + MultiAgentMode::ExplicitRequestOnly + } + None => return None, + }; + Some(Box::new(MultiAgentModeInstructions::new(mode))) + } +} diff --git a/codex-rs/core/src/context/world_state/permissions.rs b/codex-rs/core/src/context/world_state/permissions.rs new file mode 100644 index 000000000000..47e35baa7ee4 --- /dev/null +++ b/codex-rs/core/src/context/world_state/permissions.rs @@ -0,0 +1,73 @@ +use super::WorldStateSection; +use crate::context::ContextualUserFragment; +use crate::context::PermissionsInstructions; +use crate::session::turn_context::TurnContext; +use codex_execpolicy::Policy; +use codex_features::Feature; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::TurnContextItem; + +#[derive(Debug, PartialEq)] +struct PermissionValues { + permission_profile: PermissionProfile, + approval_policy: AskForApproval, +} + +#[derive(Debug)] +pub(crate) struct PermissionsState { + values: PermissionValues, + instructions: Option, +} + +impl PermissionsState { + pub(crate) fn from_turn_context(turn_context: &TurnContext, exec_policy: &Policy) -> Self { + let instructions = turn_context + .config + .include_permissions_instructions + .then(|| { + PermissionsInstructions::from_permission_profile( + &turn_context.permission_profile, + turn_context.approval_policy.value(), + turn_context.config.approvals_reviewer, + exec_policy, + #[allow(deprecated)] + &turn_context.cwd, + turn_context + .config + .features + .enabled(Feature::ExecPermissionApprovals), + turn_context + .config + .features + .enabled(Feature::RequestPermissionsTool), + ) + }); + Self { + values: PermissionValues { + permission_profile: turn_context.permission_profile(), + approval_policy: turn_context.approval_policy.value(), + }, + instructions, + } + } + + pub(crate) fn from_turn_context_item(turn_context_item: &TurnContextItem) -> Self { + Self { + values: PermissionValues { + permission_profile: turn_context_item.permission_profile(), + approval_policy: turn_context_item.approval_policy, + }, + instructions: None, + } + } +} + +impl WorldStateSection for PermissionsState { + fn render_diff(&self, previous: Option<&Self>) -> Option> { + let instructions = self.instructions.as_ref()?; + previous + .is_none_or(|previous| self.values != previous.values) + .then(|| Box::new(instructions.clone()) as Box) + } +} diff --git a/codex-rs/core/src/context/world_state/personality.rs b/codex-rs/core/src/context/world_state/personality.rs new file mode 100644 index 000000000000..a36882e85bac --- /dev/null +++ b/codex-rs/core/src/context/world_state/personality.rs @@ -0,0 +1,61 @@ +use super::WorldStateSection; +use crate::context::ContextualUserFragment; +use crate::context::PersonalitySpecInstructions; +use crate::session::turn_context::TurnContext; +use codex_protocol::config_types::Personality; +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::protocol::TurnContextItem; + +#[derive(Debug)] +pub(crate) struct PersonalityState { + enabled: bool, + model: String, + personality: Option, + spec: Option, +} + +impl PersonalityState { + pub(crate) fn from_turn_context(turn_context: &TurnContext, enabled: bool) -> Self { + let personality = turn_context.personality; + Self { + enabled, + model: turn_context.model_info.slug.clone(), + personality, + spec: personality + .and_then(|personality| Self::message(&turn_context.model_info, personality)), + } + } + + pub(crate) fn from_turn_context_item(turn_context_item: &TurnContextItem) -> Self { + Self { + enabled: false, + model: turn_context_item.model.clone(), + personality: turn_context_item.personality, + spec: None, + } + } + + pub(crate) fn message(model_info: &ModelInfo, personality: Personality) -> Option { + model_info + .model_messages + .as_ref() + .and_then(|spec| spec.get_personality_message(Some(personality))) + .filter(|message| !message.is_empty()) + } +} + +impl WorldStateSection for PersonalityState { + fn render_diff(&self, previous: Option<&Self>) -> Option> { + if !self.enabled + || previous.is_some_and(|previous| { + self.model != previous.model || self.personality == previous.personality + }) + { + return None; + } + self.personality?; + self.spec.as_ref().map(|spec| { + Box::new(PersonalitySpecInstructions::new(spec)) as Box + }) + } +} diff --git a/codex-rs/core/src/context/world_state/realtime.rs b/codex-rs/core/src/context/world_state/realtime.rs new file mode 100644 index 000000000000..327c482938d4 --- /dev/null +++ b/codex-rs/core/src/context/world_state/realtime.rs @@ -0,0 +1,58 @@ +use super::WorldStateSection; +use crate::context::ContextualUserFragment; +use crate::context::RealtimeEndInstructions; +use crate::context::RealtimeStartInstructions; +use crate::context::RealtimeStartWithInstructions; +use crate::session::turn_context::TurnContext; + +#[derive(Debug)] +pub(crate) struct RealtimeState { + active: Option, + active_fallback: Option, + start_instructions: Option, +} + +impl RealtimeState { + pub(crate) fn from_turn_context(turn_context: &TurnContext) -> Self { + Self { + active: Some(turn_context.realtime_active), + active_fallback: None, + start_instructions: turn_context + .config + .experimental_realtime_start_instructions + .clone(), + } + } + + pub(crate) fn from_previous(active: Option, active_fallback: Option) -> Self { + Self { + active, + active_fallback, + start_instructions: None, + } + } +} + +impl WorldStateSection for RealtimeState { + fn render_diff(&self, previous: Option<&Self>) -> Option> { + let previous_active = previous.and_then(|previous| previous.active); + match (previous_active, self.active.unwrap_or(false)) { + (Some(true), false) => Some(Box::new(RealtimeEndInstructions::new("inactive"))), + (Some(false), true) | (None, true) => { + if let Some(instructions) = self.start_instructions.as_deref() { + Some(Box::new(RealtimeStartWithInstructions::new(instructions))) + } else { + Some(Box::new(RealtimeStartInstructions)) + } + } + (Some(true), true) | (Some(false), false) => None, + (None, false) => previous + .and_then(|previous| previous.active_fallback) + .filter(|active| *active) + .map(|_| { + Box::new(RealtimeEndInstructions::new("inactive")) + as Box + }), + } + } +} diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 6ce1c6148f89..546aa326f652 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -1,3 +1,4 @@ +use crate::context::world_state::WorldState; use crate::context_manager::normalize; use crate::event_mapping::has_non_contextual_dev_message_content; use crate::event_mapping::is_contextual_dev_message_content; @@ -27,6 +28,7 @@ use codex_utils_output_truncation::truncate_function_output_items_with_policy; use codex_utils_output_truncation::truncate_text; use std::num::NonZeroUsize; use std::ops::Deref; +use std::sync::Arc; use std::sync::LazyLock; /// Transcript of thread history @@ -48,6 +50,8 @@ pub(crate) struct ContextManager { /// also clear this when it trims a mixed initial-context developer bundle /// whose non-diff fragments no longer exist in the surviving history. reference_context_item: Option, + /// Latest in-memory world snapshot represented by this history. + world_state: Option>, } impl ContextManager { @@ -59,6 +63,7 @@ impl ContextManager { &None, &None, /*model_context_window*/ None, ), reference_context_item: None, + world_state: None, } } @@ -71,6 +76,9 @@ impl ContextManager { } pub(crate) fn set_reference_context_item(&mut self, item: Option) { + if item.is_none() { + self.world_state = None; + } self.reference_context_item = item; } @@ -78,6 +86,14 @@ impl ContextManager { self.reference_context_item.clone() } + pub(crate) fn set_world_state(&mut self, world_state: Arc) { + self.world_state = Some(world_state); + } + + pub(crate) fn world_state(&self) -> Option> { + self.world_state.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), @@ -412,6 +428,7 @@ impl ContextManager { // steady-state diffs once trimmed, so the next real turn must fully // reinject context instead of diffing against a stale baseline. self.reference_context_item = None; + self.world_state = None; } cut_idx -= 1; } diff --git a/codex-rs/core/src/context_manager/updates.rs b/codex-rs/core/src/context_manager/updates.rs index 7cf877a72856..9f79e002935d 100644 --- a/codex-rs/core/src/context_manager/updates.rs +++ b/codex-rs/core/src/context_manager/updates.rs @@ -1,214 +1,25 @@ -use crate::context::CollaborationModeInstructions; use crate::context::ContextualUserFragment; -use crate::context::EnvironmentContext; -use crate::context::ModelSwitchInstructions; -use crate::context::MultiAgentModeInstructions; -use crate::context::PermissionsInstructions; -use crate::context::PersonalitySpecInstructions; -use crate::context::RealtimeEndInstructions; -use crate::context::RealtimeStartInstructions; -use crate::context::RealtimeStartWithInstructions; -use crate::session::PreviousTurnSettings; -use crate::session::turn_context::TurnContext; -use crate::shell::Shell; -use codex_execpolicy::Policy; -use codex_features::Feature; -use codex_protocol::config_types::MultiAgentMode; -use codex_protocol::config_types::Personality; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; -use codex_protocol::openai_models::ModelInfo; -use codex_protocol::protocol::TurnContextItem; -fn build_environment_update_item( - previous: Option<&TurnContextItem>, - next: &TurnContext, - shell: &Shell, -) -> Option { - if !next.config.include_environment_context { - return None; - } - - let prev = previous?; - let prev_context = EnvironmentContext::from_turn_context_item(prev, shell.name().to_string()); - let next_context = EnvironmentContext::from_turn_context(next, shell); - if prev_context.equals_except_shell(&next_context) { - return None; - } - - Some(ContextualUserFragment::into( - EnvironmentContext::diff_from_turn_context_item(prev, &next_context), - )) -} - -fn build_permissions_update_item( - previous: Option<&TurnContextItem>, - next: &TurnContext, - exec_policy: &Policy, -) -> Option { - if !next.config.include_permissions_instructions { - return None; - } - - let prev = previous?; - if prev.permission_profile() == next.permission_profile() - && prev.approval_policy == next.approval_policy.value() - { - return None; - } - - Some( - PermissionsInstructions::from_permission_profile( - &next.permission_profile, - next.approval_policy.value(), - next.config.approvals_reviewer, - exec_policy, - #[allow(deprecated)] - &next.cwd, - next.config - .features - .enabled(Feature::ExecPermissionApprovals), - next.config - .features - .enabled(Feature::RequestPermissionsTool), - ) - .render(), - ) -} - -fn build_collaboration_mode_update_item( - previous: Option<&TurnContextItem>, - next: &TurnContext, -) -> Option { - if !next.config.include_collaboration_mode_instructions { - return None; - } - - let prev = previous?; - 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( - CollaborationModeInstructions::from_collaboration_mode(&next.collaboration_mode)? - .render(), - ) - } else { - None - } -} - -fn build_multi_agent_mode_update_item( - previous: Option<&TurnContextItem>, - next: &TurnContext, -) -> Option { - let effective_multi_agent_mode = crate::session::multi_agents::effective_multi_agent_mode( - next.multi_agent_version, - &next.config.multi_agent_v2, - &next.session_source, - next.multi_agent_mode, - next.config.features.enabled(Feature::MultiAgentMode), - ); - let previous = previous?; - if previous.multi_agent_mode == effective_multi_agent_mode { - return None; - } - - match effective_multi_agent_mode { - Some(multi_agent_mode) => Some(MultiAgentModeInstructions::new(multi_agent_mode).render()), - None if previous.multi_agent_mode == Some(MultiAgentMode::Proactive) => { - Some(MultiAgentModeInstructions::new(MultiAgentMode::ExplicitRequestOnly).render()) +pub(crate) fn merge_contextual_fragments( + fragments: Vec>, +) -> Vec { + let mut messages: Vec<(&str, Vec)> = Vec::with_capacity(fragments.len()); + for fragment in fragments { + let role = fragment.role(); + let text = fragment.render(); + match messages.last_mut() { + Some((previous_role, text_sections)) if *previous_role == role => { + text_sections.push(text); + } + _ => messages.push((role, vec![text])), } - None => None, - } -} - -pub(crate) fn build_realtime_update_item( - previous: Option<&TurnContextItem>, - previous_turn_settings: Option<&PreviousTurnSettings>, - next: &TurnContext, -) -> Option { - match ( - previous.and_then(|item| item.realtime_active), - next.realtime_active, - ) { - (Some(true), false) => Some(RealtimeEndInstructions::new("inactive").render()), - (Some(false), true) | (None, true) => Some( - if let Some(instructions) = next - .config - .experimental_realtime_start_instructions - .as_deref() - { - RealtimeStartWithInstructions::new(instructions).render() - } else { - RealtimeStartInstructions.render() - }, - ), - (Some(true), true) | (Some(false), false) => None, - (None, false) => previous_turn_settings - .and_then(|settings| settings.realtime_active) - .filter(|realtime_active| *realtime_active) - .map(|_| RealtimeEndInstructions::new("inactive").render()), - } -} - -pub(crate) fn build_initial_realtime_item( - previous: Option<&TurnContextItem>, - previous_turn_settings: Option<&PreviousTurnSettings>, - next: &TurnContext, -) -> Option { - build_realtime_update_item(previous, previous_turn_settings, next) -} - -fn build_personality_update_item( - previous: Option<&TurnContextItem>, - next: &TurnContext, - personality_feature_enabled: bool, -) -> Option { - if !personality_feature_enabled { - return None; - } - let previous = previous?; - if next.model_info.slug != previous.model { - return None; - } - - if let Some(personality) = next.personality - && next.personality != previous.personality - { - let model_info = &next.model_info; - let personality_message = personality_message_for(model_info, personality); - personality_message.map(|message| PersonalitySpecInstructions::new(message).render()) - } else { - None - } -} - -pub(crate) fn personality_message_for( - model_info: &ModelInfo, - personality: Personality, -) -> Option { - model_info - .model_messages - .as_ref() - .and_then(|spec| spec.get_personality_message(Some(personality))) - .filter(|message| !message.is_empty()) -} - -pub(crate) fn build_model_instructions_update_item( - previous_turn_settings: Option<&PreviousTurnSettings>, - next: &TurnContext, -) -> Option { - let previous_turn_settings = previous_turn_settings?; - if previous_turn_settings.model == next.model_info.slug { - return None; } - - let model_instructions = next.model_info.get_model_instructions(next.personality); - if model_instructions.is_empty() { - return None; - } - - Some(ModelSwitchInstructions::new(model_instructions).render()) + messages + .into_iter() + .filter_map(|(role, text_sections)| build_text_message(role, text_sections)) + .collect() } pub(crate) fn build_developer_update_item(text_sections: Vec) -> Option { @@ -237,40 +48,3 @@ fn build_text_message(role: &str, text_sections: Vec) -> Option, - previous_turn_settings: Option<&PreviousTurnSettings>, - next: &TurnContext, - shell: &Shell, - exec_policy: &Policy, - personality_feature_enabled: bool, -) -> Vec { - // TODO(ccunningham): build_settings_update_items still does not cover every - // model-visible item emitted by build_initial_context. Persist the remaining - // inputs or add explicit replay events so fork/resume can diff everything - // deterministically. - let contextual_user_message = build_environment_update_item(previous, next, shell); - let developer_update_sections = [ - // Keep model-switch instructions first so model-specific guidance is read before - // any other context diffs on this turn. - build_model_instructions_update_item(previous_turn_settings, next), - build_permissions_update_item(previous, next, exec_policy), - build_collaboration_mode_update_item(previous, next), - build_multi_agent_mode_update_item(previous, next), - build_realtime_update_item(previous, previous_turn_settings, next), - build_personality_update_item(previous, next, personality_feature_enabled), - ] - .into_iter() - .flatten() - .collect(); - - let mut items = Vec::with_capacity(2); - if let Some(developer_message) = build_developer_update_item(developer_update_sections) { - items.push(developer_message); - } - if let Some(contextual_user_message) = contextual_user_message { - items.push(contextual_user_message); - } - items -} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 0792334b307f..d0a7849e3300 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -25,13 +25,12 @@ use crate::context::ApprovedCommandPrefixSaved; use crate::context::AppsInstructions; use crate::context::AvailablePluginsInstructions; use crate::context::AvailableSkillsInstructions; -use crate::context::CollaborationModeInstructions; use crate::context::ContextualUserFragment; -use crate::context::MultiAgentModeInstructions; use crate::context::NetworkRuleSaved; -use crate::context::PermissionsInstructions; -use crate::context::PersonalitySpecInstructions; use crate::context::RecommendedPluginsInstructions; +use crate::context::world_state::ModelState; +use crate::context::world_state::RealtimeState; +use crate::context::world_state::WorldStateSection; use crate::current_time::TimeProvider; use crate::default_skill_metadata_budget; use crate::environment_selection::TurnEnvironmentSnapshot; @@ -222,6 +221,7 @@ pub(crate) mod time_reminder; mod token_budget; pub(crate) mod turn; pub(crate) mod turn_context; +mod world_state; use self::config_lock::export_config_lock_if_configured; use self::config_lock::validate_config_lock_if_configured; #[cfg(test)] @@ -242,6 +242,7 @@ use self::turn::collect_explicit_app_ids_from_skill_items; use self::turn::realtime_text_for_event; use self::turn_context::TurnContext; use self::turn_context::TurnSkillsContext; +pub(crate) use self::world_state::build_world_state_from_turn_context_item; #[cfg(test)] mod rollout_reconstruction_tests; @@ -1381,10 +1382,10 @@ impl Session { } { let mut state = self.state.lock().await; + state.set_previous_turn_settings(previous_turn_settings.clone()); state.replace_history(history, reference_context_item); let window_id = window_id.unwrap_or_else(|| state.auto_compact_window_id()); state.restore_auto_compact_window(window_number, window_id); - state.set_previous_turn_settings(previous_turn_settings.clone()); } let prefix_tokens = if matches!( turn_context.config.model_auto_compact_token_limit_scope, @@ -1666,23 +1667,25 @@ impl Session { reference_context_item: Option<&TurnContextItem>, current_context: &TurnContext, ) -> Vec { - // TODO: Make context updates a pure diff of persisted previous/current TurnContextItem - // state so replay/backtracking is deterministic. Runtime inputs that affect model-visible - // context (shell, exec policy, feature gates, previous-turn bridge) should be persisted - // state or explicit non-state replay events. - let previous_turn_settings = { + let (previous_world_state, previous_turn_settings) = { let state = self.state.lock().await; - state.previous_turn_settings() + (state.world_state(), state.previous_turn_settings()) }; - let shell = self.user_shell(); - let exec_policy = self.services.exec_policy.current(); - crate::context_manager::updates::build_settings_update_items( - reference_context_item, - previous_turn_settings.as_ref(), - current_context, - shell.as_ref(), - exec_policy.as_ref(), - self.features.enabled(Feature::Personality), + let previous_world_state = previous_world_state.or_else(|| { + reference_context_item.map(|item| { + Arc::new(build_world_state_from_turn_context_item( + item, + previous_turn_settings.as_ref(), + )) + }) + }); + let Some(previous_world_state) = previous_world_state else { + return Vec::new(); + }; + + crate::context_manager::updates::merge_contextual_fragments( + self.build_world_state(current_context) + .render_diff(previous_world_state.as_ref()), ) } @@ -2873,6 +2876,9 @@ impl Session { { let mut state = self.state.lock().await; state.replace_history(items, reference_context_item.clone()); + if reference_context_item.is_some() { + state.set_world_state(Arc::new(self.build_world_state(turn_context))); + } } self.persist_rollout_items(&[RolloutItem::Compacted(compacted_item)]) @@ -3008,7 +3014,6 @@ impl Session { let ( reference_context_item, previous_turn_settings, - collaboration_mode, base_instructions, session_source, auto_compact_window_id, @@ -3017,41 +3022,24 @@ impl Session { ( state.reference_context_item(), state.previous_turn_settings(), - state.session_configuration.collaboration_mode.clone(), state.session_configuration.base_instructions.clone(), state.session_configuration.session_source.clone(), state.auto_compact_window_id(), ) }; - if let Some(model_switch_message) = - crate::context_manager::updates::build_model_instructions_update_item( - previous_turn_settings.as_ref(), - turn_context, - ) - { - developer_sections.push(model_switch_message); - } - if turn_context.config.include_permissions_instructions { - developer_sections.push( - PermissionsInstructions::from_permission_profile( - &turn_context.permission_profile, - turn_context.approval_policy.value(), - turn_context.config.approvals_reviewer, - self.services.exec_policy.current().as_ref(), - #[allow(deprecated)] - &turn_context.cwd, - turn_context - .config - .features - .enabled(Feature::ExecPermissionApprovals), - turn_context - .config - .features - .enabled(Feature::RequestPermissionsTool), - ) - .render(), - ); - } + let environment_subagents = if turn_context.config.include_environment_context { + self.services + .agent_control + .format_environment_context_subagents(self.thread_id) + .await + } else { + String::new() + }; + let initial_world_state = self.build_initial_world_state( + turn_context, + &environment_subagents, + &base_instructions, + ); let separate_guardian_developer_message = crate::guardian::is_guardian_reviewer_source(&session_source); // Keep the guardian policy prompt out of the aggregated developer bundle so it @@ -3062,37 +3050,6 @@ impl Session { { developer_sections.push(developer_instructions.to_string()); } - // Add developer instructions from collaboration_mode if they exist and are non-empty - if turn_context.config.include_collaboration_mode_instructions - && let Some(collab_instructions) = - CollaborationModeInstructions::from_collaboration_mode(&collaboration_mode) - { - developer_sections.push(collab_instructions.render()); - } - if let Some(realtime_update) = crate::context_manager::updates::build_initial_realtime_item( - reference_context_item.as_ref(), - previous_turn_settings.as_ref(), - turn_context, - ) { - developer_sections.push(realtime_update); - } - if self.features.enabled(Feature::Personality) - && let Some(personality) = turn_context.personality - { - let model_info = turn_context.model_info.clone(); - let has_baked_personality = model_info.supports_personality() - && base_instructions == model_info.get_model_instructions(Some(personality)); - if !has_baked_personality - && let Some(personality_message) = - crate::context_manager::updates::personality_message_for( - &model_info, - personality, - ) - { - developer_sections - .push(PersonalitySpecInstructions::new(personality_message).render()); - } - } if turn_context.config.include_apps_instructions && turn_context.apps_enabled() { let mcp_connection_manager = self.services.mcp_connection_manager.load_full(); let accessible_and_enabled_connectors = @@ -3216,24 +3173,35 @@ impl Session { .render(), ); } - if turn_context.config.include_environment_context { - let shell = self.user_shell(); - let subagents = self - .services - .agent_control - .format_environment_context_subagents(self.thread_id) - .await; - contextual_user_sections.push( - crate::context::EnvironmentContext::from_turn_context(turn_context, shell.as_ref()) - .with_subagents(subagents) - .render(), - ); - } - let multi_agent_v2_usage_hint_text = multi_agents::usage_hint_text(turn_context, &session_source); - let mut items = Vec::with_capacity(4); + let mut world_fragments = Vec::new(); + let previous_model_state = previous_turn_settings + .as_ref() + .map(|settings| ModelState::from_previous_model(&settings.model)); + if let Some(model_switch_fragment) = + ModelState::from_turn_context(turn_context).render_diff(previous_model_state.as_ref()) + { + world_fragments.push(model_switch_fragment); + } + let previous_realtime_state = RealtimeState::from_previous( + reference_context_item + .as_ref() + .and_then(|item| item.realtime_active), + previous_turn_settings + .as_ref() + .and_then(|settings| settings.realtime_active), + ); + if !turn_context.realtime_active + && let Some(realtime_fragment) = RealtimeState::from_turn_context(turn_context) + .render_diff(Some(&previous_realtime_state)) + { + world_fragments.push(realtime_fragment); + } + world_fragments.extend(initial_world_state.render_full()); + let mut items = + crate::context_manager::updates::merge_contextual_fragments(world_fragments); if let Some(developer_message) = crate::context_manager::updates::build_developer_update_item(developer_sections) { @@ -3254,20 +3222,6 @@ impl Session { { items.push(usage_hint_message); } - if let Some(multi_agent_mode) = multi_agents::effective_multi_agent_mode( - turn_context.multi_agent_version, - &turn_context.config.multi_agent_v2, - &session_source, - turn_context.multi_agent_mode, - turn_context - .config - .features - .enabled(Feature::MultiAgentMode), - ) { - items.push(ContextualUserFragment::into( - MultiAgentModeInstructions::new(multi_agent_mode), - )); - } if let Some(contextual_user_message) = crate::context_manager::updates::build_contextual_user_message(contextual_user_sections) { @@ -3335,6 +3289,7 @@ impl Session { { let mut state = self.state.lock().await; state.replace_history(replacement_history.clone(), Some(turn_context_item.clone())); + state.set_world_state(Arc::new(self.build_world_state(turn_context))); }; self.persist_rollout_items(&[ RolloutItem::Compacted(CompactedItem { @@ -3413,6 +3368,7 @@ impl Session { // context items. This keeps later runtime diffing aligned with the current turn state. let mut state = self.state.lock().await; state.set_reference_context_item(Some(turn_context_item)); + state.set_world_state(Arc::new(self.build_world_state(turn_context))); } pub(crate) async fn update_token_usage_info( diff --git a/codex-rs/core/src/session/snapshots/codex_core__codex_tests__fork_startup_context_then_first_turn_diff.snap b/codex-rs/core/src/session/snapshots/codex_core__codex_tests__fork_startup_context_then_first_turn_diff.snap index c51cb93ac60e..0396a748a143 100644 --- a/codex-rs/core/src/session/snapshots/codex_core__codex_tests__fork_startup_context_then_first_turn_diff.snap +++ b/codex-rs/core/src/session/snapshots/codex_core__codex_tests__fork_startup_context_then_first_turn_diff.snap @@ -7,8 +7,9 @@ Scenario: First request after fork when startup preserves the parent baseline, t ## First Forked Turn Request 00:message/developer: 01:message/user:> -02:message/user:fork seed -03:message/developer[2]: +02:message/developer: +03:message/user:fork seed +04:message/developer[2]: [01] [02] Fork turn collaboration instructions. -04:message/user:after fork +05:message/user:after fork diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index bca6be6a7441..e9792f9d886d 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -7325,6 +7325,26 @@ async fn spawn_task_does_not_update_previous_turn_settings_for_non_run_turn_task assert_eq!(sess.previous_turn_settings().await, None); } +#[tokio::test] +async fn build_settings_update_items_emits_nothing_for_matching_world_state() { + let (session, turn_context) = make_session_and_context().await; + let reference_context_item = turn_context.to_turn_context_item(); + session + .set_previous_turn_settings(Some(PreviousTurnSettings { + model: turn_context.model_info.slug.clone(), + comp_hash: turn_context.model_info.comp_hash.clone(), + realtime_active: Some(turn_context.realtime_active), + })) + .await; + + assert_eq!( + session + .build_settings_update_items(Some(&reference_context_item), &turn_context) + .await, + Vec::new() + ); +} + #[tokio::test] async fn build_settings_update_items_emits_environment_item_for_network_changes() { let (session, previous_context) = make_session_and_context().await; @@ -7388,46 +7408,38 @@ async fn build_settings_update_items_emits_environment_item_for_network_changes( } #[tokio::test] -async fn environment_context_uses_session_shell_when_environment_shell_is_absent() { - let (mut session, mut turn_context) = make_session_and_context().await; - session.services.user_shell = Arc::new(crate::shell::Shell { - shell_type: crate::shell::ShellType::PowerShell, - shell_path: PathBuf::from("powershell"), - }); - for environment in &mut turn_context.environments.turn_environments { - environment.shell = None; - } - - let session_shell = session.user_shell(); - let environment_context = crate::context::EnvironmentContext::from_turn_context( - &turn_context, - session_shell.as_ref(), - ) - .render(); - assert!( - environment_context.contains("powershell"), - "{environment_context}" +async fn build_settings_update_items_emits_environment_item_for_cwd_changes() { + let (session, previous_context) = make_session_and_context().await; + let previous_context = Arc::new(previous_context); + let mut current_context = previous_context + .with_model( + previous_context.model_info.slug.clone(), + &session.services.models_manager, + ) + .await; + let cwd = test_path_buf("/new-repo").abs(); + let environment = current_context.environments.turn_environments[0].clone(); + current_context.environments.turn_environments[0] = TurnEnvironment::new( + environment.environment_id, + environment.environment, + PathUri::from_abs_path(&cwd), + environment.shell, ); - let primary_environment = turn_context - .environments - .turn_environments - .first_mut() - .expect("primary environment"); - primary_environment.shell = Some(crate::shell::Shell { - shell_type: crate::shell::ShellType::Cmd, - shell_path: PathBuf::from("cmd"), - }); + let reference_context_item = previous_context.to_turn_context_item(); + let update_items = session + .build_settings_update_items(Some(&reference_context_item), ¤t_context) + .await; - let environment_context = crate::context::EnvironmentContext::from_turn_context( - &turn_context, - session_shell.as_ref(), - ) - .render(); + let environment_update = user_input_texts(&update_items) + .into_iter() + .find(|text| text.contains("")) + .expect("environment update item should be emitted"); assert!( - environment_context.contains("cmd"), - "{environment_context}" + environment_update.contains(&format!("{}", cwd.display())), + "{environment_update}" ); + assert!(!environment_update.contains("")); } #[tokio::test] @@ -7469,7 +7481,13 @@ async fn build_settings_update_items_omits_environment_item_when_disabled() { let mut config = (*current_context.config).clone(); config.include_environment_context = false; current_context.config = Arc::new(config); - current_context.current_date = Some("2026-02-27".to_string()); + let environment = current_context.environments.turn_environments[0].clone(); + current_context.environments.turn_environments[0] = TurnEnvironment::new( + environment.environment_id, + environment.environment, + PathUri::from_abs_path(&test_path_buf("/new-repo").abs()), + environment.shell, + ); let reference_context_item = previous_context.to_turn_context_item(); let update_items = session @@ -7576,7 +7594,7 @@ async fn build_settings_update_items_uses_previous_turn_settings_for_realtime_en } #[tokio::test] -async fn build_initial_context_uses_previous_realtime_state() { +async fn build_initial_context_renders_full_realtime_state() { let (session, mut turn_context) = make_session_and_context().await; turn_context.realtime_active = true; @@ -7597,10 +7615,10 @@ async fn build_initial_context_uses_previous_realtime_state() { let resumed_context = session.build_initial_context(&turn_context).await; let resumed_developer_texts = developer_input_texts(&resumed_context); assert!( - !resumed_developer_texts + resumed_developer_texts .iter() .any(|text| text.contains("")), - "did not expect a duplicate realtime update, got {resumed_developer_texts:?}" + "expected full initial context to restate active realtime state, got {resumed_developer_texts:?}" ); } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 28c0f6561293..8cd627c3344f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -240,6 +240,8 @@ pub(crate) async fn run_turn( ) .await?; + sess.record_world_state_diff(turn_context.as_ref()).await; + // Construct the input that we will send to the model. let sampling_request_input: Vec = async { sess.clone_history() diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs new file mode 100644 index 000000000000..769911f7717a --- /dev/null +++ b/codex-rs/core/src/session/world_state.rs @@ -0,0 +1,137 @@ +use super::PreviousTurnSettings; +use super::Session; +use super::turn_context::TurnContext; +use crate::context::EnvironmentsState; +use crate::context::world_state::CollaborationModeState; +use crate::context::world_state::ModelState; +use crate::context::world_state::MultiAgentModeState; +use crate::context::world_state::PermissionsState; +use crate::context::world_state::PersonalityState; +use crate::context::world_state::RealtimeState; +use crate::context::world_state::WorldState; +use crate::environment_selection::TurnEnvironmentSnapshot; +use codex_execpolicy::Policy; +use codex_features::Feature; +use codex_protocol::protocol::TurnContextItem; +use std::sync::Arc; + +fn build_world_state_from_turn_context_with_environments( + turn_context: &TurnContext, + environments: &TurnEnvironmentSnapshot, + exec_policy: &Policy, + personality_feature_enabled: bool, + environment_subagents: &str, +) -> WorldState { + let mut world_state = WorldState::default(); + world_state.add_section(ModelState::from_turn_context(turn_context)); + world_state.add_section(PermissionsState::from_turn_context( + turn_context, + exec_policy, + )); + world_state.add_section(CollaborationModeState::from_turn_context(turn_context)); + world_state.add_section(MultiAgentModeState::from_turn_context(turn_context)); + world_state.add_section(RealtimeState::from_turn_context(turn_context)); + world_state.add_section(PersonalityState::from_turn_context( + turn_context, + personality_feature_enabled, + )); + if turn_context.config.include_environment_context { + world_state.add_section( + EnvironmentsState::from_turn_context_with_environments(turn_context, environments) + .with_subagents(environment_subagents.to_string()), + ); + } + world_state +} + +pub(crate) fn build_world_state_from_turn_context_item( + turn_context_item: &TurnContextItem, + previous_turn_settings: Option<&PreviousTurnSettings>, +) -> WorldState { + let mut world_state = WorldState::default(); + if let Some(previous_turn_settings) = previous_turn_settings { + world_state.add_section(ModelState::from_previous_model( + &previous_turn_settings.model, + )); + } + world_state.add_section(PermissionsState::from_turn_context_item(turn_context_item)); + world_state.add_section(CollaborationModeState::from_turn_context_item( + turn_context_item, + )); + world_state.add_section(MultiAgentModeState::from_turn_context_item( + turn_context_item, + )); + world_state.add_section(RealtimeState::from_previous( + turn_context_item.realtime_active, + previous_turn_settings.and_then(|settings| settings.realtime_active), + )); + world_state.add_section(PersonalityState::from_turn_context_item(turn_context_item)); + world_state.add_section(EnvironmentsState::from_turn_context_item(turn_context_item)); + world_state +} + +impl Session { + pub(super) fn build_world_state(&self, turn_context: &TurnContext) -> WorldState { + let exec_policy = self.services.exec_policy.current(); + build_world_state_from_turn_context_with_environments( + turn_context, + &turn_context.environments, + exec_policy.as_ref(), + self.features.enabled(Feature::Personality), + "", + ) + } + + pub(super) fn build_initial_world_state( + &self, + turn_context: &TurnContext, + environment_subagents: &str, + base_instructions: &str, + ) -> WorldState { + let personality_is_baked = turn_context.personality.is_some_and(|personality| { + turn_context.model_info.supports_personality() + && base_instructions + == turn_context + .model_info + .get_model_instructions(Some(personality)) + }); + let exec_policy = self.services.exec_policy.current(); + build_world_state_from_turn_context_with_environments( + turn_context, + &turn_context.environments, + exec_policy.as_ref(), + self.features.enabled(Feature::Personality) && !personality_is_baked, + environment_subagents, + ) + } + + async fn build_live_world_state(&self, turn_context: &TurnContext) -> WorldState { + let environments = self.services.turn_environments.snapshot().await; + let exec_policy = self.services.exec_policy.current(); + build_world_state_from_turn_context_with_environments( + turn_context, + &environments, + exec_policy.as_ref(), + self.features.enabled(Feature::Personality), + "", + ) + } + + pub(crate) async fn record_world_state_diff(&self, turn_context: &TurnContext) { + let world_state = Arc::new(self.build_live_world_state(turn_context).await); + let previous = { + let state = self.state.lock().await; + state.world_state() + }; + let fragments = match previous.as_deref() { + Some(previous) => world_state.render_diff(previous), + None => world_state.render_full(), + }; + let items = crate::context_manager::updates::merge_contextual_fragments(fragments); + if !items.is_empty() { + self.record_conversation_items(turn_context, &items).await; + } + let mut state = self.state.lock().await; + state.set_world_state(world_state); + } +} diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 83cef227f823..c6f5b1e3e77c 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -6,13 +6,16 @@ use codex_sandboxing::policy_transforms::merge_permission_profiles; use std::collections::HashMap; use std::collections::HashSet; use std::collections::VecDeque; +use std::sync::Arc; use uuid::Uuid; use super::AdditionalContextStore; use super::auto_compact_window::AutoCompactWindow; use super::auto_compact_window::AutoCompactWindowSnapshot; +use crate::context::world_state::WorldState; use crate::context_manager::ContextManager; use crate::session::PreviousTurnSettings; +use crate::session::build_world_state_from_turn_context_item; use crate::session::session::SessionConfiguration; use crate::session::time_reminder::CurrentTimeReminderState; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; @@ -105,9 +108,15 @@ impl SessionState { items: Vec, reference_context_item: Option, ) { + let world_state = reference_context_item.as_ref().map(|item| { + build_world_state_from_turn_context_item(item, self.previous_turn_settings.as_ref()) + }); self.history.replace(items); self.history .set_reference_context_item(reference_context_item); + if let Some(world_state) = world_state { + self.history.set_world_state(Arc::new(world_state)); + } self.auto_compact_window.clear_prefill(); } @@ -116,13 +125,27 @@ impl SessionState { } pub(crate) fn set_reference_context_item(&mut self, item: Option) { + let world_state = item.as_ref().map(|item| { + build_world_state_from_turn_context_item(item, self.previous_turn_settings.as_ref()) + }); self.history.set_reference_context_item(item); + if let Some(world_state) = world_state { + self.history.set_world_state(Arc::new(world_state)); + } } pub(crate) fn reference_context_item(&self) -> Option { self.history.reference_context_item() } + pub(crate) fn set_world_state(&mut self, world_state: Arc) { + self.history.set_world_state(world_state); + } + + pub(crate) fn world_state(&self) -> Option> { + self.history.world_state() + } + // Token/rate limit helpers pub(crate) fn update_token_info_from_usage( &mut self, diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index 269beae63833..dffae023235d 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -479,6 +479,65 @@ fn assert_two_responses_input_snapshot(snapshot_name: &str, requests: &[Vec] insta::assert_snapshot!(snapshot_name, snapshot); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn turn_iteration_injects_environment_diff_after_thread_environment_changes() { + let (gate_first_completed_tx, gate_first_completed_rx) = oneshot::channel(); + let first_chunks = vec![ + chunk(ev_response_created("resp-1")), + chunk(ev_message_item_done("msg-1", "first answer")), + gated_chunk(gate_first_completed_rx, vec![ev_completed("resp-1")]), + ]; + let second_chunks = vec![ + chunk(ev_response_created("resp-2")), + chunk(ev_message_item_done("msg-2", "second answer")), + chunk(ev_completed("resp-2")), + ]; + let (server, _completions) = + start_streaming_sse_server(vec![first_chunks, second_chunks]).await; + let test = test_codex() + .with_model("gpt-5.4") + .build_with_streaming_server(&server) + .await + .expect("build streaming Codex test session"); + + submit_user_input(&test.codex, "first prompt").await; + wait_for_agent_message(&test.codex, "first answer").await; + + let next_cwd = test.config.codex_home.clone(); + test.codex + .submit(Op::ThreadSettings { + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + environments: Some(local_selections(next_cwd.clone())), + ..Default::default() + }, + }) + .await + .expect("update thread environment"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::ThreadSettingsApplied(_)) + }) + .await; + steer_user_input(&test.codex, "second prompt").await; + let _ = gate_first_completed_tx.send(()); + + wait_for_agent_message(&test.codex, "second answer").await; + wait_for_turn_complete(&test.codex).await; + + let requests = server.requests().await; + assert_eq!(requests.len(), 2); + let second_request: Value = from_slice(&requests[1]).expect("parse second request"); + let user_texts = message_input_texts(&second_request, "user"); + let expected_cwd = format!("{}", next_cwd.display()); + let environment_update = user_texts + .iter() + .find(|text| text.starts_with("") && text.contains(&expected_cwd)) + .expect("second iteration should include the changed environment cwd"); + assert!(!environment_update.contains("")); + assert!(user_texts.iter().any(|text| text == "second prompt")); + + server.shutdown().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[ignore = "TODO(aibrahim): flaky"] async fn injected_user_input_triggers_follow_up_request_with_deltas() { diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_resume_with_personality_change.snap b/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_resume_with_personality_change.snap index 42d1cd1a9f40..305d9b5d30a5 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_resume_with_personality_change.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_resume_with_personality_change.snap @@ -5,21 +5,19 @@ expression: "format_labeled_requests_snapshot(\"First post-resume turn where res Scenario: First post-resume turn where resumed config model differs from rollout and personality changes. ## Last Request Before Resume -00:message/developer[2]: - [01] - [02] +00:message/developer: 01:message/user:> -02:message/user:seed resume history +02:message/developer: +03:message/user:seed resume history ## First Request After Resume -00:message/developer[2]: - [01] - [02] +00:message/developer: 01:message/user:> -02:message/user:seed resume history -03:message/assistant:recorded before resume -04:message/developer[2]: +02:message/developer: +03:message/user:seed resume history +04:message/assistant:recorded before resume +05:message/developer[2]: [01] \nThe user was previously using a different model. Please continue the conversatio... [02] -05:message/user: -06:message/user:resume and change personality +06:message/user: +07:message/user:resume and change personality diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_turn_overrides.snap b/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_turn_overrides.snap index 8e66e3314cc2..b4f8d53add9e 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_turn_overrides.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_turn_overrides.snap @@ -5,21 +5,19 @@ expression: "format_labeled_requests_snapshot(\"Second turn changes cwd, approva Scenario: Second turn changes cwd, approval policy, and personality while keeping model constant. ## First Request (Baseline) -00:message/developer[2]: - [01] - [02] +00:message/developer: 01:message/user:> -02:message/user:first turn +02:message/developer: +03:message/user:first turn ## Second Request (Turn Overrides) -00:message/developer[2]: - [01] - [02] +00:message/developer: 01:message/user:> -02:message/user:first turn -03:message/assistant:turn one complete -04:message/developer[2]: +02:message/developer: +03:message/user:first turn +04:message/assistant:turn one complete +05:message/developer[2]: [01] [02] The user has requested a new communication style. Future messages should adhe... -05:message/user: -06:message/user:second turn with context updates +06:message/user: +07:message/user:second turn with context updates