mirror of
https://github.com/openai/codex.git
synced 2026-10-21 17:48:44 +00:00
add model world state foundation
This commit is contained in:
pakrym-oai
@@ -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 = "<filesystem>".to_string();
|
||||
if !self.workspace_roots.is_empty() {
|
||||
rendered.push_str("<workspace_roots>");
|
||||
@@ -283,7 +283,7 @@ fn push_text_element(rendered: &mut String, name: &str, value: &str) {
|
||||
rendered.push_str(&format!("</{name}>"));
|
||||
}
|
||||
|
||||
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 = "<network enabled=\"true\">".to_string();
|
||||
Self::push_rendered_domain_element(&mut rendered, "allowed", &self.allowed_domains);
|
||||
Self::push_rendered_domain_element(&mut rendered, "denied", &self.denied_domains);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<String, EnvironmentState>,
|
||||
#[serde(skip)]
|
||||
current_date: Option<String>,
|
||||
#[serde(skip)]
|
||||
timezone: Option<String>,
|
||||
#[serde(skip)]
|
||||
network: Option<NetworkContext>,
|
||||
#[serde(skip)]
|
||||
filesystem: Option<FileSystemContext>,
|
||||
#[serde(skip)]
|
||||
subagents: Option<String>,
|
||||
}
|
||||
|
||||
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<ResponseItem> {
|
||||
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::<BTreeMap<_, _>>();
|
||||
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<String, EnvironmentUpdate>,
|
||||
current_date: Option<String>,
|
||||
timezone: Option<String>,
|
||||
network: Option<NetworkContext>,
|
||||
filesystem: Option<FileSystemContext>,
|
||||
subagents: Option<String>,
|
||||
}
|
||||
|
||||
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(" <environments>\n");
|
||||
for (id, update) in &self.updates {
|
||||
match update {
|
||||
EnvironmentUpdate::Current(environment) => {
|
||||
rendered.push_str(" <environment id=\"");
|
||||
push_xml_escaped_text(&mut rendered, id);
|
||||
rendered.push('"');
|
||||
if let Some(status) = environment.status {
|
||||
rendered.push_str(" status=\"");
|
||||
rendered.push_str(status.as_str());
|
||||
rendered.push('"');
|
||||
}
|
||||
rendered.push_str(">\n <cwd>");
|
||||
push_xml_escaped_text(
|
||||
&mut rendered,
|
||||
&environment.cwd.inferred_native_path_string(),
|
||||
);
|
||||
rendered.push_str("</cwd>");
|
||||
if let Some(shell) = &environment.shell {
|
||||
rendered.push_str("\n <shell>");
|
||||
push_xml_escaped_text(&mut rendered, shell);
|
||||
rendered.push_str("</shell>");
|
||||
}
|
||||
rendered.push_str("\n </environment>\n");
|
||||
}
|
||||
EnvironmentUpdate::Unavailable => {
|
||||
rendered.push_str(" <environment id=\"");
|
||||
push_xml_escaped_text(&mut rendered, id);
|
||||
rendered.push_str("\" status=\"unavailable\" />\n");
|
||||
}
|
||||
}
|
||||
}
|
||||
rendered.push_str(" </environments>\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(" <subagents>\n");
|
||||
for line in subagents.lines() {
|
||||
rendered.push_str(" ");
|
||||
rendered.push_str(line);
|
||||
rendered.push('\n');
|
||||
}
|
||||
rendered.push_str(" </subagents>\n");
|
||||
}
|
||||
rendered
|
||||
}
|
||||
}
|
||||
|
||||
fn insert_environment(
|
||||
environments: &mut BTreeMap<String, EnvironmentState>,
|
||||
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("</");
|
||||
rendered.push_str(name);
|
||||
rendered.push_str(">\n");
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
struct EnvironmentState {
|
||||
cwd: PathUri,
|
||||
#[serde(skip)]
|
||||
status: Option<EnvironmentStatus>,
|
||||
#[serde(skip)]
|
||||
shell: Option<String>,
|
||||
}
|
||||
|
||||
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;
|
||||
@@ -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#"<environment_context>
|
||||
<environments>
|
||||
<environment id="devbox" status="available">
|
||||
<cwd>/workspace</cwd>
|
||||
<shell>bash</shell>
|
||||
</environment>
|
||||
<environment id="laptop" status="available">
|
||||
<cwd>/repo</cwd>
|
||||
<shell>zsh</shell>
|
||||
</environment>
|
||||
</environments>
|
||||
<current_date>2026-06-20</current_date>
|
||||
<timezone>America/Los_Angeles</timezone>
|
||||
<network enabled="true"><allowed>api.example.com</allowed><denied>blocked.example.com</denied></network>
|
||||
<filesystem><permission_profile type="disabled"><file_system type="unrestricted" /></permission_profile></filesystem>
|
||||
<subagents>
|
||||
task_1: running
|
||||
task_2: completed
|
||||
</subagents>
|
||||
</environment_context>"#,
|
||||
)],
|
||||
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#"<environment_context>
|
||||
<environments>
|
||||
<environment id="laptop" status="available">
|
||||
<cwd>/new-repo</cwd>
|
||||
<shell>zsh</shell>
|
||||
</environment>
|
||||
<environment id="old" status="unavailable" />
|
||||
<environment id="remote" status="starting">
|
||||
<cwd>/remote</cwd>
|
||||
</environment>
|
||||
</environments>
|
||||
</environment_context>"#,
|
||||
)],
|
||||
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::<ResponseItem>::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::<EnvironmentsState>(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::<ResponseItem>::new(),
|
||||
live_world.render_diff(&loaded_world),
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn available(cwd: &str, shell: &str) -> Result<EnvironmentState> {
|
||||
Ok(EnvironmentState {
|
||||
cwd: PathUri::parse(cwd)?,
|
||||
status: Some(EnvironmentStatus::Available),
|
||||
shell: Some(shell.to_string()),
|
||||
})
|
||||
}
|
||||
|
||||
fn starting(cwd: &str) -> Result<EnvironmentState> {
|
||||
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,
|
||||
}
|
||||
}
|
||||
@@ -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<Value>;
|
||||
|
||||
fn render_diff(&self, previous: Option<&dyn Any>) -> Option<ResponseItem>;
|
||||
}
|
||||
|
||||
impl<S: WorldStateSection> ErasedWorldStateSection for S {
|
||||
fn as_any(&self) -> &dyn Any {
|
||||
self
|
||||
}
|
||||
|
||||
fn name(&self) -> &'static str {
|
||||
S::NAME
|
||||
}
|
||||
|
||||
fn json(&self) -> serde_json::Result<Value> {
|
||||
serde_json::to_value(self)
|
||||
}
|
||||
|
||||
fn render_diff(&self, previous: Option<&dyn Any>) -> Option<ResponseItem> {
|
||||
match previous {
|
||||
Some(previous) => {
|
||||
let Some(previous) = previous.downcast_ref::<S>() 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<ResponseItem>;
|
||||
}
|
||||
|
||||
type LoadSection = fn(Value) -> serde_json::Result<Box<dyn ErasedWorldStateSection>>;
|
||||
|
||||
struct SectionRegistration {
|
||||
name: &'static str,
|
||||
load: LoadSection,
|
||||
}
|
||||
|
||||
const SECTION_REGISTRY: &[SectionRegistration] = &[
|
||||
SectionRegistration {
|
||||
name: EnvironmentsState::NAME,
|
||||
load: load_section::<EnvironmentsState>,
|
||||
},
|
||||
#[cfg(test)]
|
||||
SectionRegistration {
|
||||
name: tests::AlphaSection::NAME,
|
||||
load: load_section::<tests::AlphaSection>,
|
||||
},
|
||||
#[cfg(test)]
|
||||
SectionRegistration {
|
||||
name: tests::BetaSection::NAME,
|
||||
load: load_section::<tests::BetaSection>,
|
||||
},
|
||||
];
|
||||
|
||||
fn load_section<S: WorldStateSection>(
|
||||
value: Value,
|
||||
) -> serde_json::Result<Box<dyn ErasedWorldStateSection>> {
|
||||
serde_json::from_value::<S>(value)
|
||||
.map(|section| Box::new(section) as Box<dyn ErasedWorldStateSection>)
|
||||
}
|
||||
|
||||
/// A snapshot of the model-visible world with one section per concrete type.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct WorldState {
|
||||
sections: IndexMap<TypeId, Box<dyn ErasedWorldStateSection>>,
|
||||
}
|
||||
|
||||
impl WorldState {
|
||||
pub(crate) fn add_section<S: WorldStateSection>(&mut self, section: S) {
|
||||
self.sections.insert(TypeId::of::<S>(), Box::new(section));
|
||||
}
|
||||
|
||||
pub(crate) fn json_full(&self) -> serde_json::Result<Value> {
|
||||
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<Value> {
|
||||
Ok(json_merge_patch_diff(
|
||||
&previous.json_full()?,
|
||||
&self.json_full()?,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) fn from_json(json: Value) -> serde_json::Result<Self> {
|
||||
let stored_sections = serde_json::from_value::<BTreeMap<String, 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<ResponseItem> {
|
||||
self.render_diff(&Self::default())
|
||||
}
|
||||
|
||||
pub(crate) fn render_diff(&self, previous: &Self) -> Vec<ResponseItem> {
|
||||
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;
|
||||
@@ -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<ResponseItem> {
|
||||
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(())
|
||||
}
|
||||
Reference in New Issue
Block a user