use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::{
atomic::{AtomicBool, AtomicU64, Ordering},
Arc, Mutex,
};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use ri_agent_graph::config::GraphConfig;
use ri_agent_graph::event_sink::GraphEvent;
use ri_agent_graph::state::{AgentState, StateLimits};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::sync::Notify;
use crate::compiler::{compile, CompileContext};
use crate::evidence::{bundle, digest, redact, validate_witness_dependencies};
use crate::spec::{ensure_size, GraphSpec, MAX_OUTPUT_BYTES, MAX_STATE_BYTES};
use crate::store::PersistentStore;
const MAX_RUNS: usize = 100;
const MAX_ACTIVE_RUNS: usize = 8;
fn is_terminal(status: &str) -> bool {
matches!(
status,
"completed" | "failed" | "cancelled" | "checkpointed"
)
}
#[derive(Clone)]
pub struct RunRecord {
pub run_id: String,
pub trace: String,
pub graph_id: String,
pub graph_version: String,
pub status: String,
pub success: Option<bool>,
pub input: Value,
pub state: Value,
pub final_state: Value,
pub steps: Vec<Value>,
pub error: Option<String>,
pub events: VecDeque<Value>,
pub next_cursor: u64,
pub dropped_events: u64,
pub receipt: Value,
pub bundle: Value,
pub persistence_status: String,
pub persistence_error: Option<String>,
pub budgets: Option<RunBudgets>,
pub budget_counters: BudgetCounters,
pub budget_exhausted: Option<String>,
pub checkpoint_id: Option<String>,
pub checkpoint_digest: Option<String>,
pub approval: Option<Value>,
pub resumed: bool,
pub cancelled: Arc<AtomicBool>,
pub cancellation: Arc<Notify>,
}
impl RunRecord {
pub fn public(&self) -> Value {
serde_json::json!({
"run_id":self.run_id,"trace":self.trace,"graph_id":self.graph_id,"graph_version":self.graph_version,
"storage_class": if self.persistence_status == "durable_terminal" { "sqlite_terminal_projection" } else { "volatile" },
"persistence_status":self.persistence_status,"persistence_error":self.persistence_error,
"status":self.status,"success":self.success,"final_state":self.final_state,
"state":self.state,"steps":self.steps,"error":self.error,"receipt":self.receipt,
"budgets":self.budgets,"budget_counters":self.budget_counters,
"budget_exhausted":self.budget_exhausted,
"checkpoint": self.checkpoint_id.as_ref().zip(self.checkpoint_digest.as_ref()).map(|(id, digest)| serde_json::json!({"checkpoint_id":id,"checkpoint_digest":digest})),
"replay_capability":self.receipt.get("replay_capability").and_then(Value::as_str).unwrap_or("integrity_only")
})
}
}
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct RunBudgets {
#[serde(skip_serializing_if = "Option::is_none")]
pub max_wall_clock_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_nodes: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_llm_calls: Option<u64>,
}
impl RunBudgets {
pub fn parse(value: Option<&Value>) -> Result<Option<Self>, String> {
let Some(value) = value.filter(|value| !value.is_null()) else {
return Ok(None);
};
let object = value
.as_object()
.ok_or_else(|| "budgets must be an object".to_owned())?;
if object.is_empty() {
return Err("budgets must contain at least one supported positive integer".into());
}
let mut budgets = Self {
max_wall_clock_ms: None,
max_nodes: None,
max_llm_calls: None,
};
for (key, raw) in object {
let slot = match key.as_str() {
"max_wall_clock_ms" => &mut budgets.max_wall_clock_ms,
"max_nodes" => &mut budgets.max_nodes,
"max_llm_calls" => &mut budgets.max_llm_calls,
_ => return Err(format!("unknown budget field '{key}'")),
};
let number = raw
.as_u64()
.filter(|number| *number > 0)
.ok_or_else(|| format!("budget '{key}' must be a positive integer"))?;
*slot = Some(number);
}
if budgets.max_llm_calls.is_some() {
return Err(
"max_llm_calls is unavailable: no real LLM invocation hook exists in the permitted runtime path"
.into(),
);
}
if budgets.max_wall_clock_ms.is_none() && budgets.max_nodes.is_none() {
return Err("budgets must contain an enforceable supported field".into());
}
Ok(Some(budgets))
}
pub fn requested_value(&self) -> Value {
serde_json::to_value(self).unwrap_or(Value::Null)
}
}
#[derive(Debug, Clone, Default, Serialize, PartialEq, Eq)]
pub struct BudgetCounters {
pub nodes: u64,
pub llm_calls: u64,
pub wall_clock_ms: u64,
}
#[derive(Clone)]
pub struct RunManager {
inner: Arc<Mutex<Inner>>,
counter: Arc<AtomicU64>,
}
struct Inner {
runs: HashMap<String, RunRecord>,
order: VecDeque<String>,
reserved_async_slots: usize,
}
impl Default for RunManager {
fn default() -> Self {
Self {
inner: Arc::new(Mutex::new(Inner {
runs: HashMap::new(),
order: VecDeque::new(),
reserved_async_slots: 0,
})),
counter: Arc::new(AtomicU64::new(1)),
}
}
}
impl RunManager {
pub fn allocate(
&self,
graph_id: &str,
graph_version: &str,
input: Value,
) -> Result<String, String> {
self.allocate_with_budgets(graph_id, graph_version, input, None)
}
pub fn allocate_with_budgets(
&self,
graph_id: &str,
graph_version: &str,
input: Value,
budgets: Option<RunBudgets>,
) -> Result<String, String> {
let n = self.counter.fetch_add(1, Ordering::SeqCst);
let millis = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
let run_id = format!("run-{millis:x}-{n:x}");
let record = RunRecord {
run_id: run_id.clone(),
trace: format!("trace-{millis:x}-{n:x}"),
graph_id: graph_id.into(),
graph_version: graph_version.into(),
status: "accepted".into(),
success: None,
input,
state: Value::Null,
final_state: Value::Null,
steps: vec![],
error: None,
events: VecDeque::new(),
next_cursor: 0,
dropped_events: 0,
receipt: Value::Null,
bundle: Value::Null,
persistence_status: "volatile_active".into(),
persistence_error: None,
budgets,
budget_counters: BudgetCounters::default(),
budget_exhausted: None,
checkpoint_id: None,
checkpoint_digest: None,
approval: None,
resumed: false,
cancelled: Arc::new(AtomicBool::new(false)),
cancellation: Arc::new(Notify::new()),
};
self.insert_record(record)
}
fn insert_record(&self, record: RunRecord) -> Result<String, String> {
let run_id = record.run_id.clone();
let mut inner = self.inner.lock().expect("run registry poisoned");
if inner.order.len() == MAX_RUNS {
let Some(old) = inner.order.iter().find_map(|id| {
inner
.runs
.get(id)
.filter(|run| is_terminal(&run.status))
.map(|_| id.clone())
}) else {
return Err(format!(
"run retention capacity reached: {MAX_RUNS} live runs cannot be evicted"
));
};
inner.order.retain(|id| id != &old);
inner.runs.remove(&old);
}
inner.order.push_back(run_id.clone());
inner.runs.insert(run_id.clone(), record);
Ok(run_id)
}
pub fn allocate_resumed(
&self,
run_id: &str,
graph_id: &str,
graph_version: &str,
input: Value,
state: Value,
budgets: Option<RunBudgets>,
checkpoint_id: &str,
checkpoint_digest: &str,
approval: Option<Value>,
) -> Result<String, String> {
let millis = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
let record = RunRecord {
run_id: run_id.to_owned(),
trace: format!("trace-{millis:x}-resume"),
graph_id: graph_id.into(),
graph_version: graph_version.into(),
status: "accepted".into(),
success: None,
input,
state,
final_state: Value::Null,
steps: vec![],
error: None,
events: VecDeque::new(),
next_cursor: 0,
dropped_events: 0,
receipt: Value::Null,
bundle: Value::Null,
persistence_status: "volatile_active".into(),
persistence_error: None,
budgets,
budget_counters: BudgetCounters::default(),
budget_exhausted: None,
checkpoint_id: Some(checkpoint_id.to_owned()),
checkpoint_digest: Some(checkpoint_digest.to_owned()),
approval,
resumed: true,
cancelled: Arc::new(AtomicBool::new(false)),
cancellation: Arc::new(Notify::new()),
};
self.insert_record(record)
}
pub fn mark_checkpointed(
&self,
id: &str,
checkpoint_id: &str,
checkpoint_digest: &str,
) -> Result<(), String> {
self.update(id, |run| {
run.status = "checkpointed".into();
run.checkpoint_id = Some(checkpoint_id.to_owned());
run.checkpoint_digest = Some(checkpoint_digest.to_owned());
})
}
pub fn admit_async(&self, id: &str) -> Result<(), String> {
let mut inner = self.inner.lock().map_err(|_| "run registry poisoned")?;
let active_runs = inner
.runs
.values()
.filter(|run| run.status == "running")
.count();
if active_runs + inner.reserved_async_slots >= MAX_ACTIVE_RUNS {
return Err(format!(
"active run capacity reached: {MAX_ACTIVE_RUNS} concurrent runs"
));
}
let run = inner.runs.get_mut(id).ok_or("run not found")?;
if run.status != "accepted" {
return Err(format!("run '{id}' is not awaiting admission"));
}
run.status = "running".into();
Ok(())
}
pub fn reserve_async_slot(&self) -> Result<(), String> {
let mut inner = self.inner.lock().map_err(|_| "run registry poisoned")?;
let active_runs = inner
.runs
.values()
.filter(|run| run.status == "running")
.count();
if active_runs + inner.reserved_async_slots >= MAX_ACTIVE_RUNS {
return Err(format!(
"active run capacity reached: {MAX_ACTIVE_RUNS} concurrent runs"
));
}
inner.reserved_async_slots += 1;
Ok(())
}
pub fn release_async_slot(&self) {
if let Ok(mut inner) = self.inner.lock() {
inner.reserved_async_slots = inner.reserved_async_slots.saturating_sub(1);
}
}
pub fn admit_reserved_async(&self, id: &str) -> Result<(), String> {
let mut inner = self.inner.lock().map_err(|_| "run registry poisoned")?;
if inner.reserved_async_slots == 0 {
return Err("no asynchronous execution slot was reserved".into());
}
let run = inner.runs.get_mut(id).ok_or("run not found")?;
if run.status != "accepted" {
return Err(format!("run '{id}' is not awaiting admission"));
}
run.status = "running".into();
inner.reserved_async_slots -= 1;
Ok(())
}
pub fn execute(
&self,
run_id: &str,
spec: GraphSpec,
base_url: String,
default_model: String,
) -> Result<Value, String> {
self.execute_with_store(run_id, spec, base_url, default_model, None)
}
pub fn execute_with_store(
&self,
run_id: &str,
spec: GraphSpec,
base_url: String,
default_model: String,
store: Option<PersistentStore>,
) -> Result<Value, String> {
self.execute_with_store_options(
run_id,
spec,
base_url,
default_model,
store,
None,
BudgetCounters::default(),
)
}
pub fn execute_with_store_options(
&self,
run_id: &str,
spec: GraphSpec,
base_url: String,
default_model: String,
store: Option<PersistentStore>,
initial_state: Option<Value>,
initial_counters: BudgetCounters,
) -> Result<Value, String> {
self.update(run_id, |r| r.status = "running".into())?;
let (
input,
cancelled,
cancellation,
budgets,
record_state,
checkpoint_id,
checkpoint_digest,
approval,
resumed,
) = {
let inner = self.inner.lock().map_err(|_| "run registry poisoned")?;
let r = inner.runs.get(run_id).ok_or("run not found")?;
(
r.input.clone(),
r.cancelled.clone(),
r.cancellation.clone(),
r.budgets.clone(),
r.state.clone(),
r.checkpoint_id.clone(),
r.checkpoint_digest.clone(),
r.approval.clone(),
r.resumed,
)
};
let started_at = Instant::now();
let provider = safe_provider_label(&base_url);
let configured_model = default_model.clone();
let events = Arc::new(Mutex::new(Vec::<GraphEvent>::new()));
let graph = compile(
&spec,
CompileContext {
base_url,
default_model,
cancelled: cancelled.clone(),
cancellation: cancellation.clone(),
events: events.clone(),
},
)?;
let starting_state = initial_state
.or_else(|| resumed.then_some(record_state))
.unwrap_or_else(|| initial_state_for_input(&input));
let Value::Object(map) = starting_state else {
return Err("checkpoint state must be a JSON object".into());
};
let initial = map.into_iter().collect::<HashMap<_, _>>();
let state = AgentState::with_data_and_limits(
initial,
StateLimits {
max_keys: 1000,
max_value_bytes: 256 * 1024,
max_history_len: 100,
lock_timeout: std::time::Duration::from_secs(5),
},
);
let snapshot = state.clone();
let rt = tokio::runtime::Runtime::new().map_err(|e| e.to_string())?;
let base_iteration_limit = spec.max_iterations.unwrap_or(64);
let budget_iteration_limit =
budgets
.as_ref()
.and_then(|budget| budget.max_nodes)
.map(|max_nodes| {
budget_iteration_limit(&spec, max_nodes.saturating_sub(initial_counters.nodes))
});
let config = GraphConfig::new()
.with_recursion_limit(
budget_iteration_limit
.unwrap_or(base_iteration_limit)
.min(base_iteration_limit),
)
.with_max_parallelism(spec.max_parallelism.unwrap_or(8));
let graph = Arc::new(graph);
let (handle, engine_cancel) =
rt.block_on(async { graph.execute_cancellable(&spec.entry, state, config) });
let timed_out = Arc::new(AtomicBool::new(false));
let timeout_notify = Arc::new(Notify::new());
let (timeout_complete, timeout_thread) = budgets
.as_ref()
.and_then(|budget| budget.max_wall_clock_ms)
.map(|limit| {
let (complete_tx, complete_rx) = std::sync::mpsc::channel::<()>();
let timed_out = timed_out.clone();
let timeout_notify = timeout_notify.clone();
let thread = std::thread::spawn(move || {
if matches!(
complete_rx.recv_timeout(Duration::from_millis(limit)),
Err(std::sync::mpsc::RecvTimeoutError::Timeout)
) {
timed_out.store(true, Ordering::SeqCst);
timeout_notify.notify_waiters();
}
});
(complete_tx, thread)
})
.map_or((None, None), |(complete_tx, thread)| {
(Some(complete_tx), Some(thread))
});
let mut handle = handle;
let result = rt.block_on(async {
if timeout_thread.is_some() {
tokio::select! {
joined = &mut handle => joined
.map_err(|error| format!("execution task failed: {error}"))?
.map_err(|error| error.to_string()),
_ = timeout_notify.notified() => {
engine_cancel.store(true, Ordering::SeqCst);
cancellation.notify_waiters();
handle.abort();
let _ = handle.await;
Err("BUDGET_EXHAUSTED".to_owned())
}
}
} else {
handle
.await
.map_err(|error| format!("execution task failed: {error}"))?
.map_err(|error| error.to_string())
}
});
if let Some(complete) = timeout_complete {
let _ = complete.send(());
}
if let Some(thread) = timeout_thread {
let _ = thread.join();
}
let core_receipt = serde_json::json!({
"api": "execute_with_summary",
"note": "receipt not yet fully typed; upgrade to execute_with_config + receipt"
});
let exported = rt.block_on(snapshot.export());
let state_value = serde_json::to_value(exported).map_err(|e| e.to_string())?;
ensure_size(&state_value, MAX_STATE_BYTES, "total state")?;
let final_state = spec
.output_key
.as_deref()
.and_then(|key| state_value.get(key).cloned())
.or_else(|| state_value.get("__input__").cloned())
.unwrap_or(Value::Null);
ensure_size(&final_state, MAX_OUTPUT_BYTES, "execution output")?;
let graph_events = events
.lock()
.map_err(|_| "event registry poisoned")?
.clone();
let mut steps: Vec<Value> = Vec::new();
for event in &graph_events {
if let GraphEvent::StateUpdate {
node_id, updates, ..
} = event
{
let output = updates
.get("__input__")
.or_else(|| updates.get("__route__"))
.cloned()
.unwrap_or_else(|| serde_json::to_value(updates).unwrap_or(Value::Null));
steps.push(serde_json::json!({
"node_id": node_id,
"output": output,
}));
}
}
let budget_counters = BudgetCounters {
nodes: initial_counters.nodes.saturating_add(
graph_events
.iter()
.filter(|event| matches!(event, GraphEvent::NodeStart { .. }))
.count() as u64,
),
llm_calls: initial_counters.llm_calls,
wall_clock_ms: initial_counters
.wall_clock_ms
.saturating_add(started_at.elapsed().as_millis().min(u64::MAX as u128) as u64),
};
let wall_exhausted = budgets.as_ref().is_some_and(|budget| {
budget
.max_wall_clock_ms
.is_some_and(|limit| budget_counters.wall_clock_ms >= limit)
|| timed_out.load(Ordering::SeqCst)
});
let node_exhausted = budgets.as_ref().and_then(|budget| {
let limit = budget.max_nodes?;
let iteration_error = result
.as_ref()
.err()
.is_some_and(|error| error.contains("iteration"));
let budget_limited_iterations =
budget_iteration_limit.is_some_and(|limit| limit <= base_iteration_limit);
(budget_counters.nodes > limit
|| (budget_counters.nodes >= limit && iteration_error && budget_limited_iterations))
.then_some("max_nodes".into())
});
let budget_exhausted = if wall_exhausted {
Some("max_wall_clock_ms".to_owned())
} else {
node_exhausted
};
let mut dependency_envelopes = Value::Array(Vec::new());
let mut dependency_envelopes_complete = false;
let mut evidence_error = None;
if result.is_ok() && spec.nodes.iter().any(|node| node.evidence_required) {
match store.as_ref() {
Some(store) => {
let mut collected = Vec::new();
for node in spec.nodes.iter().filter(|node| node.evidence_required) {
let output_key = node
.config
.get("output_key")
.and_then(Value::as_str)
.unwrap_or("");
let Some(evidence) = state_value.get(output_key) else {
evidence_error = Some("WITNESS_EVIDENCE_MISSING".to_owned());
break;
};
match validate_witness_dependencies(evidence, store) {
Ok(Value::Array(dependencies)) => collected.extend(dependencies),
Ok(_) => {
evidence_error = Some("WITNESS_EVIDENCE_INVALID".to_owned());
break;
}
Err(error) => {
evidence_error = Some(error.code);
break;
}
}
}
if evidence_error.is_none() {
let mut unique = std::collections::BTreeMap::new();
for dependency in collected {
if let Some(id) = dependency.get("witness_id").and_then(Value::as_str) {
unique.insert(id.to_owned(), dependency);
}
}
dependency_envelopes = Value::Array(unique.into_values().collect());
dependency_envelopes_complete = true;
}
}
None => evidence_error = Some("WITNESS_STORE_REQUIRED".to_owned()),
}
}
let error = if budget_exhausted.is_some() {
Some("BUDGET_EXHAUSTED".to_owned())
} else if evidence_error.is_some() {
evidence_error
} else {
result.err().map(|e| e.to_string())
};
let trace = self.get(run_id).ok_or("run not found")?.trace;
let graph_version = self.get(run_id).ok_or("run not found")?.graph_version;
let models: Vec<Value> = spec.nodes.iter().filter(|node| matches!(node.node_type, crate::spec::NodeType::Llm)).map(|node| serde_json::json!({
"node_id":node.id,"model_alias":node.model.as_deref().unwrap_or("server_default"),"prompt_digest":digest(&Value::String(node.prompt.clone().unwrap_or_else(||"{input}".into())))
})).collect();
let model_labels: Vec<Value> = models
.iter()
.filter_map(|m| m.get("model_alias").cloned())
.collect();
let receipt = serde_json::json!({"schema":"agent-graph-mcp-receipt-v1","run_id":run_id,"trace":trace,"graph_version":graph_version,
"input_digest":digest(&input),"output_digest":digest(&state_value),"step_count":steps.len(),"models":models,
"provider":provider,"default_model":configured_model,"model_labels":model_labels,
"core":core_receipt,"dependency_envelopes":dependency_envelopes,"dependency_envelopes_complete":dependency_envelopes_complete,"replay_capability":if resumed { "deterministic_local_resume" } else { "integrity_only" },
"resume_supported":resumed,
"checkpoint":checkpoint_id.as_ref().zip(checkpoint_digest.as_ref()).map(|(id, digest)| serde_json::json!({"checkpoint_id":id,"checkpoint_digest":digest})),
"approval":approval,
"evidence_authority":if dependency_envelopes_complete { "local_capture_receipt_only; source_authority_not_verified" } else { "structural_unverified" },"persistence_status":"pending",
"budgets":budgets.as_ref().map(RunBudgets::requested_value).unwrap_or(Value::Null),
"budget_counters":budget_counters,"budget_exhausted":budget_exhausted});
let artifact = bundle(run_id, &graph_version, &input, &state_value, &receipt);
self.update(run_id, |r| {
let cancellation_observed = r.cancelled.load(Ordering::SeqCst);
let terminal_error = (!cancellation_observed).then(|| error.clone()).flatten();
let terminal = terminal_outcome(cancellation_observed, terminal_error.as_deref());
r.status = terminal.status.into();
r.success = Some(terminal.success);
r.state = state_value.clone();
r.final_state = final_state.clone();
r.steps = steps.clone();
r.error = terminal_error;
r.budget_counters = budget_counters.clone();
r.budget_exhausted = budget_exhausted.clone();
r.receipt = receipt.clone();
r.bundle = artifact.clone();
r.persistence_status = "pending".into();
for event in graph_events {
push_event(r, serde_json::to_value(event).unwrap_or(Value::Null));
}
})?;
Ok(self.get(run_id).expect("updated run").public())
}
pub fn start(&self, run_id: String, spec: GraphSpec, base_url: String, model: String) {
self.start_with_completion(run_id, spec, base_url, model, |_| {});
}
pub fn start_with_completion<F>(
&self,
run_id: String,
spec: GraphSpec,
base_url: String,
model: String,
on_completion: F,
) where
F: FnOnce(RunRecord) + Send + 'static,
{
self.start_with_completion_with_store(run_id, spec, base_url, model, None, on_completion);
}
pub fn start_with_completion_with_store<F>(
&self,
run_id: String,
spec: GraphSpec,
base_url: String,
model: String,
store: Option<PersistentStore>,
on_completion: F,
) where
F: FnOnce(RunRecord) + Send + 'static,
{
let manager = self.clone();
std::thread::spawn(move || {
if let Err(error) = manager.execute_with_store(&run_id, spec, base_url, model, store) {
let _ = manager.update(&run_id, |r| {
r.status = "failed".into();
r.success = Some(false);
r.error = Some(error.clone());
});
}
if let Some(record) = manager.get(&run_id) {
on_completion(record);
}
});
}
pub fn start_resumed_with_completion<F>(
&self,
run_id: String,
spec: GraphSpec,
base_url: String,
model: String,
store: Option<PersistentStore>,
on_completion: F,
) where
F: FnOnce(RunRecord) + Send + 'static,
{
let manager = self.clone();
std::thread::spawn(move || {
if let Err(error) = manager.execute_with_store_options(
&run_id,
spec,
base_url,
model,
store,
None,
BudgetCounters::default(),
) {
let _ = manager.update(&run_id, |r| {
r.status = "failed".into();
r.success = Some(false);
r.error = Some(error.clone());
});
}
if let Some(record) = manager.get(&run_id) {
on_completion(record);
}
});
}
pub fn cancel(&self, id: &str) -> Result<Value, String> {
let cancellation = {
let mut inner = self.inner.lock().map_err(|_| "run registry poisoned")?;
let r = inner.runs.get_mut(id).ok_or("run not found")?;
if is_terminal(&r.status) {
return Err("RUN_NOT_CANCELLABLE".into());
}
r.cancelled.store(true, Ordering::SeqCst);
r.cancellation.clone()
};
cancellation.notify_waiters();
Ok(serde_json::json!({
"run_id":id,
"status":"cancellation_requested",
"cancellation_effect":"best_effort_drop_provider_future",
"provider_request_may_still_be_in_flight":true,
"effective_at":"provider_completion_or_cancellation_observation"
}))
}
pub fn get(&self, id: &str) -> Option<RunRecord> {
self.inner.lock().ok()?.runs.get(id).cloned()
}
pub fn mark_persistence(&self, id: &str, status: &str, error: Option<String>) {
let _ = self.update(id, |run| {
run.persistence_status = status.into();
run.persistence_error = error.clone();
if let Some(object) = run.receipt.as_object_mut() {
object.insert("persistence_status".into(), Value::String(status.into()));
if let Some(error) = error {
object.insert("persistence_error".into(), Value::String(error));
}
}
run.bundle = bundle(
&run.run_id,
&run.graph_version,
&run.input,
&run.state,
&run.receipt,
);
});
}
pub(crate) fn remove(&self, id: &str) -> Option<RunRecord> {
let mut inner = self.inner.lock().ok()?;
if inner
.runs
.get(id)
.is_some_and(|run| run.status == "running")
{
return None;
}
inner.order.retain(|run_id| run_id != id);
inner.runs.remove(id)
}
pub fn list(&self) -> Vec<Value> {
self.inner
.lock()
.map(|i| {
i.order
.iter()
.filter_map(|id| i.runs.get(id).map(RunRecord::public))
.collect()
})
.unwrap_or_default()
}
pub fn events(
&self,
store: Option<&PersistentStore>,
id: &str,
cursor: u64,
limit: usize,
) -> Result<Value, String> {
let Some(r) = self.get(id) else {
if let Some(store) = store {
if let Some(events) = store.load_events(id, cursor, limit)? {
return Ok(events);
}
}
return Err("run not found".into());
};
let first = r
.events
.front()
.and_then(|v| v.get("cursor"))
.and_then(Value::as_u64)
.unwrap_or(r.next_cursor);
let events: Vec<_> = r
.events
.iter()
.filter(|v| v["cursor"].as_u64().unwrap_or(0) >= cursor)
.take(limit.min(200))
.cloned()
.collect();
Ok(
serde_json::json!({"run_id":id,"events":events,"next_cursor":r.next_cursor,"gap":cursor<first,"truncated":r.dropped_events>0,"dropped":r.dropped_events}),
)
}
pub fn set_state_value(&self, id: &str, key: &str, value: Value) -> Result<(), String> {
self.update(id, |run| {
if let Some(state) = run.state.as_object_mut() {
state.insert(key.to_owned(), value.clone());
} else {
let mut state = serde_json::Map::new();
state.insert(key.to_owned(), value);
run.state = Value::Object(state);
}
})
}
fn update(&self, id: &str, f: impl FnOnce(&mut RunRecord)) -> Result<(), String> {
let mut inner = self.inner.lock().map_err(|_| "run registry poisoned")?;
let r = inner.runs.get_mut(id).ok_or("run not found")?;
f(r);
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct TerminalOutcome {
status: &'static str,
success: bool,
}
fn terminal_outcome(cancelled: bool, error: Option<&str>) -> TerminalOutcome {
if cancelled {
TerminalOutcome {
status: "cancelled",
success: false,
}
} else if error.is_some() {
TerminalOutcome {
status: "failed",
success: false,
}
} else {
TerminalOutcome {
status: "completed",
success: true,
}
}
}
/// Return the greatest number of engine supersteps whose conservative frontier
/// size stays within a node budget. The engine checks recursion before entering
/// a superstep, so this prevents a later parallel frontier from starting after
/// the budget is exhausted. Conditional edges are treated as possible edges;
/// this can stop early, but never allows an over-budget frontier.
fn budget_iteration_limit(spec: &GraphSpec, max_nodes: u64) -> usize {
let graph_limit = spec.max_iterations.unwrap_or(64);
let mut frontier = vec![spec.entry.clone()];
let mut observed = 0u64;
let mut allowed = 0usize;
for _ in 0..graph_limit {
frontier.retain(|node| node != "END");
if frontier.is_empty() {
break;
}
let width = frontier.len() as u64;
if observed.saturating_add(width) > max_nodes {
break;
}
observed = observed.saturating_add(width);
allowed += 1;
let current: HashSet<&str> = frontier.iter().map(String::as_str).collect();
let mut next = Vec::new();
for edge in &spec.edges {
if current.contains(edge.from.as_str()) && edge.to != "END" {
next.push(edge.to.clone());
}
}
let mut seen = HashSet::new();
next.retain(|node| seen.insert(node.clone()));
frontier = next;
}
allowed.max(1)
}
fn push_event(run: &mut RunRecord, event: Value) {
if run.events.len() == 512 {
run.events.pop_front();
run.dropped_events += 1;
}
let cursor = run.next_cursor;
run.next_cursor += 1;
run.events
.push_back(serde_json::json!({"cursor":cursor,"event":redact(&event)}));
}
fn safe_provider_label(url: &str) -> String {
let without_fragment = url.split(['?', '#']).next().unwrap_or(url);
if let Some((scheme, rest)) = without_fragment.split_once(": let authority_and_path = rest.rsplit_once('@').map(|(_, safe)| safe).unwrap_or(rest);
format!("{scheme}://{authority_and_path}")
} else {
"server-configured".into()
}
}
pub fn initial_state_for_input(input: &Value) -> Value {
let mut initial = serde_json::Map::new();
initial.insert("__input__".into(), input.clone());
if let Value::Object(map) = input {
initial.extend(map.clone());
}
Value::Object(initial)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn async_admission_is_bounded_atomically() {
let manager = RunManager::default();
for index in 0..MAX_ACTIVE_RUNS {
let id = manager
.allocate("graph", "version", serde_json::json!({"index": index}))
.expect("allocate admitted run");
manager.admit_async(&id).expect("admit within cap");
}
let overflow = manager
.allocate("graph", "version", Value::Null)
.expect("registry still has room");
assert!(manager.admit_async(&overflow).is_err());
manager.remove(&overflow);
assert!(manager.get(&overflow).is_none());
}
#[test]
fn retention_never_evicts_live_runs() {
let manager = RunManager::default();
let mut ids = Vec::new();
for _ in 0..MAX_RUNS {
ids.push(
manager
.allocate("graph", "version", Value::Null)
.expect("allocate retained run"),
);
}
assert!(manager.allocate("graph", "version", Value::Null).is_err());
manager
.update(&ids[0], |record| record.status = "completed".into())
.expect("mark terminal");
let replacement = manager
.allocate("graph", "version", Value::Null)
.expect("terminal record may be evicted");
assert!(manager.get(&ids[0]).is_none());
assert!(manager.get(&replacement).is_some());
}
#[test]
fn cancelled_terminal_run_is_never_successful() {
let outcome = terminal_outcome(true, None);
assert_eq!(outcome.status, "cancelled");
assert!(!outcome.success);
let completed = terminal_outcome(false, None);
assert_eq!(completed.status, "completed");
assert!(completed.success);
let failed = terminal_outcome(false, Some("provider error"));
assert_eq!(failed.status, "failed");
assert!(!failed.success);
}
#[test]
fn cancellation_is_only_a_request_until_a_node_boundary() {
let manager = RunManager::default();
let id = manager
.allocate("graph", "version", Value::Null)
.expect("allocate run");
manager.admit_async(&id).expect("admit run");
let response = manager.cancel(&id).expect("request cancellation");
assert_eq!(response["status"], "cancellation_requested");
assert_eq!(
response["cancellation_effect"],
"best_effort_drop_provider_future"
);
assert_eq!(manager.get(&id).expect("run").status, "running");
}
#[test]
fn persistence_failure_is_publicly_volatile_and_recorded() {
let manager = RunManager::default();
let id = manager
.allocate("graph", "version", Value::Null)
.expect("allocate run");
manager.mark_persistence(
&id,
"volatile_persistence_failed",
Some("database is unavailable".into()),
);
let public = manager.get(&id).expect("run").public();
assert_eq!(public["storage_class"], "volatile");
assert_eq!(public["persistence_error"], "database is unavailable");
}
}