1use super::{
4 BTreeMap, BTreeSet, BoundedWorkGraph, BudgetAllocation, ChainBudget, ChildExecutionReceipt,
5 ChildOutcome, FusedResults, MAX_DEPTH, MAX_FAN_OUT, MAX_GRAPH_NODES, MAX_LOCAL_CONCURRENCY,
6 MAX_PATH_CLAIMS, NodeStatus, ResultConflict, StopReason, TeamValueReport, WorkAttribution,
7 WorkGraphError, WorkNode, WorkNodeBudget, cascade_budget, execution_fence, validate_budget,
8 validate_cascade, validate_identifier, validate_path_claim, validate_receipt,
9};
10#[cfg(test)]
11use super::{MeasuredExecutionSpend, validate_measured_spend};
12
13impl BoundedWorkGraph {
14 #[must_use]
15 pub fn new(max_fan_out: usize, max_depth: u16) -> Self {
16 Self {
17 graph_id: String::new(),
18 nodes: BTreeMap::new(),
19 children: BTreeMap::new(),
20 chain_budgets: BTreeMap::new(),
21 pending_child_budgets: BTreeMap::new(),
22 pending_child_parents: BTreeMap::new(),
23 receipts: BTreeMap::new(),
24 failed_spend: BTreeMap::new(),
25 reaped_executions: BTreeMap::new(),
26 accepted_path: BTreeSet::new(),
27 max_local_concurrency: MAX_FAN_OUT,
28 max_fan_out: max_fan_out.clamp(1, MAX_FAN_OUT),
29 max_depth: max_depth.clamp(1, MAX_DEPTH),
30 }
31 }
32
33 #[must_use]
34 pub fn with_local_concurrency(mut self, max_local_concurrency: usize) -> Self {
35 self.max_local_concurrency = max_local_concurrency.clamp(1, MAX_LOCAL_CONCURRENCY);
36 self
37 }
38
39 pub fn with_graph_id(mut self, graph_id: &str) -> Result<Self, WorkGraphError> {
40 validate_identifier(graph_id, "graph_id")?;
41 self.graph_id = graph_id.to_string();
42 Ok(self)
43 }
44
45 pub fn bind_graph_id(&mut self, graph_id: &str) -> Result<(), WorkGraphError> {
46 validate_identifier(graph_id, "graph_id")?;
47 if !self.graph_id.is_empty() && self.graph_id != graph_id {
48 return Err(WorkGraphError::InvalidGraph("graph id mismatch".into()));
49 }
50 self.graph_id = graph_id.to_string();
51 Ok(())
52 }
53
54 pub fn graph_id(&self) -> &str {
55 &self.graph_id
56 }
57
58 pub fn observation(&self) -> serde_json::Value {
61 self.observe_if_changed(None)
62 }
63
64 pub fn observe_if_changed(&self, if_revision: Option<&str>) -> serde_json::Value {
67 let nodes: Vec<_> = self
68 .nodes
69 .values()
70 .map(|node| {
71 serde_json::json!({
72 "node_id": node.node_id,
73 "agent_id": node.agent_id,
74 "parent_node_id": node.parent_node_id,
75 "status": node.status,
76 "execution_started": node.execution_started,
77 "process_exit": if !node.execution_started { "not_started" }
78 else if self.execution_exit_unconfirmed(node) { "unconfirmed" }
79 else { "confirmed" },
80 "stop_reason": node.stop_reason,
81 "tokens_consumed": node.budget.tokens_consumed,
82 "tokens_remaining": node.budget.tokens_remaining(),
83 "cost_micros_consumed": node.budget.cost_micros_consumed,
84 "cost_micros_remaining": node.budget.cost_remaining(),
85 "accepted": self.accepted_path.contains(&node.node_id),
86 })
87 })
88 .collect();
89 let mut observation = serde_json::json!({
90 "schema_version": 1,
91 "graph_id": self.graph_id,
92 "max_concurrency": self.local_concurrency_limit(),
93 "nodes": nodes,
94 });
95 let revision = blake3::hash(observation.to_string().as_bytes())
96 .to_hex()
97 .to_string();
98 if if_revision == Some(revision.as_str()) {
99 return serde_json::json!({
100 "schema_version": 1,
101 "graph_id": self.graph_id,
102 "revision": revision,
103 "unchanged": true,
104 });
105 }
106 observation["revision"] = revision.into();
107 observation["unchanged"] = false.into();
108 observation
109 }
110
111 pub(super) fn local_concurrency_limit(&self) -> usize {
112 self.max_local_concurrency.clamp(1, MAX_LOCAL_CONCURRENCY)
113 }
114
115 pub fn add_root(
117 &mut self,
118 node_id: String,
119 agent_id: String,
120 capsule_ref: String,
121 budget: WorkNodeBudget,
122 ) -> Result<&WorkNode, WorkGraphError> {
123 if self.nodes.len() >= MAX_GRAPH_NODES {
124 return Err(WorkGraphError::CapacityExceeded);
125 }
126 validate_identifier(&node_id, "node_id")?;
127 validate_identifier(&agent_id, "agent_id")?;
128 validate_identifier(&capsule_ref, "capsule_ref")?;
129 validate_budget(&budget)?;
130 let concurrency_limit = self.local_concurrency_limit();
131 if self.active_count() >= concurrency_limit {
132 return Err(WorkGraphError::LocalConcurrencyExceeded(concurrency_limit));
133 }
134 if self.nodes.contains_key(&node_id) {
135 return Err(WorkGraphError::DuplicateNode(node_id));
136 }
137 let node = WorkNode {
138 node_id: node_id.clone(),
139 agent_id,
140 parent_node_id: None,
141 capsule_ref,
142 status: NodeStatus::Active,
143 budget,
144 depth: 0,
145 stop_reason: None,
146 outcome_ref: None,
147 claim_attempt: 0,
148 delivery_attempt: None,
149 execution_fence: None,
150 lease_expires_epoch_ms: None,
151 path_claims: Vec::new(),
152 task_ref: String::new(),
153 parent_task_ref: None,
154 policy_ref: String::new(),
155 expected_outcome_ref: String::new(),
156 execution_started: false,
157 };
158 self.nodes.insert(node_id.clone(), node);
159 self.chain_budgets.insert(
160 node_id.clone(),
161 ChainBudget {
162 chain_id: node_id.clone(),
163 root_budget_tokens: self.nodes[&node_id].budget.tokens_allocated,
164 total_consumed_tokens: 0,
165 total_allocated_tokens: self.nodes[&node_id].budget.tokens_allocated,
166 root_budget_cost_micros: self.nodes[&node_id].budget.cost_micros_allocated,
167 total_consumed_cost_micros: 0,
168 total_allocated_cost_micros: self.nodes[&node_id].budget.cost_micros_allocated,
169 unauthorized_tokens: 0,
170 unauthorized_cost_micros: 0,
171 depth: 0,
172 },
173 );
174 Ok(self.nodes.get(&node_id).expect("root node inserted above"))
175 }
176
177 pub fn delegate(
179 &mut self,
180 parent_node_id: &str,
181 child_node_id: String,
182 child_agent_id: String,
183 capsule_ref: String,
184 child_budget: WorkNodeBudget,
185 ) -> Result<&WorkNode, WorkGraphError> {
186 self.insert_child(
187 parent_node_id,
188 child_node_id,
189 child_agent_id,
190 capsule_ref,
191 child_budget,
192 NodeStatus::Active,
193 )
194 }
195
196 pub fn queue_child(
198 &mut self,
199 parent_node_id: &str,
200 child_node_id: String,
201 child_agent_id: String,
202 capsule_ref: String,
203 child_budget: WorkNodeBudget,
204 ) -> Result<&WorkNode, WorkGraphError> {
205 self.insert_child(
206 parent_node_id,
207 child_node_id,
208 child_agent_id,
209 capsule_ref,
210 child_budget,
211 NodeStatus::Pending,
212 )
213 }
214
215 pub fn claim_and_activate(&mut self, node_id: &str) -> Result<&WorkNode, WorkGraphError> {
217 self.claim_and_activate_until(node_id, None)
218 }
219
220 pub fn claim_and_activate_until(
221 &mut self,
222 node_id: &str,
223 lease_expires_epoch_ms: Option<u64>,
224 ) -> Result<&WorkNode, WorkGraphError> {
225 validate_identifier(node_id, "node_id")?;
226 let concurrency_limit = self.local_concurrency_limit();
227 if self.active_count() >= concurrency_limit {
228 return Err(WorkGraphError::LocalConcurrencyExceeded(concurrency_limit));
229 }
230 let parent_id = self
231 .nodes
232 .get(node_id)
233 .ok_or_else(|| WorkGraphError::NodeNotFound(node_id.to_string()))?
234 .parent_node_id
235 .clone();
236 let Some(parent_id) = parent_id else {
237 return Err(WorkGraphError::InvalidTransition(node_id.to_string()));
238 };
239 if self.nodes.get(&parent_id).map(|node| node.status) != Some(NodeStatus::Active) {
240 return Err(WorkGraphError::ParentNotActive(parent_id));
241 }
242 let graph_id = self.graph_id.clone();
243 let node = self.nodes.get_mut(node_id).expect("node checked above");
244 if node.status != NodeStatus::Pending {
245 return Err(WorkGraphError::InvalidTransition(node_id.to_string()));
246 }
247 node.claim_attempt = node
248 .claim_attempt
249 .checked_add(1)
250 .ok_or(WorkGraphError::Overflow)?;
251 node.execution_fence = Some(execution_fence(&graph_id, node));
252 node.delivery_attempt = None;
253 node.lease_expires_epoch_ms = lease_expires_epoch_ms;
254 node.status = NodeStatus::Active;
255 Ok(node)
256 }
257
258 pub fn claim_or_resume_active_until(
260 &mut self,
261 node_id: &str,
262 lease_expires_epoch_ms: u64,
263 ) -> Result<&WorkNode, WorkGraphError> {
264 let capsule_ref = self
265 .nodes
266 .get(node_id)
267 .ok_or_else(|| WorkGraphError::NodeNotFound(node_id.to_string()))?
268 .capsule_ref
269 .clone();
270 if self.nodes.iter().any(|(candidate_id, candidate)| {
271 candidate_id != node_id
272 && candidate.execution_started
273 && candidate.capsule_ref == capsule_ref
274 }) {
275 return Err(WorkGraphError::InvalidGraph(
276 "context capsule was already consumed by another node".into(),
277 ));
278 }
279 match self.nodes.get(node_id).map(|node| node.status) {
280 Some(NodeStatus::Pending) => {
281 self.claim_and_activate_until(node_id, Some(lease_expires_epoch_ms))?;
282 let node = self.nodes.get_mut(node_id).expect("node claimed above");
283 node.execution_started = true;
284 Ok(node)
285 }
286 Some(NodeStatus::Active) => {
287 let node = self.nodes.get_mut(node_id).expect("node checked above");
288 if node.parent_node_id.is_none()
289 || node.execution_fence.is_none()
290 || node.execution_started
291 {
292 return Err(WorkGraphError::InvalidTransition(node_id.to_string()));
293 }
294 node.lease_expires_epoch_ms = Some(lease_expires_epoch_ms);
295 node.execution_started = true;
296 Ok(node)
297 }
298 Some(_) => Err(WorkGraphError::InvalidTransition(node_id.to_string())),
299 None => Err(WorkGraphError::NodeNotFound(node_id.to_string())),
300 }
301 }
302
303 pub fn recover_expired_execution(
304 &mut self,
305 node_id: &str,
306 now_epoch_ms: u64,
307 ) -> Result<Vec<String>, WorkGraphError> {
308 let node = self
309 .nodes
310 .get(node_id)
311 .ok_or_else(|| WorkGraphError::NodeNotFound(node_id.to_string()))?;
312 let expired = node.parent_node_id.is_some()
313 && node.status == NodeStatus::Active
314 && node
315 .lease_expires_epoch_ms
316 .is_some_and(|expires| now_epoch_ms >= expires);
317 if expired {
318 return self.stop(node_id, StopReason::LeaseLost);
319 }
320 Ok(Vec::new())
321 }
322
323 pub fn set_path_claims(
325 &mut self,
326 node_id: &str,
327 mut path_claims: Vec<String>,
328 ) -> Result<(), WorkGraphError> {
329 if path_claims.len() > MAX_PATH_CLAIMS {
330 return Err(WorkGraphError::TooManyPathClaims);
331 }
332 path_claims.sort();
333 path_claims.dedup();
334 for claim in &path_claims {
335 validate_path_claim(claim)?;
336 }
337 let node = self
338 .nodes
339 .get_mut(node_id)
340 .ok_or_else(|| WorkGraphError::NodeNotFound(node_id.to_string()))?;
341 if node.status != NodeStatus::Pending {
342 return Err(WorkGraphError::InvalidTransition(node_id.to_string()));
343 }
344 node.path_claims = path_claims;
345 Ok(())
346 }
347
348 pub fn configure_root_context(
349 &mut self,
350 node_id: &str,
351 task_ref: &str,
352 policy_ref: &str,
353 expected_outcome_ref: &str,
354 ) -> Result<(), WorkGraphError> {
355 validate_identifier(task_ref, "task_ref")?;
356 validate_identifier(policy_ref, "policy_ref")?;
357 validate_identifier(expected_outcome_ref, "expected_outcome_ref")?;
358 let node = self
359 .nodes
360 .get_mut(node_id)
361 .ok_or_else(|| WorkGraphError::NodeNotFound(node_id.to_string()))?;
362 if node.parent_node_id.is_some() {
363 return Err(WorkGraphError::InvalidTransition(node_id.to_string()));
364 }
365 node.task_ref = task_ref.to_string();
366 node.parent_task_ref = None;
367 node.policy_ref = policy_ref.to_string();
368 node.expected_outcome_ref = expected_outcome_ref.to_string();
369 Ok(())
370 }
371
372 pub fn configure_child_context(
373 &mut self,
374 node_id: &str,
375 task_ref: &str,
376 expected_outcome_ref: &str,
377 ) -> Result<(), WorkGraphError> {
378 validate_identifier(task_ref, "task_ref")?;
379 validate_identifier(expected_outcome_ref, "expected_outcome_ref")?;
380 let parent_id = self
381 .nodes
382 .get(node_id)
383 .and_then(|node| node.parent_node_id.clone())
384 .ok_or_else(|| WorkGraphError::InvalidTransition(node_id.to_string()))?;
385 let parent = self
386 .nodes
387 .get(&parent_id)
388 .ok_or_else(|| WorkGraphError::NodeNotFound(parent_id.clone()))?;
389 if parent.task_ref.is_empty() || parent.policy_ref.is_empty() {
390 return Err(WorkGraphError::InvalidGraph(
391 "parent context is not configured".into(),
392 ));
393 }
394 let parent_task_ref = parent.task_ref.clone();
395 let policy_ref = parent.policy_ref.clone();
396 let node = self
397 .nodes
398 .get_mut(node_id)
399 .ok_or_else(|| WorkGraphError::NodeNotFound(node_id.to_string()))?;
400 if node.status != NodeStatus::Pending {
401 return Err(WorkGraphError::InvalidTransition(node_id.to_string()));
402 }
403 node.task_ref = task_ref.to_string();
404 node.parent_task_ref = Some(parent_task_ref);
405 node.policy_ref = policy_ref;
406 node.expected_outcome_ref = expected_outcome_ref.to_string();
407 Ok(())
408 }
409
410 fn insert_child(
411 &mut self,
412 parent_node_id: &str,
413 child_node_id: String,
414 child_agent_id: String,
415 capsule_ref: String,
416 child_budget: WorkNodeBudget,
417 status: NodeStatus,
418 ) -> Result<&WorkNode, WorkGraphError> {
419 if self.nodes.len() >= MAX_GRAPH_NODES {
420 return Err(WorkGraphError::CapacityExceeded);
421 }
422 validate_identifier(parent_node_id, "parent_node_id")?;
423 validate_identifier(&child_node_id, "child_node_id")?;
424 validate_identifier(&child_agent_id, "child_agent_id")?;
425 validate_identifier(&capsule_ref, "capsule_ref")?;
426 validate_budget(&child_budget)?;
427 if self.nodes.contains_key(&child_node_id) {
428 return Err(WorkGraphError::DuplicateNode(child_node_id));
429 }
430 let pending_budget = self.pending_child_budgets.get(&child_node_id);
431 let pending_parent = self.pending_child_parents.get(&child_node_id);
432 let pending_reservation_matches = match (pending_budget, pending_parent) {
433 (None, None) => true,
434 (Some(_), Some(parent)) => parent == parent_node_id,
435 _ => false,
436 };
437 if !pending_reservation_matches {
438 return Err(WorkGraphError::InvalidGraph(
439 "pending child reservation is incomplete or belongs to another parent".into(),
440 ));
441 }
442 let (reserved_tokens, reserved_cost) =
443 self.reserved_child_budget(parent_node_id, Some(&child_node_id))?;
444 let parent = self
445 .nodes
446 .get(parent_node_id)
447 .ok_or_else(|| WorkGraphError::NodeNotFound(parent_node_id.to_string()))?;
448 if parent.status != NodeStatus::Active {
449 return Err(WorkGraphError::ParentNotActive(parent_node_id.to_string()));
450 }
451 let new_depth = parent.depth + 1;
452 if new_depth > self.max_depth {
453 return Err(WorkGraphError::DepthExceeded(self.max_depth));
454 }
455 let available_tokens = parent
456 .budget
457 .tokens_remaining()
458 .saturating_sub(reserved_tokens);
459 let available_cost = parent.budget.cost_remaining().saturating_sub(reserved_cost);
460 if child_budget.tokens_allocated > available_tokens {
461 return Err(WorkGraphError::BudgetExceedsParent {
462 child_requested: child_budget.tokens_allocated,
463 parent_remaining: available_tokens,
464 });
465 }
466 if child_budget.cost_micros_allocated > available_cost {
467 return Err(WorkGraphError::BudgetExceedsParent {
468 child_requested: child_budget.cost_micros_allocated,
469 parent_remaining: available_cost,
470 });
471 }
472 let current_children = self.children.get(parent_node_id).map_or(0, Vec::len);
473 if current_children >= self.max_fan_out {
474 return Err(WorkGraphError::FanOutExceeded(self.max_fan_out));
475 }
476 let concurrency_limit = self.local_concurrency_limit();
477 if status == NodeStatus::Active && self.active_count() >= concurrency_limit {
478 return Err(WorkGraphError::LocalConcurrencyExceeded(concurrency_limit));
479 }
480 let child_tokens_allocated = child_budget.tokens_allocated;
481 let child_cost_allocated = child_budget.cost_micros_allocated;
482 let node = WorkNode {
483 node_id: child_node_id.clone(),
484 agent_id: child_agent_id,
485 parent_node_id: Some(parent_node_id.to_string()),
486 capsule_ref,
487 status,
488 budget: child_budget,
489 depth: new_depth,
490 stop_reason: None,
491 outcome_ref: None,
492 claim_attempt: 0,
493 delivery_attempt: None,
494 execution_fence: None,
495 lease_expires_epoch_ms: None,
496 path_claims: Vec::new(),
497 task_ref: String::new(),
498 parent_task_ref: None,
499 policy_ref: String::new(),
500 expected_outcome_ref: String::new(),
501 execution_started: false,
502 };
503 self.nodes.insert(child_node_id.clone(), node);
504 self.children
505 .entry(parent_node_id.to_string())
506 .or_default()
507 .push(child_node_id.clone());
508 let chain_id = self
509 .chain_id_for_node(&child_node_id)
510 .expect("delegated child always has a root node");
511 let pending_budget =
512 self.pending_child_budgets
513 .remove(&child_node_id)
514 .unwrap_or(WorkNodeBudget {
515 tokens_allocated: 0,
516 tokens_consumed: 0,
517 cost_micros_allocated: 0,
518 cost_micros_consumed: 0,
519 });
520 self.pending_child_parents.remove(&child_node_id);
521 self.record_chain_allocation(
522 &chain_id,
523 pending_budget.tokens_allocated,
524 child_tokens_allocated,
525 pending_budget.cost_micros_allocated,
526 child_cost_allocated,
527 new_depth,
528 );
529 Ok(self
530 .nodes
531 .get(&child_node_id)
532 .expect("child node inserted above"))
533 }
534
535 pub fn allocate_child_budget(
540 &mut self,
541 parent_id: &str,
542 child_id: &str,
543 fraction: f64,
544 ) -> Result<WorkNodeBudget, WorkGraphError> {
545 validate_identifier(parent_id, "parent_node_id")?;
546 validate_identifier(child_id, "child_node_id")?;
547 if self.nodes.contains_key(child_id) || self.pending_child_budgets.contains_key(child_id) {
548 return Err(WorkGraphError::DuplicateNode(child_id.to_string()));
549 }
550 let (parent_budget_tokens, parent_used_tokens, parent_cost_remaining, parent_depth) = {
551 let parent = self
552 .nodes
553 .get(parent_id)
554 .ok_or_else(|| WorkGraphError::NodeNotFound(parent_id.to_string()))?;
555 if parent.status != NodeStatus::Active {
556 return Err(WorkGraphError::ParentNotActive(parent_id.to_string()));
557 }
558 if parent.depth >= self.max_depth {
559 return Err(WorkGraphError::DepthExceeded(self.max_depth));
560 }
561 (
562 parent.budget.tokens_allocated,
563 parent.budget.tokens_consumed,
564 parent.budget.cost_remaining(),
565 parent.depth,
566 )
567 };
568
569 let (reserved_tokens, reserved_cost) =
570 self.reserved_child_budget(parent_id, Some(child_id))?;
571 let parent_remaining = parent_budget_tokens
572 .saturating_sub(parent_used_tokens)
573 .saturating_sub(reserved_tokens);
574 if parent_remaining == 0 {
575 return Err(WorkGraphError::BudgetExceedsParent {
576 child_requested: 1,
577 parent_remaining,
578 });
579 }
580
581 let allocation = BudgetAllocation {
582 parent_budget_tokens,
583 parent_used_tokens,
584 child_fraction: fraction,
585 minimum_budget: 0,
586 maximum_budget: parent_remaining,
587 };
588 let mut cascaded = cascade_budget(&allocation);
589 cascaded.depth = u32::from(parent_depth) + 1;
590 cascaded.lineage = self.node_lineage(parent_id);
591 cascaded.lineage.push(child_id.to_string());
592 validate_cascade(&cascaded)?;
593 if cascaded.allocated_tokens > parent_remaining {
594 return Err(WorkGraphError::BudgetExceedsParent {
595 child_requested: cascaded.allocated_tokens,
596 parent_remaining,
597 });
598 }
599
600 let parent_cost_remaining = parent_cost_remaining.saturating_sub(reserved_cost);
601 let cost_micros_allocated = if parent_cost_remaining == 0 {
602 0
603 } else {
604 ((parent_cost_remaining as f64 * fraction) as u64)
605 .max(1)
606 .min(parent_cost_remaining)
607 };
608 let child_budget = WorkNodeBudget {
609 tokens_allocated: cascaded.allocated_tokens,
610 tokens_consumed: 0,
611 cost_micros_allocated,
612 cost_micros_consumed: 0,
613 };
614 let chain_id = self
615 .chain_id_for_node(parent_id)
616 .expect("parent node always has a root node");
617 self.pending_child_budgets
618 .insert(child_id.to_string(), child_budget.clone());
619 self.pending_child_parents
620 .insert(child_id.to_string(), parent_id.to_string());
621 self.record_chain_allocation(
622 &chain_id,
623 0,
624 child_budget.tokens_allocated,
625 0,
626 child_budget.cost_micros_allocated,
627 parent_depth + 1,
628 );
629
630 Ok(child_budget)
631 }
632
633 pub fn complete(&mut self, node_id: &str, outcome_ref: String) -> Result<(), WorkGraphError> {
635 validate_identifier(&outcome_ref, "outcome_ref")?;
636 let node = self
637 .nodes
638 .get_mut(node_id)
639 .ok_or_else(|| WorkGraphError::NodeNotFound(node_id.to_string()))?;
640 if node.status != NodeStatus::Active {
641 return Err(WorkGraphError::InvalidTransition(node_id.to_string()));
642 }
643 node.status = NodeStatus::Completed;
644 node.outcome_ref = Some(outcome_ref);
645 node.lease_expires_epoch_ms = None;
646 self.end_incarnation(node_id);
647 Ok(())
648 }
649
650 pub fn record_child_receipt(
652 &mut self,
653 receipt: ChildExecutionReceipt,
654 ) -> Result<(), WorkGraphError> {
655 validate_receipt(&receipt)?;
656 if self.receipts.contains_key(&receipt.receipt_id)
657 || self
658 .receipts
659 .values()
660 .any(|known| known.node_id == receipt.node_id)
661 {
662 return Err(WorkGraphError::DuplicateReceipt(receipt.receipt_id));
663 }
664 let node = self
665 .nodes
666 .get(&receipt.node_id)
667 .ok_or_else(|| WorkGraphError::NodeNotFound(receipt.node_id.clone()))?;
668 if node.parent_node_id.is_none() || node.status != NodeStatus::Active {
669 return Err(WorkGraphError::InvalidTransition(receipt.node_id));
670 }
671 if node.execution_fence.is_some() && receipt.execution_fence != node.execution_fence {
672 return Err(WorkGraphError::StaleExecutionFence(receipt.node_id));
673 }
674 if receipt.tokens_consumed > node.budget.tokens_allocated
675 || receipt.cost_micros_consumed > node.budget.cost_micros_allocated
676 || receipt.tokens_consumed < node.budget.tokens_consumed
677 || receipt.cost_micros_consumed < node.budget.cost_micros_consumed
678 {
679 return Err(WorkGraphError::ReceiptExceedsBudget(receipt.node_id));
680 }
681 let token_delta = receipt.tokens_consumed - node.budget.tokens_consumed;
682 let cost_delta = receipt.cost_micros_consumed - node.budget.cost_micros_consumed;
683 let (reserved_tokens, reserved_cost) =
684 self.reserved_child_budget(&receipt.node_id, None)?;
685 if token_delta
686 > node
687 .budget
688 .tokens_remaining()
689 .saturating_sub(reserved_tokens)
690 || cost_delta > node.budget.cost_remaining().saturating_sub(reserved_cost)
691 {
692 return Err(WorkGraphError::ReceiptExceedsBudget(receipt.node_id));
693 }
694
695 let node_id = receipt.node_id.clone();
696 let receipt_id = receipt.receipt_id.clone();
697 let chain_id = self
698 .chain_id_for_node(&node_id)
699 .expect("receipt node always has a root");
700 self.ensure_chain_budget(&chain_id);
701 let chain = self
702 .chain_budgets
703 .get(&chain_id)
704 .expect("chain budget initialized");
705 let chain_consumed = chain
706 .total_consumed_tokens
707 .checked_add(token_delta)
708 .ok_or(WorkGraphError::Overflow)?;
709 let chain_cost = chain
710 .total_consumed_cost_micros
711 .checked_add(cost_delta)
712 .ok_or(WorkGraphError::Overflow)?;
713 if chain_consumed.saturating_add(chain.unauthorized_tokens) > chain.root_budget_tokens
714 || chain_cost.saturating_add(chain.unauthorized_cost_micros)
715 > chain.root_budget_cost_micros
716 {
717 return Err(WorkGraphError::ReceiptExceedsBudget(node_id));
718 }
719 let node = self
720 .nodes
721 .get_mut(&node_id)
722 .expect("receipt node checked above");
723 node.budget.tokens_consumed = receipt.tokens_consumed;
724 node.budget.cost_micros_consumed = receipt.cost_micros_consumed;
725 node.status = match receipt.outcome {
726 ChildOutcome::Accepted | ChildOutcome::Partial => NodeStatus::Completed,
727 ChildOutcome::Rejected => NodeStatus::Failed,
728 };
729 node.outcome_ref = Some(receipt.outcome_ref.clone());
730 node.lease_expires_epoch_ms = None;
731 self.end_incarnation(&node_id);
732 self.chain_budgets
733 .get_mut(&chain_id)
734 .expect("chain budget initialized")
735 .total_consumed_tokens = chain_consumed;
736 self.chain_budgets
737 .get_mut(&chain_id)
738 .expect("chain budget initialized")
739 .total_consumed_cost_micros = chain_cost;
740 self.receipts.insert(receipt_id, receipt);
741 Ok(())
742 }
743
744 #[cfg(test)]
753 pub fn record_failed_execution_spend(
754 &mut self,
755 spend: MeasuredExecutionSpend,
756 ) -> Result<(), WorkGraphError> {
757 validate_measured_spend(&spend)?;
758 if self.failed_spend.get(&spend.node_id) == Some(&spend) {
759 return Ok(());
760 }
761 if self.failed_spend.contains_key(&spend.node_id)
762 || self
763 .receipts
764 .values()
765 .any(|receipt| receipt.node_id == spend.node_id)
766 {
767 return Err(WorkGraphError::DuplicateReceipt(spend.spend_id));
768 }
769 let node = self
770 .nodes
771 .get(&spend.node_id)
772 .ok_or_else(|| WorkGraphError::NodeNotFound(spend.node_id.clone()))?;
773 if node.parent_node_id.is_none() || node.status != NodeStatus::Active {
774 return Err(WorkGraphError::InvalidTransition(spend.node_id));
775 }
776 if node.execution_fence.as_deref() != Some(spend.execution_fence.as_str()) {
777 return Err(WorkGraphError::StaleExecutionFence(spend.node_id));
778 }
779 if spend.tokens_consumed <= node.budget.tokens_allocated
780 && spend.cost_micros_consumed <= node.budget.cost_micros_allocated
781 {
782 return Err(WorkGraphError::InvalidBudget);
783 }
784
785 let node_id = spend.node_id.clone();
786 let chain_id = self
787 .chain_id_for_node(&node_id)
788 .expect("failed spend node always has a root");
789 let unauthorized_tokens = spend
790 .tokens_consumed
791 .saturating_sub(node.budget.tokens_allocated);
792 let unauthorized_cost = spend
793 .cost_micros_consumed
794 .saturating_sub(node.budget.cost_micros_allocated);
795 self.stop(&node_id, StopReason::BudgetExhausted)?;
796 let chain = self
797 .chain_budgets
798 .get_mut(&chain_id)
799 .expect("failed spend chain budget initialized");
800 chain.unauthorized_tokens = chain
801 .unauthorized_tokens
802 .saturating_add(unauthorized_tokens);
803 chain.unauthorized_cost_micros = chain
804 .unauthorized_cost_micros
805 .saturating_add(unauthorized_cost);
806 self.failed_spend.insert(node_id, spend);
807 Ok(())
808 }
809
810 pub fn evaluate_child_outcome(&self, node_id: &str) -> Result<ChildOutcome, WorkGraphError> {
812 validate_identifier(node_id, "node_id")?;
813 self.receipts
814 .values()
815 .find(|receipt| receipt.node_id == node_id)
816 .map(|receipt| receipt.outcome)
817 .ok_or_else(|| WorkGraphError::MissingReceipt(node_id.to_string()))
818 }
819
820 pub fn mark_accepted_path(&mut self, node_ids: &[String]) -> Result<(), WorkGraphError> {
824 let mut accepted = BTreeSet::new();
825 for node_id in node_ids {
826 validate_identifier(node_id, "node_id")?;
827 let receipt = self
828 .receipts
829 .values()
830 .find(|receipt| receipt.node_id == *node_id)
831 .ok_or_else(|| WorkGraphError::MissingReceipt(node_id.clone()))?;
832 if receipt.outcome != ChildOutcome::Accepted {
833 return Err(WorkGraphError::RejectedAcceptedPath(node_id.clone()));
834 }
835 accepted.extend(self.node_lineage(node_id));
836 }
837 self.accepted_path = accepted;
838 Ok(())
839 }
840
841 pub fn is_on_accepted_path(&self, node_id: &str) -> bool {
842 self.accepted_path.contains(node_id)
843 }
844
845 pub fn fuse_accepted_results(&self) -> Result<FusedResults, WorkGraphError> {
848 let mut candidates: BTreeMap<String, BTreeMap<String, Vec<String>>> = BTreeMap::new();
849 for receipt in self.receipts.values() {
850 if !self.accepted_path.contains(&receipt.node_id) {
851 continue;
852 }
853 for claim in &receipt.claims {
854 candidates
855 .entry(claim.key.clone())
856 .or_default()
857 .entry(claim.value.clone())
858 .or_default()
859 .push(receipt.node_id.clone());
860 }
861 }
862 let mut accepted = BTreeMap::new();
863 let mut conflicts = Vec::new();
864 for (key, values) in candidates {
865 if values.len() == 1 {
866 let (value, _) = values.into_iter().next().expect("one value");
867 accepted.insert(key, value);
868 } else {
869 let mut node_ids: Vec<String> = values.values().flatten().cloned().collect();
870 node_ids.sort();
871 node_ids.dedup();
872 conflicts.push(ResultConflict {
873 key,
874 values: values.into_keys().collect(),
875 node_ids,
876 });
877 }
878 }
879 Ok(FusedResults {
880 accepted,
881 conflicts,
882 })
883 }
884
885 pub fn attribution(&self) -> Result<WorkAttribution, WorkGraphError> {
886 let mut result = WorkAttribution {
887 accepted_tokens: 0,
888 accepted_cost_micros: 0,
889 waste_tokens: 0,
890 waste_cost_micros: 0,
891 };
892 for node in self.nodes.values() {
893 let (tokens, cost) = self.failed_spend.get(&node.node_id).map_or(
894 (
895 node.budget.tokens_consumed,
896 node.budget.cost_micros_consumed,
897 ),
898 |spend| (spend.tokens_consumed, spend.cost_micros_consumed),
899 );
900 let (token_total, cost_total) = if self.accepted_path.contains(&node.node_id) {
901 (
902 &mut result.accepted_tokens,
903 &mut result.accepted_cost_micros,
904 )
905 } else {
906 (&mut result.waste_tokens, &mut result.waste_cost_micros)
907 };
908 *token_total = token_total
909 .checked_add(tokens)
910 .ok_or(WorkGraphError::Overflow)?;
911 *cost_total = cost_total
912 .checked_add(cost)
913 .ok_or(WorkGraphError::Overflow)?;
914 }
915 Ok(result)
916 }
917
918 pub fn team_value_report(&self) -> Result<TeamValueReport, WorkGraphError> {
919 let attribution = self.attribution()?;
920 let total_tokens = attribution
921 .accepted_tokens
922 .checked_add(attribution.waste_tokens)
923 .ok_or(WorkGraphError::Overflow)?;
924 let total_cost_micros = attribution
925 .accepted_cost_micros
926 .checked_add(attribution.waste_cost_micros)
927 .ok_or(WorkGraphError::Overflow)?;
928 let accepted_cost_basis_points = if total_cost_micros == 0 {
929 0
930 } else {
931 u16::try_from(
932 u128::from(attribution.accepted_cost_micros).saturating_mul(10_000)
933 / u128::from(total_cost_micros),
934 )
935 .unwrap_or(10_000)
936 };
937 Ok(TeamValueReport {
938 accepted_tokens: attribution.accepted_tokens,
939 total_tokens,
940 accepted_cost_micros: attribution.accepted_cost_micros,
941 total_cost_micros,
942 waste_tokens: attribution.waste_tokens,
943 waste_cost_micros: attribution.waste_cost_micros,
944 accepted_cost_basis_points,
945 })
946 }
947}