diff --git a/codex-rs/core/src/context/contextual_user_message.rs b/codex-rs/core/src/context/contextual_user_message.rs index de735b6f1..390a20fd4 100644 --- a/codex-rs/core/src/context/contextual_user_message.rs +++ b/codex-rs/core/src/context/contextual_user_message.rs @@ -3,7 +3,6 @@ use codex_protocol::items::parse_hook_prompt_fragment; use codex_protocol::models::ContentItem; use super::AdditionalContextUserFragment; -use super::EnvironmentContext; use super::FragmentRegistration; use super::FragmentRegistrationProxy; use super::InternalModelContextFragment; @@ -16,10 +15,11 @@ use super::SubagentNotification; use super::TurnAborted; use super::UserInstructions; use super::UserShellCommand; +use super::world_state::EnvironmentsState; 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 5f36124b8..2f795a8ea 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -1,119 +1,13 @@ -use crate::session::step_context::StepContext; -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 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, - status: EnvironmentContextEnvironmentStatus, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -enum EnvironmentContextEnvironmentStatus { - Starting, - Ready { shell: String }, -} - -impl EnvironmentContextEnvironment { - fn from_turn_environments(environments: &[TurnEnvironment], shell: &Shell) -> Vec { - environments - .iter() - .map(|environment| Self { - id: environment.environment_id.clone(), - cwd: environment.cwd().clone(), - status: EnvironmentContextEnvironmentStatus::Ready { - shell: environment - .shell - .as_ref() - .map(|shell| shell.name().to_string()) - .unwrap_or_else(|| shell.name().to_string()), - }, - }) - .collect() - } -} - -impl EnvironmentContextEnvironmentStatus { - fn equals_except_shell(&self, other: &Self) -> bool { - matches!( - (self, other), - (Self::Starting, Self::Starting) | (Self::Ready { .. }, Self::Ready { .. }) - ) - } - - fn render(&self, indent: &str) -> String { - match self { - Self::Starting => format!("{indent}starting"), - Self::Ready { shell } => format!("{indent}{shell}"), - } - } -} - -#[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 && left.status.equals_except_shell(&right.status) - } - (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 - && left.status.equals_except_shell(&right.status) - }) - } - _ => false, - } - } -} - #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct FileSystemContext { workspace_roots: Vec, @@ -137,7 +31,7 @@ enum ManagedFileSystemContext { } impl FileSystemContext { - fn from_permission_profile( + pub(super) fn from_permission_profile( permission_profile: &PermissionProfile, workspace_roots: &[AbsolutePathBuf], ) -> Self { @@ -163,7 +57,7 @@ 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(""); @@ -303,7 +197,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("&"), @@ -330,7 +224,7 @@ impl NetworkContext { } } - 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); @@ -348,302 +242,3 @@ impl NetworkContext { 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, - } - } - - pub(crate) fn from_step_context(step_context: &StepContext, shell: &Shell) -> Option { - let mut environments = EnvironmentContextEnvironment::from_turn_environments( - &step_context.environments.turn_environments, - shell, - ); - environments.extend( - step_context - .environments - .starting - .iter() - .map(|environment| EnvironmentContextEnvironment { - id: environment.selection.environment_id.clone(), - cwd: environment.selection.cwd.clone(), - status: EnvironmentContextEnvironmentStatus::Starting, - }), - ); - - Self::environment_only(environments) - } - - pub(crate) fn from_attached_environments( - environments: &[TurnEnvironment], - shell: &Shell, - ) -> Option { - Self::environment_only(EnvironmentContextEnvironment::from_turn_environments( - environments, - shell, - )) - } - - fn environment_only(environments: Vec) -> Option { - (!environments.is_empty()).then(|| { - Self::new( - environments, - /*current_date*/ None, - /*timezone*/ None, - /*network*/ None, - /*subagents*/ None, - ) - }) - } - - /// 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 { - id: String::new(), - cwd: environment.cwd.clone(), - status: environment.status.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 { - id: String::new(), - cwd: PathUri::from_abs_path(&turn_context_item.cwd), - status: EnvironmentContextEnvironmentStatus::Ready { 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(environment.status.render(" ")); - } - 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(environment.status.render(" ")); - 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()); - } - format!("\n{}\n", lines.join("\n")) - } -} - -#[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 c9319ebef..4ba6acf02 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -32,6 +32,7 @@ mod token_budget_context; mod turn_aborted; mod user_instructions; mod user_shell_command; +pub(crate) mod world_state; pub(crate) use approved_command_prefix_saved::ApprovedCommandPrefixSaved; pub(crate) use apps_instructions::AppsInstructions; @@ -47,7 +48,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; 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 000000000..589e8dbad --- /dev/null +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -0,0 +1,387 @@ +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 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 { + Self { + environments: environment_states(environments), + 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, + } + } + + 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: 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 + } + + 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) + .is_none_or(|previous| !environment.has_same_diff_value(previous)) + }) + .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 environment.status == EnvironmentStatus::Starting { + rendered.push_str(indent); + rendered.push_str("starting\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, PartialEq, Eq)] +struct EnvironmentState { + cwd: PathUri, + status: EnvironmentStatus, + shell: Option, +} + +impl EnvironmentState { + fn has_same_diff_value(&self, other: &Self) -> bool { + self.cwd == other.cwd + && self.status == other.status + && self + .shell + .as_ref() + .zip(other.shell.as_ref()) + .is_none_or(|(current, previous)| current == previous) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum EnvironmentStatus { + Starting, + Available, +} + +fn environment_states(snapshot: &TurnEnvironmentSnapshot) -> BTreeMap { + let mut environments = snapshot + .turn_environments + .iter() + .map(|environment| { + ( + environment.environment_id.clone(), + EnvironmentState { + cwd: environment.cwd().clone(), + status: EnvironmentStatus::Available, + shell: environment + .shell + .as_ref() + .map(|shell| shell.name().to_string()), + }, + ) + }) + .collect::>(); + for environment in &snapshot.starting { + environments + .entry(environment.selection.environment_id.clone()) + .or_insert_with(|| EnvironmentState { + cwd: environment.selection.cwd.clone(), + status: EnvironmentStatus::Starting, + shell: None, + }); + } + environments +} + +fn is_legacy_single(environments: &BTreeMap) -> bool { + environments.len() == 1 + && environments + .values() + .all(|environment| environment.status == 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 64% 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 4d9f1a729..be127de55 100644 --- a/codex-rs/core/src/context/environment_context_tests.rs +++ b/codex-rs/core/src/context/world_state/environment_render_tests.rs @@ -1,6 +1,5 @@ use crate::shell::ShellType; -use super::EnvironmentContextEnvironmentStatus::Ready; use super::*; use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::FileSystemAccessMode; @@ -31,17 +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: 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()), - status: Ready { - 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, @@ -63,14 +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"), - status: Ready { - 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, @@ -92,14 +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")), - status: Ready { - 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), @@ -156,14 +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")), - status: Ready { - 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, @@ -219,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!( @@ -246,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()), @@ -262,111 +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")), - status: Ready { - 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")), - status: Ready { - 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")), - status: Ready { - 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")), - status: Ready { - 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")), - status: Ready { - 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")), - status: Ready { - 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")), - status: Ready { - 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, @@ -394,22 +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()), - status: Ready { - shell: "bash".to_string(), - }, - }, - EnvironmentContextEnvironment { - id: "remote".to_string(), - cwd: PathUri::from_abs_path(&remote_cwd.abs()), - status: Ready { - 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()), @@ -443,22 +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()), - status: Ready { - shell: "powershell".to_string(), - }, - }, - EnvironmentContextEnvironment { - id: "remote".to_string(), - cwd: PathUri::from_abs_path(&remote_cwd.abs()), - status: Ready { - 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 000000000..77a3adde1 --- /dev/null +++ b/codex-rs/core/src/context/world_state/environment_tests.rs @@ -0,0 +1,256 @@ +use super::*; +use crate::context::ContextualUserFragment; +use crate::context::world_state::WorldState; +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(), + ..Default::default() + }; + + let mut world_state = WorldState::default(); + world_state.add_section(context); + + assert_eq!( + vec![user_message( + r#" + + + /workspace + bash + + + /repo + zsh + + +"#, + )], + render_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:///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#" + + + /workspace + powershell + + + /repo + zsh + + + + /remote + starting + + +"#, + )], + render_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 + +"#, + )], + render_fragments(current.render_diff(&previous)), + ); + Ok(()) +} + +#[test] +fn single_environment_diff_ignores_unknown_shell() -> Result<()> { + let previous = EnvironmentsState { + environments: [( + LOCAL_ENVIRONMENT_ID.to_string(), + EnvironmentState { + cwd: PathUri::parse("file:///repo")?, + status: EnvironmentStatus::Available, + shell: None, + }, + )] + .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, + render_fragment(WorldStateSection::render_diff(¤t, 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#" + + + +"#, + )), + render_fragment(WorldStateSection::render_diff( + &EnvironmentsState::default(), + Some(&previous), + )), + ); + Ok(()) +} + +fn available(cwd: &str, shell: &str) -> Result { + Ok(EnvironmentState { + cwd: PathUri::parse(cwd)?, + status: EnvironmentStatus::Available, + shell: Some(shell.to_string()), + }) +} + +fn starting(cwd: &str) -> Result { + Ok(EnvironmentState { + cwd: PathUri::parse(cwd)?, + status: EnvironmentStatus::Starting, + shell: None, + }) +} + +fn render_fragments(fragments: Vec>) -> Vec { + fragments + .into_iter() + .map(ContextualUserFragment::into_boxed_response_item) + .collect() +} + +fn render_fragment(fragment: Option>) -> Option { + fragment.map(ContextualUserFragment::into_boxed_response_item) +} + +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, + } +} 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 000000000..06f750ef8 --- /dev/null +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -0,0 +1,80 @@ +mod environment; + +use crate::context::ContextualUserFragment; +use indexmap::IndexMap; +use std::any::Any; +use std::any::TypeId; +use std::fmt; + +pub(crate) use environment::EnvironmentsState; + +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. +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, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.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_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index a639aa51b..0244f7f96 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -1,4 +1,5 @@ -use crate::context::EnvironmentContext; +use crate::context::ContextualUserFragment; +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; @@ -28,6 +29,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 @@ -49,8 +51,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, - /// Environment state most recently appended to model-visible history. - environment_context_baseline: Option, + /// World state most recently appended to model-visible history. + world_state_baseline: Option>, } impl ContextManager { @@ -62,7 +64,7 @@ impl ContextManager { &None, &None, /*model_context_window*/ None, ), reference_context_item: None, - environment_context_baseline: None, + world_state_baseline: None, } } @@ -82,15 +84,20 @@ impl ContextManager { self.reference_context_item.clone() } - pub(crate) fn update_environment_context_baseline( + pub(crate) fn update_world_state( &mut self, - context: &EnvironmentContext, - ) -> bool { - if self.environment_context_baseline.as_ref() == Some(context) { - return false; - } - self.environment_context_baseline = Some(context.clone()); - true + world_state: WorldState, + ) -> Vec> { + let fragments = self.world_state_baseline.as_deref().map_or_else( + || world_state.render_full(), + |previous| world_state.render_diff(previous), + ); + self.world_state_baseline = Some(Arc::new(world_state)); + fragments + } + + pub(crate) fn set_world_state_baseline(&mut self, world_state: WorldState) { + self.world_state_baseline = Some(Arc::new(world_state)); } pub(crate) fn set_token_usage_full(&mut self, context_window: i64) { @@ -178,14 +185,14 @@ impl ContextManager { // its corresponding counterpart to keep the invariants intact without // running a full normalization pass. normalize::remove_corresponding_for(&mut self.items, &removed); - self.environment_context_baseline = None; + self.world_state_baseline = None; } } pub(crate) fn replace(&mut self, items: Vec) { self.items = items; self.history_version = self.history_version.saturating_add(1); - self.environment_context_baseline = None; + self.world_state_baseline = None; } /// Replace image content in the last turn if it originated from a tool output. diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index b388a1a06..da2f14fd4 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -1,5 +1,6 @@ use super::*; -use crate::context::EnvironmentContext; +use crate::context::world_state::EnvironmentsState; +use crate::context::world_state::WorldState; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_protocol::AgentPath; @@ -75,17 +76,22 @@ fn create_history_with_items(items: Vec) -> ContextManager { } #[test] -fn environment_context_baseline_deduplicates_until_history_is_replaced() { - let context = - EnvironmentContext::from_turn_context_item(&reference_context_item(), "bash".to_string()); +fn world_state_baseline_deduplicates_until_history_is_replaced() { + let world_state = || { + let mut state = WorldState::default(); + state.add_section(EnvironmentsState::from_turn_context_item( + &reference_context_item(), + )); + state + }; let mut history = ContextManager::new(); - assert!(history.update_environment_context_baseline(&context)); - assert!(!history.update_environment_context_baseline(&context)); + assert_eq!(1, history.update_world_state(world_state()).len()); + assert!(history.update_world_state(world_state()).is_empty()); history.replace(Vec::new()); - assert!(history.update_environment_context_baseline(&context)); + assert_eq!(1, history.update_world_state(world_state()).len()); } fn user_msg(text: &str) -> ResponseItem { diff --git a/codex-rs/core/src/context_manager/updates.rs b/codex-rs/core/src/context_manager/updates.rs index b83c1b985..937c4bf05 100644 --- a/codex-rs/core/src/context_manager/updates.rs +++ b/codex-rs/core/src/context_manager/updates.rs @@ -1,6 +1,5 @@ use crate::context::CollaborationModeInstructions; use crate::context::ContextualUserFragment; -use crate::context::EnvironmentContext; use crate::context::ModelSwitchInstructions; use crate::context::MultiAgentModeInstructions; use crate::context::PermissionsInstructions; @@ -10,7 +9,6 @@ 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; @@ -20,27 +18,6 @@ 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, @@ -220,6 +197,26 @@ pub(crate) fn build_contextual_user_message(text_sections: Vec) -> Optio build_text_message("user", text_sections) } +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])), + } + } + messages + .into_iter() + .filter_map(|(role, text_sections)| build_text_message(role, text_sections)) + .collect() +} + fn build_text_message(role: &str, text_sections: Vec) -> Option { if text_sections.is_empty() { return None; @@ -243,7 +240,6 @@ pub(crate) fn build_settings_update_items( previous: Option<&TurnContextItem>, previous_turn_settings: Option<&PreviousTurnSettings>, next: &TurnContext, - shell: &Shell, exec_policy: &Policy, personality_feature_enabled: bool, ) -> Vec { @@ -251,7 +247,6 @@ pub(crate) fn build_settings_update_items( // 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. @@ -266,12 +261,7 @@ pub(crate) fn build_settings_update_items( .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 + build_developer_update_item(developer_update_sections) + .into_iter() + .collect() } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 32fcfd09e..4a7c7a9e4 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -32,6 +32,7 @@ use crate::context::NetworkRuleSaved; use crate::context::PermissionsInstructions; use crate::context::PersonalitySpecInstructions; use crate::context::RecommendedPluginsInstructions; +use crate::context::world_state::WorldState; use crate::current_time::TimeProvider; use crate::default_skill_metadata_budget; use crate::environment_selection::TurnEnvironmentSnapshot; @@ -223,6 +224,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)] @@ -243,6 +245,9 @@ 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; +use self::world_state::build_world_state_from_environment_snapshot; +use self::world_state::build_world_state_from_turn_context; +use self::world_state::build_world_state_from_turn_context_item; #[cfg(test)] mod rollout_reconstruction_tests; @@ -1377,9 +1382,15 @@ impl Session { // will be processed again if the rollout is reconstructed in a future session. // This meets image resizing requirements without modifying persisted rollouts. prepare_response_items(&mut history); + let world_state_baseline = reference_context_item + .as_ref() + .map(build_world_state_from_turn_context_item); { let mut state = self.state.lock().await; state.replace_history(history, reference_context_item); + if let Some(world_state) = world_state_baseline { + state.history.set_world_state_baseline(world_state); + } let fallback_ids = state.auto_compact_window_ids(); let window_id = window_id.unwrap_or(fallback_ids.window_id); state.restore_auto_compact_window( @@ -1674,19 +1685,17 @@ impl Session { ) -> 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 + // context (exec policy, feature gates, previous-turn bridge) should be persisted // state or explicit non-state replay events. let previous_turn_settings = { let state = self.state.lock().await; 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), ) @@ -2757,24 +2766,19 @@ impl Session { return; } - let shell = self.user_shell(); - let Some(environment_context) = - crate::context::EnvironmentContext::from_step_context(step_context, shell.as_ref()) - else { - return; - }; - let changed = { + let world_state = + build_world_state_from_environment_snapshot(turn_context, &step_context.environments); + let items = { let mut state = self.state.lock().await; - state - .history - .update_environment_context_baseline(&environment_context) + crate::context_manager::updates::merge_contextual_fragments( + state.history.update_world_state(world_state), + ) }; - if !changed { + if items.is_empty() { return; } - let item = ContextualUserFragment::into(environment_context); - self.record_conversation_items(turn_context, &[item]).await; + self.record_conversation_items(turn_context, &items).await; } pub(crate) async fn record_inter_agent_communication( @@ -2886,9 +2890,17 @@ impl Session { replacement_history: Some(items.clone()), ..compacted_item }; + let world_state_baseline = if reference_context_item.is_some() { + Some(self.build_world_state(turn_context).await) + } else { + None + }; { let mut state = self.state.lock().await; state.replace_history(items, reference_context_item.clone()); + if let Some(world_state) = world_state_baseline { + state.history.set_world_state_baseline(world_state); + } } self.persist_rollout_items(&[RolloutItem::Compacted(compacted_item)]) @@ -3017,6 +3029,28 @@ impl Session { pub(crate) async fn build_initial_context( &self, turn_context: &TurnContext, + ) -> Vec { + let world_state = self.build_world_state(turn_context).await; + self.build_initial_context_with_world_state(turn_context, &world_state) + .await + } + + async fn build_world_state(&self, turn_context: &TurnContext) -> WorldState { + 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() + }; + build_world_state_from_turn_context(turn_context, &environment_subagents) + } + + async fn build_initial_context_with_world_state( + &self, + turn_context: &TurnContext, + world_state: &WorldState, ) -> Vec { let mut developer_sections = Vec::::with_capacity(8); let mut contextual_user_sections = Vec::::with_capacity(2); @@ -3257,18 +3291,12 @@ 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(), - ); + for fragment in world_state.render_full() { + match fragment.role() { + "developer" => developer_sections.push(fragment.render()), + "user" => contextual_user_sections.push(fragment.render()), + _ => {} + } } let multi_agent_v2_usage_hint_text = @@ -3370,12 +3398,16 @@ impl Session { state.start_new_context_window_if_requested() }; let (window_number, window_ids) = window?; - let context_items = self.build_initial_context(turn_context).await; + let world_state = self.build_world_state(turn_context).await; + let context_items = self + .build_initial_context_with_world_state(turn_context, &world_state) + .await; let turn_context_item = turn_context.to_turn_context_item(); let replacement_history = context_items; { let mut state = self.state.lock().await; state.replace_history(replacement_history.clone(), Some(turn_context_item.clone())); + state.history.set_world_state_baseline(world_state); }; self.persist_rollout_items(&[ RolloutItem::Compacted(CompactedItem { @@ -3406,15 +3438,15 @@ impl Session { /// steady-state turns that emit model-visible context updates. /// /// When the reference snapshot is missing, this injects full initial context. Otherwise, it - /// emits only settings diff items. + /// emits only context diffs. /// /// If full context is injected and a model switch occurred, this prepends the /// `` developer message so model-specific instructions are not lost. /// /// This is the normal runtime path that establishes a new `reference_context_item`. - /// Mid-turn compaction is the other path that can re-establish that baseline when it - /// reinjects full initial context into replacement history. Other non-regular tasks - /// intentionally do not update the baseline. + /// Mid-turn compaction is the other path that can re-establish that reference when it + /// reinjects full initial context into replacement history. Live world-state changes may + /// independently advance their in-memory baseline within a turn. #[instrument(level = "trace", skip_all)] pub(crate) async fn record_context_updates_and_set_reference_context_item( &self, @@ -3425,40 +3457,43 @@ impl Session { state.reference_context_item() }; let turn_context_item = turn_context.to_turn_context_item(); - if reference_context_item.as_ref() == Some(&turn_context_item) { - return; - } + let turn_context_changed = reference_context_item.as_ref() != Some(&turn_context_item); let should_inject_full_context = reference_context_item.is_none(); + let world_state = self.build_world_state(turn_context).await; let mut context_items = if should_inject_full_context { - self.build_initial_context(turn_context).await + let context_items = self + .build_initial_context_with_world_state(turn_context, &world_state) + .await; + self.state + .lock() + .await + .history + .set_world_state_baseline(world_state); + context_items } else { // Steady-state path: append only built-in context diffs here; turn-scoped extension // context is added below. - self.build_settings_update_items(reference_context_item.as_ref(), turn_context) - .await + let mut context_items = self + .build_settings_update_items(reference_context_item.as_ref(), turn_context) + .await; + let world_state_items = { + let mut state = self.state.lock().await; + crate::context_manager::updates::merge_contextual_fragments( + state.history.update_world_state(world_state), + ) + }; + context_items.extend(world_state_items); + context_items }; - if !should_inject_full_context { + if !should_inject_full_context && turn_context_changed { context_items.extend( self.build_turn_context_contribution_items(turn_context) .await, ); } - let initial_environment_context = if should_inject_full_context - && !context_items.is_empty() - && turn_context.config.include_environment_context - && turn_context - .config - .features - .enabled(Feature::DeferredExecutor) - { - let shell = self.user_shell(); - crate::context::EnvironmentContext::from_attached_environments( - &turn_context.environments.turn_environments, - shell.as_ref(), - ) - } else { - None - }; + if !turn_context_changed && context_items.is_empty() { + return; + } if !context_items.is_empty() { self.record_conversation_items(turn_context, &context_items) .await; @@ -3468,15 +3503,10 @@ impl Session { self.persist_rollout_items(&[RolloutItem::TurnContext(turn_context_item.clone())]) .await; - // Advance the in-memory diff baseline even when this turn emitted no model-visible - // context items. This keeps later runtime diffing aligned with the current turn state. + // Advance the persisted-settings baseline even when this turn emitted no model-visible + // context items. let mut state = self.state.lock().await; state.set_reference_context_item(Some(turn_context_item)); - if let Some(environment_context) = initial_environment_context { - state - .history - .update_environment_context_baseline(&environment_context); - } } pub(crate) async fn update_token_usage_info( diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 903ea0b47..3443c233d 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -7349,7 +7349,7 @@ async fn spawn_task_does_not_update_previous_turn_settings_for_non_run_turn_task } #[tokio::test] -async fn build_settings_update_items_emits_environment_item_for_network_changes() { +async fn record_context_updates_emits_environment_item_for_network_changes() { let (session, previous_context) = make_session_and_context().await; let previous_context = Arc::new(previous_context); let mut current_context = previous_context @@ -7396,10 +7396,8 @@ async fn build_settings_update_items_emits_environment_item_for_network_changes( .expect("rebuild config layer stack with network requirements"); current_context.config = Arc::new(config); - 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 update_items = + record_context_update_items(&session, &previous_context, ¤t_context).await; let environment_update = user_input_texts(&update_items) .into_iter() @@ -7411,50 +7409,40 @@ 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 record_context_updates_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 update_items = + record_context_update_items(&session, &previous_context, ¤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] -async fn build_settings_update_items_emits_environment_item_for_time_changes() { +async fn record_context_updates_emits_environment_item_for_time_changes() { let (session, previous_context) = make_session_and_context().await; let previous_context = Arc::new(previous_context); let mut current_context = previous_context @@ -7466,10 +7454,8 @@ async fn build_settings_update_items_emits_environment_item_for_time_changes() { current_context.current_date = Some("2026-02-27".to_string()); current_context.timezone = Some("Europe/Berlin".to_string()); - 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 update_items = + record_context_update_items(&session, &previous_context, ¤t_context).await; let environment_update = user_input_texts(&update_items) .into_iter() @@ -7480,7 +7466,7 @@ async fn build_settings_update_items_emits_environment_item_for_time_changes() { } #[tokio::test] -async fn build_settings_update_items_omits_environment_item_when_disabled() { +async fn record_context_updates_omits_environment_item_when_disabled() { let (session, previous_context) = make_session_and_context().await; let previous_context = Arc::new(previous_context); let mut current_context = previous_context @@ -7492,12 +7478,16 @@ 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 - .build_settings_update_items(Some(&reference_context_item), ¤t_context) - .await; + let update_items = + record_context_update_items(&session, &previous_context, ¤t_context).await; let user_texts = user_input_texts(&update_items); assert!( @@ -7508,6 +7498,23 @@ async fn build_settings_update_items_omits_environment_item_when_disabled() { ); } +async fn record_context_update_items( + session: &Session, + previous_context: &TurnContext, + current_context: &TurnContext, +) -> Vec { + session + .record_context_updates_and_set_reference_context_item(previous_context) + .await; + let previous_len = session.clone_history().await.raw_items().len(); + + session + .record_context_updates_and_set_reference_context_item(current_context) + .await; + let history = session.clone_history().await; + history.raw_items()[previous_len..].to_vec() +} + #[tokio::test] async fn build_settings_update_items_emits_realtime_start_when_session_becomes_live() { let (session, previous_context) = make_session_and_context().await; @@ -7769,9 +7776,11 @@ async fn record_context_updates_includes_turn_context_fragments_on_steady_state_ }); let mut previous_context_item = turn_context.to_turn_context_item(); previous_context_item.turn_id = Some("previous-turn-id".to_string()); + let world_state = session.build_world_state(&turn_context).await; { let mut state = session.state.lock().await; state.set_reference_context_item(Some(previous_context_item)); + state.history.set_world_state_baseline(world_state); } session @@ -8441,9 +8450,11 @@ async fn record_context_updates_and_set_reference_context_item_persists_baseline .with_model(next_model.to_string(), &session.services.models_manager) .await; let previous_context_item = previous_context.to_turn_context_item(); + let world_state = session.build_world_state(&previous_context).await; { let mut state = session.state.lock().await; state.set_reference_context_item(Some(previous_context_item.clone())); + state.history.set_world_state_baseline(world_state); } let rollout_path = attach_thread_persistence(&mut session).await; 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 000000000..5d1285774 --- /dev/null +++ b/codex-rs/core/src/session/world_state.rs @@ -0,0 +1,41 @@ +use super::turn_context::TurnContext; +use crate::context::world_state::EnvironmentsState; +use crate::context::world_state::WorldState; +use crate::environment_selection::TurnEnvironmentSnapshot; +use codex_protocol::protocol::TurnContextItem; + +pub(super) fn build_world_state_from_turn_context( + turn_context: &TurnContext, + environment_subagents: &str, +) -> WorldState { + let mut world_state = WorldState::default(); + if turn_context.config.include_environment_context { + world_state.add_section( + EnvironmentsState::from_turn_context(turn_context) + .with_subagents(environment_subagents.to_string()), + ); + } + world_state +} + +pub(super) fn build_world_state_from_environment_snapshot( + turn_context: &TurnContext, + environments: &TurnEnvironmentSnapshot, +) -> WorldState { + let mut world_state = WorldState::default(); + if turn_context.config.include_environment_context { + world_state.add_section(EnvironmentsState::from_turn_context_with_environments( + turn_context, + environments, + )); + } + world_state +} + +pub(super) fn build_world_state_from_turn_context_item( + turn_context_item: &TurnContextItem, +) -> WorldState { + let mut world_state = WorldState::default(); + world_state.add_section(EnvironmentsState::from_turn_context_item(turn_context_item)); + world_state +}