use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use oneagentgraph::config::ConfigRef;
use onevcs::SessionRequest;
use crate::agentgraph::{Environment, GraphOutput, GraphRun, Launch};
use crate::controls::{NodeControls, WORKER_MEMBER};
use crate::error::{Error, Result};
use crate::event::{Envelope, Labels};
pub trait Executor {
fn name(&self) -> &str;
fn capabilities(&self) -> Capabilities;
fn capacity(&self) -> CapacityReport;
fn dispatch(&self, req: DispatchRequest) -> Result<Box<dyn DispatchHandle>>;
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
pub struct Capabilities {
pub vcs_sessions: bool,
}
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct CapacityReport {
pub slots_free: u32,
pub load1: f64,
pub mem_free_bytes: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct DispatchRequest {
pub graph: ConfigRef,
pub task: String,
pub labels: Labels,
pub controls: NodeControls,
pub workspace: WorkspaceSpec,
pub cancel: CancellationToken,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WorkspaceSpec {
Path(PathBuf),
VcsSession(SessionRequest),
}
#[derive(Debug, Clone, Default)]
pub struct CancellationToken(Arc<AtomicBool>);
impl CancellationToken {
pub fn new() -> Self {
Self::default()
}
pub fn cancel(&self) {
self.0.store(true, Ordering::SeqCst);
}
pub fn is_cancelled(&self) -> bool {
self.0.load(Ordering::SeqCst)
}
}
impl PartialEq for CancellationToken {
fn eq(&self, other: &Self) -> bool {
self.is_cancelled() == other.is_cancelled()
}
}
pub trait DispatchHandle {
fn events(&mut self) -> EventStream;
fn wait(&mut self) -> Result<DispatchOutcome>;
fn cancel(&self, mode: CancelMode);
}
pub type EventStream = Box<dyn Iterator<Item = Result<Envelope>> + Send>;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum CancelMode {
Cooperative,
Kill,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
#[non_exhaustive]
pub struct DispatchOutcome {
pub succeeded: bool,
pub detail: String,
pub session: Option<String>,
pub branch: Option<String>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
pub struct LocalExecutor;
impl Executor for LocalExecutor {
fn name(&self) -> &str {
"local"
}
fn capabilities(&self) -> Capabilities {
Capabilities { vcs_sessions: true }
}
fn capacity(&self) -> CapacityReport {
let load1 = load_average().unwrap_or(0.0);
let cores = std::thread::available_parallelism()
.map(std::num::NonZeroUsize::get)
.unwrap_or(1);
let busy = load1.ceil().max(0.0);
let busy = if busy.is_finite() { busy as u64 } else { 0 };
CapacityReport {
slots_free: u32::try_from(u64::try_from(cores).unwrap_or(1).saturating_sub(busy))
.unwrap_or(u32::MAX),
load1,
mem_free_bytes: available_memory().unwrap_or(u64::MAX),
}
}
fn dispatch(&self, req: DispatchRequest) -> Result<Box<dyn DispatchHandle>> {
let (dir, session) = match &req.workspace {
WorkspaceSpec::Path(path) => (path.clone(), None),
WorkspaceSpec::VcsSession(request) => {
let session = crate::vcs::session_open(request)?;
(session.worktree.clone(), Some(session))
}
};
let node_sets = node_sets(&req.labels, &req.controls)?;
let filters = launched_with(&req.labels)?
.map(|record| record.filters)
.unwrap_or_default();
let env = prepare_dispatch_env(&req.labels)?;
let mut run = GraphRun::start(&Launch {
graph: &req.graph.0,
task: &req.task,
dir: &dir,
labels: &req.labels,
env: &env,
environment: Environment::PerLaunch,
sets: &node_sets,
filter: filters.agentgraph.as_ref(),
output: GraphOutput::Relayed,
})?;
let claim = match register_dispatch(&req.labels, run.process()) {
Ok(claim) => claim,
Err(refusal) => {
run.cancel();
let _ = run.wait();
return Err(refusal);
}
};
Ok(Box::new(LocalDispatch {
run,
cancel: req.cancel,
labels: req.labels,
session,
_claim: claim,
}))
}
}
pub(crate) const NODE_SCRATCH_DIR_ENV: &str = "ONEPIPELINE_NODE_SCRATCH_DIR";
fn prepare_dispatch_env(labels: &Labels) -> Result<Vec<(String, String)>> {
let mut env: Vec<(String, String)> = labels
.run_id
.iter()
.map(|run| (crate::agentgraph::RUN_ID_ENV.to_string(), run.clone()))
.collect();
let scratch = make_node_scratch_dir(labels)?.display().to_string();
env.push((crate::channel::ASKER_ENV.to_string(), scratch.clone()));
env.push((NODE_SCRATCH_DIR_ENV.to_string(), scratch));
Ok(env)
}
fn make_node_scratch_dir(labels: &Labels) -> Result<PathBuf> {
const TRIES: u64 = 4096;
static MINTED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let base = match labels.run_id.as_deref() {
Some(run) => crate::ledger::RunPaths::under(&crate::ledger::runs_root(), run)
.dir
.join("scratch"),
None => std::env::temp_dir().join("onepipeline-scratch"),
};
let ledger = |path: &Path| {
let path = path.to_path_buf();
move |source: std::io::Error| Error::Ledger { path, source }
};
std::fs::create_dir_all(&base).map_err(ledger(&base))?;
let pid = crate::sys::pid();
for _ in 0..TRIES {
let at = base.join(format!(
"{pid}-{}",
MINTED.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
));
match std::fs::create_dir(&at) {
Ok(()) => return std::fs::canonicalize(&at).map_err(ledger(&at)),
Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {}
Err(error) => return Err(ledger(&at)(error)),
}
}
Err(Error::Ledger {
path: base.clone(),
source: std::io::Error::new(
std::io::ErrorKind::AlreadyExists,
format!(
"no scratch directory under {} could be created",
base.display()
),
),
})
}
fn register_dispatch(
labels: &Labels,
process: Option<u32>,
) -> Result<Option<crate::ledger::DispatchClaim>> {
let (Some(run), Some(node)) = (labels.run_id.as_deref(), labels.node.as_deref()) else {
return Ok(None);
};
let paths = crate::ledger::RunPaths::under(&crate::ledger::runs_root(), run);
crate::ledger::claim_dispatch(&paths, node, process.unwrap_or_else(crate::sys::pid)).map(Some)
}
fn node_sets(labels: &Labels, controls: &NodeControls) -> Result<Vec<String>> {
if labels.persona.as_deref() == Some(crate::lifecycle::PR_AUTHOR_PERSONA) {
return Ok(Vec::new());
}
let mut sets = launched_with(labels)?.map_or_else(Vec::new, |record| record.node_sets);
if let Some(persona) = &labels.persona {
sets.push(format!("members.{WORKER_MEMBER}.persona={persona}"));
}
sets.extend(controls.overrides().map_err(Error::Invalid)?);
Ok(sets)
}
fn launched_with(labels: &Labels) -> Result<Option<crate::ledger::LaunchRecord>> {
let Some(run) = labels.run_id.as_deref() else {
return Ok(None);
};
let paths = crate::ledger::RunPaths::under(&crate::ledger::runs_root(), run);
crate::ledger::read_json::<crate::ledger::LaunchRecord>(&paths.launch()).map(Some)
}
#[derive(Debug)]
struct LocalDispatch {
run: GraphRun,
cancel: CancellationToken,
labels: Labels,
session: Option<onevcs::Session>,
_claim: Option<crate::ledger::DispatchClaim>,
}
impl DispatchHandle for LocalDispatch {
fn events(&mut self) -> EventStream {
let opened = self.session.as_ref().map(|session| {
Ok(crate::vcs::session_opened_event(session, &self.labels))
});
match opened {
Some(event) => Box::new(std::iter::once(event).chain(self.run.events())),
None => self.run.events(),
}
}
fn wait(&mut self) -> Result<DispatchOutcome> {
let settled = self.run.wait()?;
Ok(DispatchOutcome {
succeeded: settled.succeeded(),
detail: settled.stderr.trim().to_string(),
session: self.session.as_ref().map(|s| s.token.0.clone()),
branch: self.session.as_ref().map(|s| s.branch.clone()),
})
}
fn cancel(&self, mode: CancelMode) {
self.cancel.cancel();
if mode == CancelMode::Kill {
self.run.cancel();
}
}
}
fn load_average() -> Option<f64> {
let text = std::fs::read_to_string("/proc/loadavg").ok()?;
text.split_whitespace()
.next()?
.parse::<f64>()
.ok()
.filter(|value| value.is_finite() && *value >= 0.0)
}
fn available_memory() -> Option<u64> {
let text = std::fs::read_to_string("/proc/meminfo").ok()?;
for line in text.lines() {
if let Some(rest) = line.strip_prefix("MemAvailable:") {
let kib = rest.split_whitespace().next()?.parse::<u64>().ok()?;
return kib.checked_mul(1024);
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_local_executor_is_named_and_capable_of_both_workspaces() {
let executor = LocalExecutor;
assert_eq!(executor.name(), "local");
assert!(executor.capabilities().vcs_sessions);
}
#[test]
fn the_capacity_probe_reports_finite_numbers_on_any_host() {
let report = LocalExecutor.capacity();
assert!(
report.load1.is_finite() && report.load1 >= 0.0,
"{report:?}"
);
assert!(report.mem_free_bytes > 0, "{report:?}");
}
#[test]
fn a_cancellation_signal_is_shared_between_the_two_sides() {
let token = CancellationToken::new();
let observer = token.clone();
assert!(!observer.is_cancelled());
token.cancel();
assert!(
observer.is_cancelled(),
"the signal did not reach the dispatch"
);
assert_eq!(token, observer);
assert_ne!(CancellationToken::new(), observer);
}
#[test]
fn a_dispatch_request_carries_both_siblings_types() {
let request = DispatchRequest {
graph: ConfigRef("./graphs/node-scope.yaml".into()),
task: "## What\ndo it".into(),
labels: Labels::default(),
controls: NodeControls::default(),
workspace: WorkspaceSpec::VcsSession(SessionRequest {
repo: "owner/repo".into(),
branch: None,
base: None,
execution_checkout: None,
}),
cancel: CancellationToken::new(),
};
assert!(matches!(request.workspace, WorkspaceSpec::VcsSession(_)));
assert_eq!(request.graph.0, "./graphs/node-scope.yaml");
}
#[test]
fn the_drafting_dispatch_composes_nothing_onto_the_graph_the_launch_named() {
let root = std::env::temp_dir().join(format!("onepipeline-drafting-{}", crate::sys::pid()));
let _ = std::fs::remove_dir_all(&root);
let paths = crate::ledger::RunPaths::under(&root, "demo");
paths.create().expect("the run directory");
let record = r#"{"run_id":"demo","plan":"p.json","node_graph":"./node.yaml",
"pr_author_graph":"./author.yaml","launcher":"l","session":"s","pid":1,
"host":"h","started_at":"now","heartbeat_interval":1,
"node_sets":["members.worker.model=m"]}"#;
std::fs::write(paths.launch(), record).expect("the launch record is written");
std::env::set_var(crate::ledger::RUNS_DIR_ENV, &root);
let sets = |persona: &str| {
node_sets(
&Labels {
run_id: Some("demo".into()),
persona: Some(persona.into()),
..Labels::default()
},
&NodeControls::default(),
)
.expect("the launch record is readable")
};
assert!(
sets(crate::lifecycle::PR_AUTHOR_PERSONA).is_empty(),
"the drafting dispatch was given a member this graph never declared"
);
assert_eq!(
sets("engineer"),
vec![
"members.worker.model=m".to_string(),
"members.worker.persona=engineer".to_string(),
]
);
std::env::remove_var(crate::ledger::RUNS_DIR_ENV);
let _ = std::fs::remove_dir_all(&root);
}
#[test]
fn every_dispatch_is_given_a_directory_of_its_own_and_no_two_share_one() {
let root = std::env::temp_dir().join(format!("onepipeline-scratch-{}", crate::sys::pid()));
let _ = std::fs::remove_dir_all(&root);
std::env::set_var(crate::ledger::RUNS_DIR_ENV, &root);
let labels = Labels {
run_id: Some("demo".into()),
node: Some("build".into()),
..Labels::default()
};
let scratch = |labels: &Labels| {
let env = prepare_dispatch_env(labels).expect("the dispatch's environment is composed");
let (_, value) = env
.iter()
.find(|(key, _)| key == NODE_SCRATCH_DIR_ENV)
.expect("every dispatch carries a scratch directory")
.clone();
PathBuf::from(value)
};
let first = scratch(&labels);
let second = scratch(&labels);
assert_ne!(
first, second,
"a node asked again was handed the directory its first attempt had"
);
for at in [&first, &second] {
assert!(at.is_absolute(), "{} is not absolute", at.display());
assert!(at.is_dir(), "{} was not created", at.display());
std::fs::write(at.join("written"), "by the dispatch")
.unwrap_or_else(|error| panic!("{} is not writable: {error}", at.display()));
}
assert!(first.join("written").is_file());
assert!(scratch(&Labels::default()).is_dir());
let blocked = root.join("blocked");
std::fs::write(&blocked, "not a directory").expect("the blocking file is written");
std::env::set_var(crate::ledger::RUNS_DIR_ENV, &blocked);
assert!(matches!(
prepare_dispatch_env(&labels),
Err(Error::Ledger { .. })
));
std::env::remove_var(crate::ledger::RUNS_DIR_ENV);
let _ = std::fs::remove_dir_all(&root);
}
#[test]
fn a_dispatch_with_no_run_still_carries_its_nodes_own_controls() {
let sets = node_sets(
&Labels {
persona: Some("engineer".into()),
..Labels::default()
},
&NodeControls {
max_turns: std::num::NonZeroU32::new(45),
},
)
.expect("both are appliable");
assert_eq!(
sets,
vec![
"members.worker.persona=engineer".to_string(),
"members.worker.max_turns=45".to_string(),
],
"the node's own control must apply after the run-wide ones"
);
}
}