diff --git a/Cargo.lock b/Cargo.lock index 31956b5..c016f71 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8,6 +8,239 @@ version = "1.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" + +[[package]] +name = "bitflags" +version = "2.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3" + +[[package]] +name = "blake3" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3888aaa89e4b2a40fca9848e400f6a658a5a3978de7be858e209cafa8be9a4a0" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", +] + +[[package]] +name = "cc" +version = "1.2.50" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f50d563227a1c37cc0a263f64eca3334388c01c5e4c4861a9def205c614383c" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "constant_time_eq" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "fastrand" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" + +[[package]] +name = "find-msvc-tools" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a3076410a55c90011c298b04d0cfa770b00fa04e1e3c97d3f6c9de105a03844" + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "itoa" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" + +[[package]] +name = "libc" +version = "0.2.178" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37c93d8daa9d8a012fd8ab92f088405fb202ea0b6ab73ee2482ae66af4f42091" + +[[package]] +name = "linux-raw-sys" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039" + +[[package]] +name = "memchr" +version = "2.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" + +[[package]] +name = "once_cell" +version = "1.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" + +[[package]] +name = "proc-macro2" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ee95bc4ef87b8d5ba32e8b7714ccc834865276eab0aed5c9958d00ec45f49e8" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a338cc41d27e6cc6dce6cefc13a0729dfbb81c262b1f519331575dd80ef3067f" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "rustix" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd15f8a2c5551a84d56efdc1cd049089e409ac19a3072d5037a17fd70719ff3e" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "ryu" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62049b2877bf12821e8f9ad256ee38fdc31db7387ec2d3b3f403024de2034aea" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.145" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "402a6f66d8c709116cf22f558eab210f5a50187f702eb4d7e5ef38d9a7f1c79c" +dependencies = [ + "itoa", + "memchr", + "ryu", + "serde", + "serde_core", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "syn" +version = "2.0.111" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "390cc9a294ab71bdb1aa2e99d13be9c753cd2d7bd6560c77118597410c4d2e87" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tempfile" +version = "3.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d31c77bdf42a745371d260a26ca7163f1e0924b64afa0b688e61b5a9fa02f16" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys", +] + [[package]] name = "treeagent" version = "0.1.0" @@ -22,6 +255,10 @@ name = "treeagent-model" version = "0.1.0" dependencies = [ "anyhow", + "blake3", + "serde", + "serde_json", + "tempfile", ] [[package]] @@ -39,3 +276,39 @@ dependencies = [ "anyhow", "treeagent-model", ] + +[[package]] +name = "unicode-ident" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" + +[[package]] +name = "wasip2" +version = "1.0.1+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0562428422c63773dad2c345a1882263bbf4d65cf3f42e90921f787ef5ad58e7" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "wit-bindgen" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" diff --git a/crates/treeagent-model/Cargo.toml b/crates/treeagent-model/Cargo.toml index 1b49c41..3b03f01 100644 --- a/crates/treeagent-model/Cargo.toml +++ b/crates/treeagent-model/Cargo.toml @@ -4,6 +4,16 @@ authors.workspace = true edition.workspace = true license.workspace = true version.workspace = true +build = "build.rs" [dependencies] anyhow.workspace = true +blake3 = "1.5" +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" + +[dev-dependencies] +tempfile = "3.10" + +[build-dependencies] +anyhow.workspace = true diff --git a/crates/treeagent-model/build.rs b/crates/treeagent-model/build.rs new file mode 100644 index 0000000..44f4335 --- /dev/null +++ b/crates/treeagent-model/build.rs @@ -0,0 +1,20 @@ +use std::collections::BTreeSet; +use std::env; + +fn main() -> anyhow::Result<()> { + let target = env::var("TARGET").unwrap_or_else(|_| "unknown-target".to_string()); + println!("cargo:rustc-env=TREEAGENT_BUILD_TARGET={target}"); + + let profile = env::var("PROFILE").unwrap_or_else(|_| "unknown-profile".to_string()); + println!("cargo:rustc-env=TREEAGENT_BUILD_PROFILE={profile}"); + + let mut features = BTreeSet::new(); + for (key, _) in env::vars() { + if let Some(feature) = key.strip_prefix("CARGO_FEATURE_") { + features.insert(feature.to_ascii_lowercase().replace('_', "-")); + } + } + let feature_list = features.into_iter().collect::>().join(","); + println!("cargo:rustc-env=TREEAGENT_BUILD_FEATURES={feature_list}"); + Ok(()) +} diff --git a/crates/treeagent-model/src/lib.rs b/crates/treeagent-model/src/lib.rs index 586168c..d9ac8ea 100644 --- a/crates/treeagent-model/src/lib.rs +++ b/crates/treeagent-model/src/lib.rs @@ -1,3 +1,10 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::fs; +use std::path::{Path, PathBuf}; + +use anyhow::{bail, Context, Result}; +use serde::{Deserialize, Serialize}; + #[derive(Debug, Default, Clone, PartialEq, Eq)] pub struct GraphSpec { name: String, @@ -13,13 +20,394 @@ impl GraphSpec { } } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +pub struct GraphIR { + #[serde(default)] + pub spawn_rules: BTreeMap, + #[serde(default)] + pub execution: ExecutionLimits, + #[serde(default)] + pub checkpoint: CheckpointMetadata, +} + +impl GraphIR { + pub fn canonicalized(mut self) -> Self { + self.spawn_rules = self + .spawn_rules + .into_iter() + .map(|(task, rule)| (task, rule.canonicalized())) + .collect(); + self.execution = self.execution.canonicalized(); + self.checkpoint = self.checkpoint.canonicalized(); + self + } + + pub fn validate(&self) -> Result<()> { + for (parent, rule) in &self.spawn_rules { + if parent.trim().is_empty() { + bail!("spawn rule names cannot be empty"); + } + rule.validate(parent)?; + } + self.execution.validate()?; + self.checkpoint.validate()?; + Ok(()) + } + + pub fn canonical_json(&self) -> Result> { + let canonical = self.clone().canonicalized(); + serde_json::to_vec(&canonical).context("failed to serialize canonical graph IR") + } + + pub fn blake3_hash(&self, manifest: &BuildManifest) -> Result { + let mut hasher = blake3::Hasher::new(); + hasher.update(&self.canonical_json()?); + hasher.update(&manifest.canonical_bytes()?); + Ok(hasher.finalize()) + } + + pub fn blake3_hex(&self, manifest: &BuildManifest) -> Result { + Ok(self.blake3_hash(manifest)?.to_hex().to_string()) + } + + pub fn load_spawn_rules_from_disk() -> Result { + let cwd = std::env::current_dir().context("failed to resolve current directory")?; + let path = Self::resolve_spawn_rules_path(&cwd); + let spawn_rules = match path { + Some(path) => { + let data = fs::read_to_string(&path) + .with_context(|| format!("failed to read spawn rules at {}", path.display()))?; + let parsed: BTreeMap = serde_json::from_str(&data) + .with_context(|| { + format!("failed to parse spawn rules from {}", path.display()) + })?; + parsed + } + None => BTreeMap::new(), + }; + + let ir = GraphIR { + spawn_rules, + ..GraphIR::default() + } + .canonicalized(); + + ir.validate().context("invalid spawn rule configuration")?; + Ok(ir) + } + + fn resolve_spawn_rules_path(base: &Path) -> Option { + let primary = base.join("spawn_rules.json"); + if primary.exists() { + return Some(primary); + } + let fallback = base.join("config").join("spawn_rules.json"); + if fallback.exists() { + return Some(fallback); + } + None + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SpawnRule { + #[serde(default)] + pub can_spawn: BTreeMap, + #[serde(default = "default_self_spawn")] + pub self_spawn: bool, + #[serde(default)] + pub max_children: Option, +} + +impl Default for SpawnRule { + fn default() -> Self { + Self { + can_spawn: BTreeMap::new(), + self_spawn: default_self_spawn(), + max_children: None, + } + } +} + +impl SpawnRule { + fn canonicalized(mut self) -> Self { + self.can_spawn = self.can_spawn.into_iter().collect(); + if let Some(0) = self.max_children { + self.max_children = None; + } + self + } + + fn validate(&self, parent: &str) -> Result<()> { + for (child, limit) in &self.can_spawn { + if child.trim().is_empty() { + bail!("spawn rule '{}' has an empty child task name", parent); + } + if *limit == 0 { + bail!( + "spawn rule '{}' cannot allow zero spawns for '{}'", + parent, + child + ); + } + } + if let Some(limit) = self.max_children { + if limit == 0 { + bail!("spawn rule '{}' has a zero max_children limit", parent); + } + let total: u32 = self.can_spawn.values().copied().sum(); + if total > 0 && limit < total { + bail!( + "spawn rule '{}' max_children ({}) is less than the sum of child limits ({})", + parent, + limit, + total + ); + } + } + Ok(()) + } +} + +fn default_self_spawn() -> bool { + true +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +pub struct ExecutionLimits { + #[serde(default)] + pub max_depth: Option, + #[serde(default)] + pub max_nodes: Option, + #[serde(default)] + pub max_children_per_node: Option, +} + +impl ExecutionLimits { + fn canonicalized(mut self) -> Self { + if let Some(0) = self.max_depth { + self.max_depth = None; + } + if let Some(0) = self.max_nodes { + self.max_nodes = None; + } + if let Some(0) = self.max_children_per_node { + self.max_children_per_node = None; + } + self + } + + fn validate(&self) -> Result<()> { + if let Some(depth) = self.max_depth { + if depth == 0 { + bail!("execution limit max_depth cannot be zero"); + } + } + if let Some(nodes) = self.max_nodes { + if nodes == 0 { + bail!("execution limit max_nodes cannot be zero"); + } + } + if let Some(children) = self.max_children_per_node { + if children == 0 { + bail!("execution limit max_children_per_node cannot be zero"); + } + } + Ok(()) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +pub struct CheckpointMetadata { + #[serde(default)] + pub enabled: bool, + #[serde(default)] + pub interval: Option, + #[serde(default)] + pub tags: BTreeMap, +} + +impl CheckpointMetadata { + fn canonicalized(mut self) -> Self { + self.tags = self.tags.into_iter().collect(); + if let Some(0) = self.interval { + self.interval = None; + } + self + } + + fn validate(&self) -> Result<()> { + if let Some(interval) = self.interval { + if interval == 0 { + bail!("checkpoint interval cannot be zero"); + } + } + Ok(()) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct BuildManifest { + pub runtime_version: String, + pub target_triple: String, + pub profile: String, + pub features: BTreeSet, +} + +impl BuildManifest { + pub fn from_env() -> Self { + let runtime_version = option_env!("TREEAGENT_RUNTIME_VERSION") + .unwrap_or(env!("CARGO_PKG_VERSION")) + .to_string(); + let target_triple = option_env!("TREEAGENT_BUILD_TARGET") + .unwrap_or("unknown-target") + .to_string(); + let profile = option_env!("TREEAGENT_BUILD_PROFILE") + .unwrap_or("unknown-profile") + .to_string(); + let features = option_env!("TREEAGENT_BUILD_FEATURES") + .unwrap_or_default() + .split(',') + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .collect(); + Self { + runtime_version, + target_triple, + profile, + features, + } + } + + fn canonicalized(mut self) -> Self { + self.features = self.features.into_iter().collect(); + self + } + + fn canonical_bytes(&self) -> Result> { + let canonical = self.clone().canonicalized(); + serde_json::to_vec(&canonical).context("failed to serialize build manifest") + } +} + #[cfg(test)] mod tests { use super::*; + use tempfile::TempDir; #[test] fn stores_name() { let spec = GraphSpec::new("example"); assert_eq!(spec.name(), "example"); } + + #[test] + fn canonicalization_sets_defaults() { + let ir = GraphIR { + spawn_rules: BTreeMap::from([( + "HLD".into(), + SpawnRule { + can_spawn: BTreeMap::from([("LLD".into(), 2)]), + self_spawn: false, + max_children: Some(0), + }, + )]), + execution: ExecutionLimits { + max_depth: Some(0), + ..ExecutionLimits::default() + }, + checkpoint: CheckpointMetadata { + interval: Some(0), + tags: BTreeMap::from([("a".into(), "1".into()), ("b".into(), "2".into())]), + ..CheckpointMetadata::default() + }, + } + .canonicalized(); + + assert!(ir.spawn_rules["HLD"].max_children.is_none()); + assert!(ir.execution.max_depth.is_none()); + assert!(ir.checkpoint.interval.is_none()); + let tag_keys: Vec<_> = ir.checkpoint.tags.keys().cloned().collect(); + assert_eq!(tag_keys, vec!["a".to_string(), "b".to_string()]); + } + + #[test] + fn loader_prefers_cwd_rules() { + let tmp = TempDir::new().unwrap(); + let cwd = tmp.path(); + + fs::create_dir_all(cwd.join("config")).unwrap(); + fs::write( + cwd.join("config").join("spawn_rules.json"), + r#"{ "HLD": { "can_spawn": { "LLD": 1 } } }"#, + ) + .unwrap(); + fs::write( + cwd.join("spawn_rules.json"), + r#"{ "HLD": { "can_spawn": { "TEST": 2 }, "self_spawn": false } }"#, + ) + .unwrap(); + + let original_dir = std::env::current_dir().unwrap(); + std::env::set_current_dir(cwd).unwrap(); + let loaded = GraphIR::load_spawn_rules_from_disk().unwrap(); + std::env::set_current_dir(original_dir).unwrap(); + + let hld = loaded.spawn_rules.get("HLD").unwrap(); + assert_eq!( + hld.can_spawn.keys().cloned().collect::>(), + vec!["TEST"] + ); + assert!(!hld.self_spawn); + } + + #[test] + fn loader_surfaces_validation_errors() { + let tmp = TempDir::new().unwrap(); + let cwd = tmp.path(); + + fs::write( + cwd.join("spawn_rules.json"), + r#"{ "HLD": { "can_spawn": { "LLD": 0 } } }"#, + ) + .unwrap(); + + let original_dir = std::env::current_dir().unwrap(); + std::env::set_current_dir(cwd).unwrap(); + let result = GraphIR::load_spawn_rules_from_disk(); + std::env::set_current_dir(original_dir).unwrap(); + + assert!(result.is_err()); + } + + #[test] + fn manifest_changes_hash() { + let ir = GraphIR { + spawn_rules: BTreeMap::from([( + "HLD".into(), + SpawnRule { + can_spawn: BTreeMap::from([("LLD".into(), 1)]), + ..SpawnRule::default() + }, + )]), + ..GraphIR::default() + }; + + let manifest_a = BuildManifest { + runtime_version: "0.1.0".into(), + target_triple: "x86_64-unknown-linux-gnu".into(), + profile: "debug".into(), + features: BTreeSet::from(["foo".into()]), + }; + let manifest_b = BuildManifest { + runtime_version: "0.1.0".into(), + target_triple: "x86_64-unknown-linux-gnu".into(), + profile: "release".into(), + features: BTreeSet::new(), + }; + + let hash_a = ir.blake3_hex(&manifest_a).unwrap(); + let hash_b = ir.blake3_hex(&manifest_b).unwrap(); + assert_ne!(hash_a, hash_b); + } } diff --git a/crates/treeagent-runtime/src/lib.rs b/crates/treeagent-runtime/src/lib.rs index 21a1db5..5b40714 100644 --- a/crates/treeagent-runtime/src/lib.rs +++ b/crates/treeagent-runtime/src/lib.rs @@ -13,7 +13,10 @@ impl WorkspaceRuntime { } pub fn summary(&self) -> String { - format!("TreeAgent runtime initialized for graph '{}'", self.graph_spec.name()) + format!( + "TreeAgent runtime initialized for graph '{}'", + self.graph_spec.name() + ) } pub fn print_summary(&self) -> anyhow::Result<()> {