diff --git a/codex-rs/core/src/context/environment_context.rs b/codex-rs/core/src/context/environment_context.rs index e93ffb88ff..b313326561 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -117,7 +117,7 @@ enum ManagedFileSystemContext { } impl FileSystemContext { - fn from_permission_profile( + pub(super) fn from_permission_profile( permission_profile: &PermissionProfile, workspace_roots: &[AbsolutePathBuf], ) -> Self { @@ -143,7 +143,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(""); @@ -283,7 +283,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(super) fn push_xml_escaped_text(rendered: &mut String, value: &str) { for ch in value.chars() { match ch { '&' => rendered.push_str("&"), @@ -310,7 +310,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); diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 3347358bb5..494909b02a 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; 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 0000000000..23eca0d98e --- /dev/null +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -0,0 +1,271 @@ +use super::WorldStateSection; +use crate::context::ContextualUserFragment; +use crate::context::environment_context::EnvironmentContext; +use crate::context::environment_context::EnvironmentContextEnvironment; +use crate::context::environment_context::EnvironmentContextEnvironments; +use crate::context::environment_context::FileSystemContext; +use crate::context::environment_context::NetworkContext; +use crate::context::environment_context::push_xml_escaped_text; +use crate::session::turn_context::TurnContext; +use crate::shell::Shell; +use codex_protocol::models::ResponseItem; +use codex_utils_path_uri::PathUri; +use serde::Deserialize; +use serde::Serialize; +use std::collections::BTreeMap; + +/// Live environment values; only identifiers and working directories persist. +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct EnvironmentsState { + environments: BTreeMap, + #[serde(skip)] + current_date: Option, + #[serde(skip)] + timezone: Option, + #[serde(skip)] + network: Option, + #[serde(skip)] + filesystem: Option, + #[serde(skip)] + subagents: Option, +} + +impl PartialEq for EnvironmentsState { + fn eq(&self, other: &Self) -> bool { + self.environments == other.environments + } +} + +impl Eq for EnvironmentsState {} + +impl EnvironmentsState { + pub(crate) fn from_turn_context( + turn_context: &TurnContext, + default_shell: &Shell, + subagents: String, + ) -> Self { + let context = EnvironmentContext::from_turn_context(turn_context, default_shell) + .with_subagents(subagents); + let mut state = Self::from_environment_context(&context); + for environment in &turn_context.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_environment_context(context: &EnvironmentContext) -> Self { + let mut environments = BTreeMap::new(); + match &context.environments { + EnvironmentContextEnvironments::None => {} + EnvironmentContextEnvironments::Single(environment) => { + insert_environment(&mut environments, environment); + } + EnvironmentContextEnvironments::Multiple(current) => { + for environment in current { + insert_environment(&mut environments, environment); + } + } + } + Self { + environments, + current_date: context.current_date.clone(), + timezone: context.timezone.clone(), + network: context.network.clone(), + filesystem: context.filesystem.clone(), + subagents: context.subagents.clone(), + } + } +} + +impl WorldStateSection for EnvironmentsState { + const NAME: &'static str = "environments"; + + fn render_diff(&self, previous: &Self) -> Option { + 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)), + ); + (!updates.is_empty()).then(|| { + ContextualUserFragment::into(RenderedEnvironments { + updates, + current_date: self.current_date.clone(), + timezone: self.timezone.clone(), + network: self.network.clone(), + filesystem: self.filesystem.clone(), + subagents: self.subagents.clone(), + }) + }) + } +} + +struct RenderedEnvironments { + updates: BTreeMap, + 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) { + ( + codex_protocol::protocol::ENVIRONMENT_CONTEXT_OPEN_TAG, + codex_protocol::protocol::ENVIRONMENT_CONTEXT_CLOSE_TAG, + ) + } + + fn body(&self) -> String { + let mut rendered = "\n".to_string(); + 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_xml_escaped_text( + &mut rendered, + &environment.cwd.inferred_native_path_string(), + ); + rendered.push_str(""); + if let Some(shell) = &environment.shell { + rendered.push_str("\n "); + push_xml_escaped_text(&mut rendered, shell); + rendered.push_str(""); + } + rendered.push_str("\n \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 insert_environment( + environments: &mut BTreeMap, + environment: &EnvironmentContextEnvironment, +) { + environments.insert( + environment.id.clone(), + EnvironmentState { + cwd: environment.cwd.clone(), + status: Some(EnvironmentStatus::Available), + shell: Some(environment.shell.clone()), + }, + ); +} + +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, Serialize, Deserialize)] +struct EnvironmentState { + cwd: PathUri, + #[serde(skip)] + status: Option, + #[serde(skip)] + 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, +} + +impl EnvironmentStatus { + fn as_str(self) -> &'static str { + match self { + Self::Starting => "starting", + Self::Available => "available", + } + } +} + +#[cfg(test)] +#[path = "environment_tests.rs"] +mod tests; 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 0000000000..2d1233ae18 --- /dev/null +++ b/codex-rs/core/src/context/world_state/environment_tests.rs @@ -0,0 +1,228 @@ +use super::*; +use crate::context::world_state::WorldState; +use anyhow::Result; +use codex_protocol::models::ContentItem; +use codex_protocol::models::PermissionProfile; +use pretty_assertions::assert_eq; +use serde_json::json; + +#[test] +fn renders_full_environment_state() -> Result<()> { + let mut context = EnvironmentContext::new( + vec![ + EnvironmentContextEnvironment { + id: "laptop".to_string(), + cwd: PathUri::parse("file:///repo")?, + shell: "zsh".to_string(), + }, + EnvironmentContextEnvironment { + id: "devbox".to_string(), + cwd: PathUri::parse("file:///workspace")?, + shell: "bash".to_string(), + }, + ], + Some("2026-06-20".to_string()), + Some("America/Los_Angeles".to_string()), + Some(NetworkContext::new( + vec!["api.example.com".to_string()], + vec!["blocked.example.com".to_string()], + )), + Some("task_1: running\ntask_2: completed".to_string()), + ); + context.filesystem = Some(FileSystemContext::from_permission_profile( + &PermissionProfile::Disabled, + &[], + )); + + let mut world_state = WorldState::default(); + world_state.add_section(EnvironmentsState::from_environment_context(&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 + +"#, + )], + 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 + + +"#, + )], + current.render_diff(&previous), + ); + Ok(()) +} + +#[test] +fn unchanged_environments_do_not_render_a_diff() -> Result<()> { + let environments = EnvironmentsState { + environments: [("laptop".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, + &[], + )), + subagents: Some("task_1: running".to_string()), + ..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: None, + subagents: Some("task_1: completed".to_string()), + ..environments + }); + + assert_eq!(Vec::::new(), current.render_diff(&previous)); + Ok(()) +} + +#[test] +fn loaded_environment_state_produces_no_diff_with_live_state() -> Result<()> { + let live_state = EnvironmentsState { + environments: [ + ("laptop".to_string(), available("file:///repo", "zsh")?), + ("devbox".to_string(), starting("file:///workspace")?), + ] + .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".to_string()), + }; + + let stored = serde_json::to_value(&live_state)?; + assert_eq!( + json!({ + "devbox": { + "cwd": "file:///workspace", + }, + "laptop": { + "cwd": "file:///repo", + }, + }), + stored, + ); + let loaded_state = serde_json::from_value::(stored)?; + let mut live_world = WorldState::default(); + live_world.add_section(live_state); + let mut loaded_world = WorldState::default(); + loaded_world.add_section(loaded_state); + assert_eq!( + Vec::::new(), + live_world.render_diff(&loaded_world), + ); + 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, + } +} 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 0000000000..7cc239e8df --- /dev/null +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -0,0 +1,215 @@ +mod environment; + +use codex_protocol::models::ResponseItem; +use environment::EnvironmentsState; +use indexmap::IndexMap; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value; +use std::any::Any; +use std::any::TypeId; +use std::collections::BTreeMap; + +trait ErasedWorldStateSection: Send + Sync { + fn as_any(&self) -> &dyn Any; + + fn name(&self) -> &'static str; + + fn json(&self) -> serde_json::Result; + + fn render_diff(&self, previous: Option<&dyn Any>) -> Option; +} + +impl ErasedWorldStateSection for S { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &'static str { + S::NAME + } + + fn json(&self) -> serde_json::Result { + serde_json::to_value(self) + } + + fn render_diff(&self, previous: Option<&dyn Any>) -> Option { + match previous { + Some(previous) => { + let Some(previous) = previous.downcast_ref::() else { + unreachable!("world-state section type must match its type ID"); + }; + WorldStateSection::render_diff(self, previous) + } + None => WorldStateSection::render_diff(self, &S::default()), + } + } +} + +/// 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. `NAME` is the stable key used in +/// persisted JSON and must be registered in `SECTION_REGISTRY`. +pub(crate) trait WorldStateSection: + Any + Default + DeserializeOwned + Serialize + Send + Sync +{ + const NAME: &'static str; + + fn render_diff(&self, previous: &Self) -> Option; +} + +type LoadSection = fn(Value) -> serde_json::Result>; + +struct SectionRegistration { + name: &'static str, + load: LoadSection, +} + +const SECTION_REGISTRY: &[SectionRegistration] = &[ + SectionRegistration { + name: EnvironmentsState::NAME, + load: load_section::, + }, + #[cfg(test)] + SectionRegistration { + name: tests::AlphaSection::NAME, + load: load_section::, + }, + #[cfg(test)] + SectionRegistration { + name: tests::BetaSection::NAME, + load: load_section::, + }, +]; + +fn load_section( + value: Value, +) -> serde_json::Result> { + serde_json::from_value::(value) + .map(|section| Box::new(section) as Box) +} + +/// A snapshot of the model-visible world with one section per concrete type. +#[derive(Default)] +pub(crate) struct WorldState { + sections: IndexMap>, +} + +impl WorldState { + pub(crate) fn add_section(&mut self, section: S) { + self.sections.insert(TypeId::of::(), Box::new(section)); + } + + pub(crate) fn json_full(&self) -> serde_json::Result { + let mut sections = BTreeMap::new(); + for section in self.sections.values() { + sections.insert(section.name(), section.json()?); + } + serde_json::to_value(sections) + } + + /// Returns the JSON Merge Patch that transforms `previous` into this state. + pub(crate) fn json_diff(&self, previous: &Self) -> serde_json::Result { + Ok(json_merge_patch_diff( + &previous.json_full()?, + &self.json_full()?, + )) + } + + pub(crate) fn from_json(json: Value) -> serde_json::Result { + let stored_sections = serde_json::from_value::>(json)?; + let mut world_state = Self::default(); + for (name, value) in stored_sections { + let Some(registration) = SECTION_REGISTRY + .iter() + .find(|registration| registration.name == name) + else { + continue; + }; + let section = (registration.load)(value)?; + world_state + .sections + .insert(section.as_any().type_id(), section); + } + Ok(world_state) + } + + /// Applies a persisted diff, discarding values that are not represented in JSON. + pub(crate) fn apply_json_diff(&mut self, diff: &Value) -> serde_json::Result<()> { + let mut json = self.json_full()?; + apply_json_merge_patch(&mut json, diff); + *self = Self::from_json(json)?; + Ok(()) + } + + 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() + } +} + +fn json_merge_patch_diff(previous: &Value, current: &Value) -> Value { + if previous == current { + return Value::Object(serde_json::Map::new()); + } + + let (Value::Object(previous), Value::Object(current)) = (previous, current) else { + return current.clone(); + }; + + let mut diff = serde_json::Map::new(); + for key in previous.keys() { + if !current.contains_key(key) { + diff.insert(key.clone(), Value::Null); + } + } + for (key, current) in current { + match previous.get(key) { + Some(previous) if previous == current => {} + Some(previous) => { + diff.insert(key.clone(), json_merge_patch_diff(previous, current)); + } + None => { + diff.insert(key.clone(), current.clone()); + } + } + } + Value::Object(diff) +} + +fn apply_json_merge_patch(target: &mut Value, patch: &Value) { + let Value::Object(patch) = patch else { + *target = patch.clone(); + return; + }; + if !target.is_object() { + *target = Value::Object(serde_json::Map::new()); + } + let Value::Object(target) = target else { + unreachable!("target was initialized as an object"); + }; + for (key, value) in patch { + if value.is_null() { + target.remove(key); + } else { + apply_json_merge_patch(target.entry(key.clone()).or_insert(Value::Null), value); + } + } +} + +#[cfg(test)] +#[path = "world_state_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/context/world_state/world_state_tests.rs b/codex-rs/core/src/context/world_state/world_state_tests.rs new file mode 100644 index 0000000000..d9d510f371 --- /dev/null +++ b/codex-rs/core/src/context/world_state/world_state_tests.rs @@ -0,0 +1,76 @@ +use super::*; +use anyhow::Result; +use pretty_assertions::assert_eq; +use serde::Deserialize; +use serde::Serialize; +use serde_json::json; + +macro_rules! mock_section { + ($type:ident, $name:literal) => { + #[derive(Debug, Default, Deserialize, Eq, PartialEq, Serialize)] + pub(super) struct $type { + value: String, + } + + impl $type { + fn new(value: &str) -> Self { + Self { + value: value.to_string(), + } + } + } + + impl WorldStateSection for $type { + const NAME: &'static str = $name; + + fn render_diff(&self, _previous: &Self) -> Option { + None + } + } + }; +} + +mock_section!(AlphaSection, "alpha"); +mock_section!(BetaSection, "beta"); + +#[test] +fn full_json_reloads_registered_sections() -> Result<()> { + let mut state = WorldState::default(); + state.add_section(BetaSection::new("two")); + state.add_section(AlphaSection::new("one")); + + let full = state.json_full()?; + assert_eq!( + json!({ + "alpha": { "value": "one" }, + "beta": { "value": "two" }, + }), + full, + ); + + let loaded = WorldState::from_json(full.clone())?; + assert_eq!(full, loaded.json_full()?); + Ok(()) +} + +#[test] +fn json_diff_reconstructs_current_state() -> Result<()> { + let mut previous = WorldState::default(); + previous.add_section(AlphaSection::new("old")); + previous.add_section(BetaSection::new("removed")); + let mut current = WorldState::default(); + current.add_section(AlphaSection::new("new")); + + let diff = current.json_diff(&previous)?; + assert_eq!( + json!({ + "alpha": { "value": "new" }, + "beta": null, + }), + diff, + ); + + previous.apply_json_diff(&diff)?; + assert_eq!(current.json_full()?, previous.json_full()?); + Ok(()) +}