1use std::collections::{BTreeMap, BTreeSet};
10
11use compact_str::CompactString;
12use serde::{Deserialize, Serialize};
13
14use crate::proc::ProcessState;
15use crate::scheduler::policy::SchedulerBudget;
16use crate::types::agent::{AgentIsolation, AgentRole, ContextInheritance, IsolationManifest};
17use crate::types::result::{SubAgentResult, TerminationReason};
18
19pub type TaskId = CompactString;
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
27#[serde(rename_all = "snake_case")]
28pub enum TaskLifecycle {
29 PendingLaunch,
37 Starting,
41 Ready,
43 Running,
45 Suspended,
47 Done(TerminationReason),
49}
50
51impl TaskLifecycle {
52 pub fn label(self) -> &'static str {
53 match self {
54 Self::PendingLaunch => "pending_launch",
55 Self::Starting => "starting",
56 Self::Ready => "ready",
57 Self::Running => "running",
58 Self::Suspended => "suspended",
59 Self::Done(_) => "done",
60 }
61 }
62
63 pub fn is_terminal(self) -> bool {
64 matches!(self, Self::Done(_))
65 }
66
67 pub fn occupies_slot(self) -> bool {
73 matches!(self, Self::PendingLaunch | Self::Starting | Self::Running)
74 }
75}
76
77#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
78#[serde(rename_all = "snake_case")]
79pub enum RunnableCause {
80 #[default]
81 NestedTask,
82 TimerWaiter,
83 MessageWaiter,
84 EventWaiter,
85}
86
87impl From<ProcessState> for TaskLifecycle {
89 fn from(state: ProcessState) -> Self {
90 match state {
91 ProcessState::Running => TaskLifecycle::Running,
92 ProcessState::Joined => TaskLifecycle::Done(TerminationReason::Completed),
93 ProcessState::Failed => TaskLifecycle::Done(TerminationReason::Error),
96 }
97 }
98}
99
100#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
104pub struct ApprovalId(pub CompactString);
105
106#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
107pub struct SignalFilter(pub CompactString);
108
109#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
112pub struct LogicalDeadline(pub u64);
113
114#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
115pub struct ChannelId(pub CompactString);
116
117#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
118pub struct ResourceKey(pub CompactString);
119
120#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
121pub struct SubscriptionId(pub CompactString);
122
123#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
125pub enum WaitCondition {
126 Effect(crate::runtime::kernel::wire::EffectId),
127 Child(TaskId),
128 Children(Vec<TaskId>),
129 Approval(ApprovalId),
130 Signal(SignalFilter),
131 Timer(LogicalDeadline),
132 Channel(ChannelId),
133 Resource(ResourceKey),
134 External(SubscriptionId),
135}
136
137#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
138pub enum WaitMode {
139 Any,
140 All,
141}
142
143#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
145pub struct WaitSet {
146 pub mode: WaitMode,
147 pub conditions: Vec<WaitCondition>,
148}
149
150#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
153pub struct DurableWaitSet {
154 pub mode: WaitMode,
155 pub conditions: Vec<WaitCondition>,
156 pub satisfied: BTreeSet<usize>,
157}
158
159impl From<WaitSet> for DurableWaitSet {
160 fn from(wait_set: WaitSet) -> Self {
161 let conditions = wait_set
162 .conditions
163 .into_iter()
164 .flat_map(|condition| match condition {
165 WaitCondition::Children(ids) => ids
166 .into_iter()
167 .map(WaitCondition::Child)
168 .collect::<Vec<_>>(),
169 condition => vec![condition],
170 })
171 .collect();
172 Self {
173 mode: wait_set.mode,
174 conditions,
175 satisfied: BTreeSet::new(),
176 }
177 }
178}
179
180#[derive(Debug, Clone)]
183pub struct BudgetLedger {
184 pub limits: SchedulerBudget,
185 pub turns: u32,
186 pub total_tokens: u64,
187 pub started_at_ms: Option<u64>,
188}
189
190impl BudgetLedger {
191 pub fn new(limits: SchedulerBudget) -> Self {
192 Self {
193 limits,
194 turns: 0,
195 total_tokens: 0,
196 started_at_ms: None,
197 }
198 }
199
200 pub fn exceeded(&self, now_ms: Option<u64>) -> Option<&'static str> {
202 self.limits
203 .should_terminate(self.turns, self.total_tokens, now_ms, self.started_at_ms)
204 }
205}
206
207impl Default for BudgetLedger {
208 fn default() -> Self {
209 Self::new(SchedulerBudget::default())
210 }
211}
212
213#[derive(Debug, Clone)]
222pub struct ProcInfo {
223 pub role: AgentRole,
224 pub isolation: AgentIsolation,
225 pub context_inheritance: ContextInheritance,
226 pub result: Option<SubAgentResult>,
228}
229
230#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
234#[serde(rename_all = "snake_case")]
235pub enum ChildFailurePolicy {
236 #[default]
237 Propagate,
238 Isolate,
239 Restart,
240 Retry,
241 Ignore,
242}
243
244#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
246pub struct SupervisionPolicy {
247 pub child_failure: ChildFailurePolicy,
248 pub max_restarts: Option<u32>,
249 pub cancel_children_on_exit: bool,
250}
251
252#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
253pub struct SupervisionEvent {
254 pub attempt: u32,
255 pub strategy: ChildFailurePolicy,
256 pub reason: CompactString,
257 pub terminal: bool,
259 pub relaunched: bool,
260}
261
262impl Default for SupervisionPolicy {
263 fn default() -> Self {
264 Self {
265 child_failure: ChildFailurePolicy::default(),
266 max_restarts: None,
267 cancel_children_on_exit: true,
268 }
269 }
270}
271
272#[derive(Debug, Clone)]
274pub struct Tcb {
275 pub id: TaskId,
276 pub parent: Option<TaskId>,
277 pub children: BTreeSet<TaskId>,
280 pub state: TaskLifecycle,
281 pub runnable_cause: RunnableCause,
283 pub budget: BudgetLedger,
284 pub wait_set: Option<DurableWaitSet>,
286 pub caps: Vec<CompactString>,
288 pub capabilities: Vec<crate::types::capability::Capability>,
293 pub proc: Option<ProcInfo>,
295 pub supervision: SupervisionPolicy,
298 pub supervision_events: Vec<SupervisionEvent>,
300 pub detached: bool,
304 pub child_budget_remaining: Option<super::budget_grant::ResourceBudget>,
309 pub budget_grant: Option<super::budget_grant::BudgetGrant>,
313 pub mailbox: super::mailbox::Mailbox,
316}
317
318impl Tcb {
319 pub fn root(id: impl Into<TaskId>, budget: SchedulerBudget) -> Self {
321 Self {
322 id: id.into(),
323 parent: None,
324 children: BTreeSet::new(),
325 state: TaskLifecycle::Ready,
326 runnable_cause: RunnableCause::NestedTask,
327 budget: BudgetLedger::new(budget),
328 wait_set: None,
329 caps: Vec::new(),
330 capabilities: Vec::new(),
331 proc: None,
332 supervision: SupervisionPolicy::default(),
333 supervision_events: Vec::new(),
334 detached: false,
335 child_budget_remaining: None,
336 budget_grant: None,
337 mailbox: super::mailbox::Mailbox::new(),
338 }
339 }
340
341 pub fn spawned(manifest: &IsolationManifest, budget: SchedulerBudget) -> Self {
344 Self::spawned_in(
345 manifest,
346 budget,
347 TaskLifecycle::Running,
348 Some(TaskId::from("root")),
349 )
350 }
351
352 pub fn spawned_in(
362 manifest: &IsolationManifest,
363 budget: SchedulerBudget,
364 state: TaskLifecycle,
365 parent: Option<TaskId>,
366 ) -> Self {
367 Self {
368 id: manifest.agent_id.clone(),
369 parent,
370 children: BTreeSet::new(),
371 state,
372 runnable_cause: RunnableCause::NestedTask,
373 budget: BudgetLedger::new(budget),
374 wait_set: None,
375 caps: manifest.permitted_capability_ids.clone(),
376 capabilities: manifest.requested_capabilities.clone(),
377 proc: Some(ProcInfo {
378 role: manifest.role,
379 isolation: manifest.isolation,
380 context_inheritance: manifest.context_inheritance,
381 result: None,
382 }),
383 supervision: SupervisionPolicy::default(),
384 supervision_events: Vec::new(),
385 detached: false,
386 child_budget_remaining: None,
387 budget_grant: None,
388 mailbox: super::mailbox::Mailbox::new(),
389 }
390 }
391}
392
393#[derive(Debug, Clone, Default)]
396pub struct TaskTable {
397 tasks: Vec<Tcb>,
398 wait_index: super::wait_index::WaitIndex,
400 channels: BTreeMap<ChannelId, super::mailbox::Channel>,
401 objects: BTreeMap<crate::mm::handle::ObjectId, crate::mm::handle::ObjectDescriptor>,
402}
403
404#[derive(Debug, Clone, Copy, PartialEq, Eq)]
406pub(crate) enum TaskSpawnError {
407 UnknownCaller,
408 CallerTerminal,
409 DuplicateTask,
410}
411
412#[derive(Debug, Clone, Copy, PartialEq, Eq)]
413pub(crate) enum LocalIpcError {
414 UnknownCaller,
415 CallerTerminal,
416 UnknownRecipient,
417 ChannelSubscribersMismatch,
418 NotSubscriber,
419 Full,
420 Expired,
421 ObjectConflict,
422}
423
424impl TaskTable {
425 pub fn new() -> Self {
426 Self::default()
427 }
428
429 pub fn insert(&mut self, tcb: Tcb) {
430 let is_new = !self.tasks.iter().any(|t| t.id == tcb.id);
433 if is_new
434 && let Some(parent_id) = tcb.parent.clone()
435 && let Some(parent) = self.tasks.iter_mut().find(|t| t.id == parent_id)
436 {
437 parent.children.insert(tcb.id.clone());
438 }
439 if let Some(existing) = self.tasks.iter_mut().find(|t| t.id == tcb.id) {
440 *existing = tcb;
441 } else {
442 self.tasks.push(tcb);
443 }
444 }
445
446 pub(crate) fn spawn_child(
452 &mut self,
453 caller: &str,
454 manifest: &IsolationManifest,
455 budget: SchedulerBudget,
456 state: TaskLifecycle,
457 ) -> Result<TaskId, TaskSpawnError> {
458 let parent = self.get(caller).ok_or(TaskSpawnError::UnknownCaller)?;
459 if parent.state.is_terminal() {
460 return Err(TaskSpawnError::CallerTerminal);
461 }
462 if self.get(manifest.agent_id.as_str()).is_some() {
463 return Err(TaskSpawnError::DuplicateTask);
464 }
465
466 let child_id = manifest.agent_id.clone();
467 let child = Tcb::spawned_in(manifest, budget, state, Some(caller.into()));
468 self.insert(child);
469 Ok(child_id)
470 }
471
472 pub fn get(&self, id: &str) -> Option<&Tcb> {
473 self.tasks.iter().find(|t| t.id.as_str() == id)
474 }
475
476 pub fn get_mut(&mut self, id: &str) -> Option<&mut Tcb> {
477 self.tasks.iter_mut().find(|t| t.id.as_str() == id)
478 }
479
480 pub fn all(&self) -> &[Tcb] {
481 &self.tasks
482 }
483
484 pub(crate) fn runnable_candidates(&self) -> Vec<super::runnable::LocalRunnable> {
485 let mut tasks: Vec<_> = self
486 .tasks
487 .iter()
488 .filter(|task| task.proc.is_some() && task.state == TaskLifecycle::Ready)
489 .collect();
490 tasks.sort_by(|left, right| left.id.cmp(&right.id));
491 tasks
492 .into_iter()
493 .map(|task| super::runnable::LocalRunnable {
494 id: task.id.to_string(),
495 kind: match task.runnable_cause {
496 RunnableCause::NestedTask => super::runnable::LocalRunnableKind::NestedTask,
497 RunnableCause::TimerWaiter => super::runnable::LocalRunnableKind::TimerWaiter,
498 RunnableCause::MessageWaiter => {
499 super::runnable::LocalRunnableKind::MessageWaiter
500 }
501 RunnableCause::EventWaiter => super::runnable::LocalRunnableKind::EventWaiter,
502 },
503 source_rank: 0,
504 })
505 .collect()
506 }
507
508 pub fn children_of(&self, parent: &str) -> Vec<&Tcb> {
509 self.tasks
510 .iter()
511 .filter(|t| t.parent.as_deref() == Some(parent))
512 .collect()
513 }
514
515 pub(crate) fn lineage_depth(&self, task_id: &str) -> Option<usize> {
518 let mut current = self.get(task_id)?;
519 let mut depth = 0usize;
520 while let Some(parent) = current.parent.as_ref() {
521 depth = depth.checked_add(1)?;
522 if depth > self.tasks.len() {
523 return None;
524 }
525 current = self.get(parent.as_str())?;
526 }
527 Some(depth)
528 }
529
530 pub fn wait_index(&self) -> &super::wait_index::WaitIndex {
531 &self.wait_index
532 }
533
534 pub fn wait_for_timer(&mut self, task_id: &str, deadline: LogicalDeadline) {
536 self.register_wait_set(
537 task_id,
538 WaitSet {
539 mode: WaitMode::Any,
540 conditions: vec![WaitCondition::Timer(deadline)],
541 },
542 );
543 }
544
545 pub fn wake_expired_timers(&mut self, now_ms: u64) -> Vec<TaskId> {
547 self.wait_index
548 .due_timer_keys(now_ms)
549 .into_iter()
550 .flat_map(|key| self.notify(&key))
551 .collect()
552 }
553
554 pub fn wait_for_condition(&mut self, task_id: &str, condition: &WaitCondition) {
556 self.register_wait_set(
557 task_id,
558 WaitSet {
559 mode: WaitMode::Any,
560 conditions: vec![condition.clone()],
561 },
562 );
563 }
564
565 pub fn wake(&mut self, key: &super::wait_index::WaitKey) -> Vec<TaskId> {
568 self.notify(key)
569 }
570
571 pub fn register_wait_set(&mut self, task_id: &str, wait_set: WaitSet) {
574 let Some(id) = self.get(task_id).map(|t| t.id.clone()) else {
575 return;
576 };
577 if wait_set.conditions.is_empty()
578 || self
579 .get(task_id)
580 .is_some_and(|task| task.state.is_terminal())
581 {
582 return;
583 }
584 self.clear_durable_wait(&id);
585 self.wait_index
586 .register_wait_set(id.clone(), wait_set.clone());
587 if let Some(task) = self.get_mut(task_id) {
588 task.state = TaskLifecycle::Suspended;
589 task.wait_set = Some(wait_set.into());
590 }
591 }
592
593 pub fn clear_wait(&mut self, task_id: &str) {
595 let Some(id) = self.get(task_id).map(|task| task.id.clone()) else {
596 return;
597 };
598 self.clear_durable_wait(&id);
599 }
600
601 pub fn notify(&mut self, key: &super::wait_index::WaitKey) -> Vec<TaskId> {
605 let mut candidates = self.wait_index.lookup(key).to_vec();
606 candidates.sort_unstable();
610 let mut woken = Vec::new();
611 for task_id in candidates {
612 let terminal = self
613 .get(task_id.as_str())
614 .is_none_or(|task| task.state.is_terminal());
615 if terminal {
616 self.clear_durable_wait(&task_id);
617 continue;
618 }
619
620 let satisfied = if let Some(task) = self.get_mut(task_id.as_str())
621 && let Some(wait_set) = task.wait_set.as_mut()
622 {
623 for (index, condition) in wait_set.conditions.iter().enumerate() {
624 if key.matches(condition) {
625 wait_set.satisfied.insert(index);
626 }
627 }
628 match wait_set.mode {
629 WaitMode::Any => !wait_set.satisfied.is_empty(),
630 WaitMode::All => wait_set.satisfied.len() == wait_set.conditions.len(),
631 }
632 } else {
633 false
634 };
635
636 if satisfied {
637 let cause = match key {
638 super::wait_index::WaitKey::Timer(_) => RunnableCause::TimerWaiter,
639 super::wait_index::WaitKey::Channel(_)
640 | super::wait_index::WaitKey::External(_) => RunnableCause::MessageWaiter,
641 _ => RunnableCause::EventWaiter,
642 };
643 self.clear_durable_wait(&task_id);
644 if let Some(task) = self.get_mut(task_id.as_str())
645 && task.state == TaskLifecycle::Suspended
646 {
647 task.state = TaskLifecycle::Ready;
648 task.runnable_cause = cause;
649 woken.push(task_id);
650 }
651 }
652 }
653 woken
654 }
655
656 fn clear_durable_wait(&mut self, task_id: &TaskId) {
657 let wait_set = self
658 .get_mut(task_id.as_str())
659 .and_then(|task| task.wait_set.take());
660 if let Some(wait_set) = wait_set {
661 for condition in &wait_set.conditions {
662 self.wait_index.remove(task_id, condition);
663 }
664 }
665 }
666
667 pub fn root_id(&self) -> Option<TaskId> {
671 self.tasks
672 .iter()
673 .find(|t| t.parent.is_none())
674 .map(|t| t.id.clone())
675 }
676
677 pub fn cancel_subtree(&mut self, task_id: &str) {
680 let children: Vec<TaskId> = self
681 .get(task_id)
682 .map(|t| t.children.iter().cloned().collect())
683 .unwrap_or_default();
684 if let Some(task) = self.get_mut(task_id) {
685 task.state = TaskLifecycle::Done(TerminationReason::UserAbort);
686 }
687 for child_id in children {
688 let is_detached = self
692 .get(child_id.as_str())
693 .is_some_and(|child| child.detached);
694 if !is_detached {
695 self.cancel_subtree(child_id.as_str());
696 }
697 }
698 }
699
700 pub fn cancel_children(&mut self, task_id: &str) {
708 let children: Vec<TaskId> = self
709 .get(task_id)
710 .map(|t| t.children.iter().cloned().collect())
711 .unwrap_or_default();
712 for child_id in children {
713 if self
714 .get(child_id.as_str())
715 .is_some_and(|t| !t.state.is_terminal() && !t.detached)
716 {
717 self.cancel_subtree(child_id.as_str());
718 }
719 }
720 }
721
722 pub(crate) fn prepare_supervised_relaunch(
723 &mut self,
724 task_id: &str,
725 strategy: ChildFailurePolicy,
726 ) {
727 self.clear_wait(task_id);
728 if let Some(task) = self.get_mut(task_id) {
729 task.state = TaskLifecycle::PendingLaunch;
730 if let Some(proc) = task.proc.as_mut() {
731 proc.result = None;
732 }
733 if strategy == ChildFailurePolicy::Restart {
734 task.budget.turns = 0;
735 task.budget.total_tokens = 0;
736 }
737 }
738 }
739
740 pub fn return_child_budget(&mut self, child_id: &str) {
748 let Some(grant) = self.get(child_id).and_then(|t| t.budget_grant.clone()) else {
749 return;
750 };
751 if grant.settled {
752 return;
753 }
754 let unused = super::budget_grant::return_unused(&grant);
755 if let Some(child) = self.get_mut(child_id)
756 && let Some(child_grant) = child.budget_grant.as_mut()
757 {
758 child_grant.returned = unused;
759 child_grant.settled = true;
760 }
761 if let Some(parent) = self.get_mut(grant.parent.as_str()) {
762 if let Some(remaining) = parent.child_budget_remaining {
763 parent.child_budget_remaining =
764 Some(super::budget_grant::credit(&remaining, &unused));
765 }
766 if let Some(parent_grant) = parent.budget_grant.as_mut()
769 && !parent_grant.settled
770 {
771 parent_grant.consumed =
772 super::budget_grant::accumulate_usage(&parent_grant.consumed, &grant.consumed);
773 }
774 }
775 }
776
777 pub(crate) fn attach_child_budget_grant(
780 &mut self,
781 child_id: &str,
782 grant: super::budget_grant::BudgetGrant,
783 ) {
784 if grant.child.as_str() != child_id {
785 return;
786 }
787 if let Some(child) = self.get_mut(child_id) {
788 child.child_budget_remaining = Some(grant.reserved);
789 child.budget_grant = Some(grant);
790 }
791 }
792
793 pub fn send_message(&mut self, msg: super::mailbox::MailboxMessage) {
797 if let Some(recipient) = self.get_mut(msg.to.as_str()) {
798 recipient.mailbox.send(msg);
799 }
800 }
801
802 pub(crate) fn send_message_from(
805 &mut self,
806 caller: &str,
807 mut msg: super::mailbox::MailboxMessage,
808 now: super::mailbox::LogicalTime,
809 ) -> Result<bool, LocalIpcError> {
810 let caller_task = self.get(caller).ok_or(LocalIpcError::UnknownCaller)?;
811 if caller_task.state.is_terminal() {
812 return Err(LocalIpcError::CallerTerminal);
813 }
814 if self.get(msg.to.as_str()).is_none() {
815 return Err(LocalIpcError::UnknownRecipient);
816 }
817 msg.from = caller.into();
818 let recipient_id = msg.to.clone();
819 let outcome = self
820 .get_mut(recipient_id.as_str())
821 .expect("recipient was validated")
822 .mailbox
823 .try_send(msg, now);
824 match outcome {
825 super::mailbox::IpcEnqueueOutcome::Accepted => {
826 self.notify(&super::wait_index::WaitKey::External(SubscriptionId(
827 format!("mailbox:{recipient_id}").into(),
828 )));
829 Ok(true)
830 }
831 super::mailbox::IpcEnqueueOutcome::Duplicate => Ok(false),
832 super::mailbox::IpcEnqueueOutcome::Full => Err(LocalIpcError::Full),
833 super::mailbox::IpcEnqueueOutcome::Expired => Err(LocalIpcError::Expired),
834 }
835 }
836
837 pub(crate) fn receive_mailbox(
838 &mut self,
839 caller: &str,
840 now: super::mailbox::LogicalTime,
841 max: usize,
842 ) -> Result<Vec<super::mailbox::MailboxMessage>, LocalIpcError> {
843 let task = self.get_mut(caller).ok_or(LocalIpcError::UnknownCaller)?;
844 let mut messages = Vec::new();
845 for _ in 0..max {
846 let Some(message) = task.mailbox.receive_at(now) else {
847 break;
848 };
849 messages.push(message);
850 }
851 Ok(messages)
852 }
853
854 pub(crate) fn publish_channel(
855 &mut self,
856 caller: &str,
857 channel_id: ChannelId,
858 subscribers: Vec<TaskId>,
859 mut msg: super::mailbox::MailboxMessage,
860 now: super::mailbox::LogicalTime,
861 ) -> Result<bool, LocalIpcError> {
862 let caller_task = self.get(caller).ok_or(LocalIpcError::UnknownCaller)?;
863 if caller_task.state.is_terminal() {
864 return Err(LocalIpcError::CallerTerminal);
865 }
866 if subscribers.iter().any(|id| self.get(id.as_str()).is_none()) {
867 return Err(LocalIpcError::UnknownRecipient);
868 }
869 msg.from = caller.into();
870 let channel = self
871 .channels
872 .entry(channel_id.clone())
873 .or_insert_with(|| super::mailbox::Channel::new(subscribers.clone()));
874 if channel.subscribers != subscribers {
875 return Err(LocalIpcError::ChannelSubscribersMismatch);
876 }
877 match channel.publish_at(msg, now) {
878 super::mailbox::IpcEnqueueOutcome::Accepted => {
879 self.notify(&super::wait_index::WaitKey::Channel(channel_id));
880 Ok(true)
881 }
882 super::mailbox::IpcEnqueueOutcome::Duplicate => Ok(false),
883 super::mailbox::IpcEnqueueOutcome::Full => Err(LocalIpcError::Full),
884 super::mailbox::IpcEnqueueOutcome::Expired => Err(LocalIpcError::Expired),
885 }
886 }
887
888 pub(crate) fn receive_channel(
889 &mut self,
890 caller: &str,
891 channel_id: &ChannelId,
892 now: super::mailbox::LogicalTime,
893 ) -> Result<Vec<super::mailbox::MailboxMessage>, LocalIpcError> {
894 if self.get(caller).is_none() {
895 return Err(LocalIpcError::UnknownCaller);
896 }
897 let channel = self
898 .channels
899 .get_mut(channel_id)
900 .ok_or(LocalIpcError::NotSubscriber)?;
901 if !channel.subscribers.iter().any(|id| id.as_str() == caller) {
902 return Err(LocalIpcError::NotSubscriber);
903 }
904 Ok(channel.drain_for_at(caller.into(), now))
905 }
906
907 pub(crate) fn channels(&self) -> &BTreeMap<ChannelId, super::mailbox::Channel> {
908 &self.channels
909 }
910
911 pub(crate) fn restore_channels(
912 &mut self,
913 channels: BTreeMap<ChannelId, super::mailbox::Channel>,
914 ) {
915 self.channels = channels;
916 }
917
918 pub(crate) fn register_object(
919 &mut self,
920 caller: &str,
921 descriptor: crate::mm::handle::ObjectDescriptor,
922 ) -> Result<bool, LocalIpcError> {
923 let caller_task = self.get(caller).ok_or(LocalIpcError::UnknownCaller)?;
924 if caller_task.state.is_terminal() {
925 return Err(LocalIpcError::CallerTerminal);
926 }
927 if descriptor.owner.as_str() != caller {
928 return Err(LocalIpcError::UnknownCaller);
929 }
930 if let Some(existing) = self.objects.get(&descriptor.id) {
931 return if existing == &descriptor {
932 Ok(false)
933 } else {
934 Err(LocalIpcError::ObjectConflict)
935 };
936 }
937 let resource = ResourceKey(format!("object:{}/{}", descriptor.owner, descriptor.id).into());
938 self.objects.insert(descriptor.id, descriptor);
939 self.notify(&super::wait_index::WaitKey::Resource(resource));
940 Ok(true)
941 }
942
943 pub(crate) fn object(
944 &self,
945 object_id: crate::mm::handle::ObjectId,
946 ) -> Option<&crate::mm::handle::ObjectDescriptor> {
947 self.objects.get(&object_id)
948 }
949
950 pub(crate) fn objects(
951 &self,
952 ) -> &BTreeMap<crate::mm::handle::ObjectId, crate::mm::handle::ObjectDescriptor> {
953 &self.objects
954 }
955
956 pub(crate) fn restore_objects(
957 &mut self,
958 objects: BTreeMap<crate::mm::handle::ObjectId, crate::mm::handle::ObjectDescriptor>,
959 ) {
960 self.objects = objects;
961 }
962
963 pub fn rebuild_children(&mut self) {
969 for task in &mut self.tasks {
970 task.children.clear();
971 }
972 let edges: Vec<(TaskId, TaskId)> = self
973 .tasks
974 .iter()
975 .filter_map(|t| t.parent.clone().map(|parent| (parent, t.id.clone())))
976 .collect();
977 for (parent_id, child_id) in edges {
978 if let Some(parent) = self.tasks.iter_mut().find(|t| t.id == parent_id) {
979 parent.children.insert(child_id);
980 }
981 }
982 }
983
984 pub fn rebuild_wait_index(&mut self) {
986 self.wait_index = super::wait_index::WaitIndex::new();
987 let durable: Vec<(TaskId, WaitSet)> = self
988 .tasks
989 .iter()
990 .filter_map(|task| {
991 task.wait_set.as_ref().map(|wait_set| {
992 (
993 task.id.clone(),
994 WaitSet {
995 mode: wait_set.mode,
996 conditions: wait_set
997 .conditions
998 .iter()
999 .enumerate()
1000 .filter(|(index, _)| !wait_set.satisfied.contains(index))
1001 .map(|(_, condition)| condition.clone())
1002 .collect(),
1003 },
1004 )
1005 })
1006 })
1007 .collect();
1008 for (id, wait_set) in durable {
1009 self.wait_index.register_wait_set(id, wait_set);
1010 }
1011 }
1012
1013 pub fn has_cycle(&self) -> bool {
1017 for task in &self.tasks {
1018 let mut visited: BTreeSet<TaskId> = BTreeSet::new();
1019 let mut current = Some(task.id.clone());
1020 while let Some(id) = current {
1021 if !visited.insert(id.clone()) {
1022 return true;
1023 }
1024 current = self.get(id.as_str()).and_then(|t| t.parent.clone());
1025 }
1026 }
1027 false
1028 }
1029}
1030
1031pub fn budget_verdict(task: &Tcb, now_ms: Option<u64>) -> Option<TerminationReason> {
1035 task.budget.exceeded(now_ms).map(|axis| match axis {
1036 "max_turns" => TerminationReason::MaxTurns,
1037 "wall_time" => TerminationReason::Timeout,
1038 _ => TerminationReason::TokenBudget,
1039 })
1040}
1041
1042#[cfg(test)]
1043mod tests {
1044 use super::*;
1045 use crate::types::agent::{AgentIdentity, AgentRole, AgentRunSpec};
1046 use crate::types::capability::CapabilityManifest;
1047
1048 fn manifest_for(id: &str) -> IsolationManifest {
1049 let spec = AgentRunSpec::new(
1050 AgentIdentity::sub_agent(id, format!("{id}-session")),
1051 AgentRole::Implement,
1052 "do work",
1053 );
1054 IsolationManifest::from_spec(&spec, &CapabilityManifest::new())
1055 }
1056
1057 #[test]
1058 fn spawned_in_uses_explicit_parent() {
1059 let manifest = manifest_for("agent-7-child");
1060 let tcb = Tcb::spawned_in(
1061 &manifest,
1062 SchedulerBudget::default(),
1063 TaskLifecycle::Running,
1064 Some(TaskId::from("agent-7")),
1065 );
1066 assert_eq!(tcb.parent, Some(TaskId::from("agent-7")));
1067 }
1068
1069 #[test]
1070 fn spawn_child_derives_recursive_lineage_and_registers_each_edge_once() {
1071 let mut table = TaskTable::new();
1072 table.insert(Tcb::root("root", SchedulerBudget::default()));
1073
1074 let child_manifest = manifest_for("child");
1075 let child_id = table
1076 .spawn_child(
1077 "root",
1078 &child_manifest,
1079 SchedulerBudget::default(),
1080 TaskLifecycle::Running,
1081 )
1082 .expect("root can spawn a child");
1083
1084 let grandchild_manifest = manifest_for("grandchild");
1085 let grandchild_id = table
1086 .spawn_child(
1087 child_id.as_str(),
1088 &grandchild_manifest,
1089 SchedulerBudget::default(),
1090 TaskLifecycle::Running,
1091 )
1092 .expect("a live child can spawn its own child");
1093
1094 assert_eq!(
1095 table
1096 .get(child_id.as_str())
1097 .and_then(|task| task.parent.clone()),
1098 Some(TaskId::from("root"))
1099 );
1100 assert_eq!(
1101 table
1102 .get(grandchild_id.as_str())
1103 .and_then(|task| task.parent.clone()),
1104 Some(child_id.clone())
1105 );
1106 assert_eq!(
1107 table
1108 .children_of("root")
1109 .iter()
1110 .map(|task| task.id.clone())
1111 .collect::<Vec<_>>(),
1112 vec![child_id.clone()]
1113 );
1114 assert_eq!(
1115 table
1116 .children_of(child_id.as_str())
1117 .iter()
1118 .map(|task| task.id.clone())
1119 .collect::<Vec<_>>(),
1120 vec![grandchild_id]
1121 );
1122 assert!(!table.has_cycle());
1123 }
1124
1125 #[test]
1126 fn spawn_child_rejects_unknown_terminal_and_duplicate_callers() {
1127 let mut table = TaskTable::new();
1128 table.insert(Tcb::root("root", SchedulerBudget::default()));
1129 let manifest = manifest_for("child");
1130
1131 assert_eq!(
1132 table.spawn_child(
1133 "missing",
1134 &manifest,
1135 SchedulerBudget::default(),
1136 TaskLifecycle::Running,
1137 ),
1138 Err(TaskSpawnError::UnknownCaller)
1139 );
1140
1141 table
1142 .spawn_child(
1143 "root",
1144 &manifest,
1145 SchedulerBudget::default(),
1146 TaskLifecycle::Running,
1147 )
1148 .expect("first child creation succeeds");
1149 assert_eq!(
1150 table.spawn_child(
1151 "root",
1152 &manifest,
1153 SchedulerBudget::default(),
1154 TaskLifecycle::Running,
1155 ),
1156 Err(TaskSpawnError::DuplicateTask)
1157 );
1158
1159 table.get_mut("root").unwrap().state = TaskLifecycle::Done(TerminationReason::Completed);
1160 let terminal_manifest = manifest_for("terminal-child");
1161 assert_eq!(
1162 table.spawn_child(
1163 "root",
1164 &terminal_manifest,
1165 SchedulerBudget::default(),
1166 TaskLifecycle::Running,
1167 ),
1168 Err(TaskSpawnError::CallerTerminal)
1169 );
1170 }
1171
1172 #[test]
1173 fn spawned_in_carries_requested_capabilities_as_the_grant() {
1174 use crate::types::capability::{
1175 ActionSet, Capability, CapabilityId, ConstraintSet, Principal, ResourceSelector,
1176 };
1177
1178 let mut manifest = manifest_for("child");
1179 let grant = Capability {
1180 id: CapabilityId("cap-1".into()),
1181 kind: crate::types::capability::CapabilityKind::Tool,
1182 resource: ResourceSelector("/repo/src/**".into()),
1183 actions: ActionSet(["read".into()].into_iter().collect()),
1184 constraints: ConstraintSet::default(),
1185 lease: None,
1186 delegatable: true,
1187 issuer: Principal("root".into()),
1188 };
1189 manifest.requested_capabilities = vec![grant.clone()];
1190
1191 let tcb = Tcb::spawned_in(
1192 &manifest,
1193 SchedulerBudget::default(),
1194 TaskLifecycle::Running,
1195 Some(TaskId::from("root")),
1196 );
1197 assert_eq!(tcb.capabilities, vec![grant]);
1198 }
1199
1200 #[test]
1201 fn root_task_has_no_capabilities_by_default() {
1202 let tcb = Tcb::root("root", SchedulerBudget::default());
1203 assert!(tcb.capabilities.is_empty());
1204 }
1205
1206 #[test]
1207 fn nested_spawn_records_real_parent() {
1208 let mut table_a = TaskTable::new();
1211 table_a.insert(Tcb::root("A", SchedulerBudget::default()));
1212 let manifest_b = manifest_for("B");
1213 let b = Tcb::spawned_in(
1214 &manifest_b,
1215 SchedulerBudget::default(),
1216 TaskLifecycle::Running,
1217 table_a.root_id(),
1218 );
1219 assert_eq!(b.parent, Some(TaskId::from("A")));
1220
1221 let mut table_b = TaskTable::new();
1224 table_b.insert(Tcb::root(b.id.clone(), SchedulerBudget::default()));
1225 let manifest_c = manifest_for("C");
1226 let c = Tcb::spawned_in(
1227 &manifest_c,
1228 SchedulerBudget::default(),
1229 TaskLifecycle::Running,
1230 table_b.root_id(),
1231 );
1232 assert_eq!(c.parent, Some(TaskId::from("B")));
1233 }
1234
1235 #[test]
1236 fn waiting_task_not_runnable() {
1237 let mut table = TaskTable::new();
1238 table.insert(Tcb::root("root", SchedulerBudget::default()));
1239 table.register_wait_set(
1240 "root",
1241 WaitSet {
1242 mode: WaitMode::Any,
1243 conditions: vec![WaitCondition::Approval(ApprovalId("pending".into()))],
1244 },
1245 );
1246
1247 assert!(
1248 !table.get("root").unwrap().state.occupies_slot(),
1249 "a Suspended/Waiting task must not occupy a concurrency slot"
1250 );
1251 assert_ne!(table.get("root").unwrap().state, TaskLifecycle::Ready);
1252 }
1253
1254 #[test]
1255 fn wake_is_idempotent() {
1256 let mut table = TaskTable::new();
1257 table.insert(Tcb::root("root", SchedulerBudget::default()));
1258 let effect = crate::runtime::kernel::wire::EffectId::new("e1").unwrap();
1259 table.wait_for_condition("root", &WaitCondition::Effect(effect.clone()));
1260 let key = crate::scheduler::wait_index::WaitKey::Effect(effect);
1261
1262 let first = table.wake(&key);
1263 assert_eq!(first, vec![TaskId::from("root")]);
1264
1265 let second = table.wake(&key);
1267 assert!(
1268 second.is_empty(),
1269 "a second wake for the same key must be a no-op"
1270 );
1271 }
1272
1273 #[test]
1274 fn unrelated_event_does_not_wake() {
1275 let mut table = TaskTable::new();
1276 table.insert(Tcb::root("root", SchedulerBudget::default()));
1277 let e1 = crate::runtime::kernel::wire::EffectId::new("e1").unwrap();
1278 let e2 = crate::runtime::kernel::wire::EffectId::new("e2").unwrap();
1279 table.wait_for_condition("root", &WaitCondition::Effect(e1.clone()));
1280
1281 let woken = table.wake(&crate::scheduler::wait_index::WaitKey::Effect(e2));
1282 assert!(woken.is_empty());
1283 assert_eq!(
1284 table
1285 .wait_index()
1286 .lookup(&crate::scheduler::wait_index::WaitKey::Effect(e1)),
1287 &[TaskId::from("root")],
1288 "task must still be registered as waiting on e1"
1289 );
1290 }
1291
1292 #[test]
1293 fn task_table_wait_set_any_mode_wakes_through_the_public_api() {
1294 let mut table = TaskTable::new();
1295 table.insert(Tcb::root("root", SchedulerBudget::default()));
1296 let e1 = crate::runtime::kernel::wire::EffectId::new("e1").unwrap();
1297 let e2 = crate::runtime::kernel::wire::EffectId::new("e2").unwrap();
1298 table.register_wait_set(
1299 "root",
1300 WaitSet {
1301 mode: WaitMode::Any,
1302 conditions: vec![WaitCondition::Effect(e1.clone()), WaitCondition::Effect(e2)],
1303 },
1304 );
1305
1306 let woken = table.notify(&crate::scheduler::wait_index::WaitKey::Effect(e1));
1307 assert_eq!(woken, vec![TaskId::from("root")]);
1308 }
1309
1310 #[test]
1311 fn durable_wait_set_tracks_all_progress_on_the_tcb_and_wakes_ready_once() {
1312 use crate::scheduler::wait_index::WaitKey;
1313
1314 let mut table = TaskTable::new();
1315 table.insert(Tcb::root("root", SchedulerBudget::default()));
1316 let first = crate::runtime::kernel::wire::EffectId::new("effect-1").unwrap();
1317 let second = crate::runtime::kernel::wire::EffectId::new("effect-2").unwrap();
1318 table.register_wait_set(
1319 "root",
1320 WaitSet {
1321 mode: WaitMode::All,
1322 conditions: vec![
1323 WaitCondition::Effect(first.clone()),
1324 WaitCondition::Effect(second.clone()),
1325 ],
1326 },
1327 );
1328
1329 assert_eq!(table.get("root").unwrap().state, TaskLifecycle::Suspended);
1330 assert_eq!(
1331 table.notify(&WaitKey::Effect(first.clone())),
1332 Vec::<TaskId>::new()
1333 );
1334 assert_eq!(
1335 table
1336 .get("root")
1337 .unwrap()
1338 .wait_set
1339 .as_ref()
1340 .unwrap()
1341 .satisfied
1342 .iter()
1343 .copied()
1344 .collect::<Vec<_>>(),
1345 vec![0],
1346 "partial All progress is durable task state, not reverse-index state"
1347 );
1348
1349 assert_eq!(
1350 table.notify(&WaitKey::Effect(first)),
1351 Vec::<TaskId>::new(),
1352 "a duplicate event cannot satisfy the missing condition"
1353 );
1354 assert_eq!(
1355 table.notify(&WaitKey::Effect(second)),
1356 vec![TaskId::from("root")]
1357 );
1358 assert_eq!(table.get("root").unwrap().state, TaskLifecycle::Ready);
1359 assert!(table.get("root").unwrap().wait_set.is_none());
1360 assert_eq!(
1361 table.notify(&WaitKey::Effect(
1362 crate::runtime::kernel::wire::EffectId::new("effect-2").unwrap()
1363 )),
1364 Vec::<TaskId>::new(),
1365 "a satisfied WaitSet wakes exactly once"
1366 );
1367 }
1368
1369 #[test]
1370 fn children_condition_in_all_mode_waits_for_every_child() {
1371 let mut table = TaskTable::new();
1372 table.insert(Tcb::root("root", SchedulerBudget::default()));
1373 table.register_wait_set(
1374 "root",
1375 WaitSet {
1376 mode: WaitMode::All,
1377 conditions: vec![WaitCondition::Children(vec!["a".into(), "b".into()])],
1378 },
1379 );
1380
1381 assert!(
1382 table
1383 .notify(&super::super::wait_index::WaitKey::Child("a".into()))
1384 .is_empty()
1385 );
1386 assert_eq!(
1387 table.notify(&super::super::wait_index::WaitKey::Child("b".into())),
1388 vec![TaskId::from("root")]
1389 );
1390 }
1391
1392 #[test]
1393 fn terminal_waiter_is_removed_without_becoming_runnable() {
1394 use crate::scheduler::wait_index::WaitKey;
1395
1396 let mut table = TaskTable::new();
1397 table.insert(Tcb::root("root", SchedulerBudget::default()));
1398 let effect = crate::runtime::kernel::wire::EffectId::new("effect-terminal").unwrap();
1399 table.register_wait_set(
1400 "root",
1401 WaitSet {
1402 mode: WaitMode::Any,
1403 conditions: vec![WaitCondition::Effect(effect.clone())],
1404 },
1405 );
1406 table.get_mut("root").unwrap().state = TaskLifecycle::Done(TerminationReason::UserAbort);
1407
1408 assert!(table.notify(&WaitKey::Effect(effect.clone())).is_empty());
1409 assert_eq!(
1410 table.get("root").unwrap().state,
1411 TaskLifecycle::Done(TerminationReason::UserAbort)
1412 );
1413 assert!(table.get("root").unwrap().wait_set.is_none());
1414 assert!(
1415 table
1416 .wait_index()
1417 .lookup(&WaitKey::Effect(effect))
1418 .is_empty()
1419 );
1420 }
1421
1422 #[test]
1423 fn wake_order_is_identical_before_and_after_reverse_index_rebuild() {
1424 use crate::scheduler::wait_index::WaitKey;
1425
1426 let mut live = TaskTable::new();
1427 live.insert(Tcb::root("root", SchedulerBudget::default()));
1428 live.insert(Tcb::root("task-a", SchedulerBudget::default()));
1429 live.insert(Tcb::root("task-b", SchedulerBudget::default()));
1430 let effect = crate::runtime::kernel::wire::EffectId::new("shared-effect").unwrap();
1431
1432 live.wait_for_condition("task-b", &WaitCondition::Effect(effect.clone()));
1435 live.wait_for_condition("task-a", &WaitCondition::Effect(effect.clone()));
1436 let mut restored = live.clone();
1437 restored.rebuild_wait_index();
1438
1439 let key = WaitKey::Effect(effect);
1440 let live_order = live.notify(&key);
1441 let restored_order = restored.notify(&key);
1442 assert_eq!(live_order, restored_order);
1443 assert_eq!(
1444 live_order,
1445 vec![TaskId::from("task-a"), TaskId::from("task-b")]
1446 );
1447 }
1448
1449 #[test]
1450 fn terminal_waiter_loses_all_wait_state_without_resuming() {
1451 use crate::scheduler::wait_index::WaitKey;
1452
1453 let mut table = TaskTable::new();
1454 table.insert(Tcb::root("root", SchedulerBudget::default()));
1455 table.register_wait_set(
1456 "root",
1457 WaitSet {
1458 mode: WaitMode::Any,
1459 conditions: vec![WaitCondition::Approval(ApprovalId("pending".into()))],
1460 },
1461 );
1462 table.get_mut("root").unwrap().state = TaskLifecycle::Done(TerminationReason::UserAbort);
1463
1464 let key = WaitKey::Approval(ApprovalId(CompactString::from("pending")));
1465 assert!(table.notify(&key).is_empty());
1466 let task = table.get("root").unwrap();
1467 assert_eq!(
1468 task.state,
1469 TaskLifecycle::Done(TerminationReason::UserAbort)
1470 );
1471 assert!(task.wait_set.is_none());
1472 assert!(table.wait_index().lookup(&key).is_empty());
1473 }
1474
1475 #[test]
1476 fn timer_wakes_exactly_at_deadline_not_before() {
1477 let mut table = TaskTable::new();
1478 table.insert(Tcb::root("root", SchedulerBudget::default()));
1479 table.wait_for_timer("root", LogicalDeadline(1_000));
1480 let key = crate::scheduler::wait_index::WaitKey::Timer(LogicalDeadline(1_000));
1481 assert_eq!(table.wait_index().lookup(&key), &[TaskId::from("root")]);
1482
1483 let woken_early = table.wake_expired_timers(999);
1484 assert!(woken_early.is_empty(), "must not wake before the deadline");
1485 assert_eq!(table.wait_index().lookup(&key), &[TaskId::from("root")]);
1486
1487 let woken = table.wake_expired_timers(1_000);
1488 assert_eq!(woken, vec![TaskId::from("root")]);
1489 assert!(table.wait_index().lookup(&key).is_empty());
1490 }
1491
1492 #[test]
1493 fn wait_set_syncs_the_wait_index() {
1494 let mut table = TaskTable::new();
1495 table.insert(Tcb::root("root", SchedulerBudget::default()));
1496
1497 let approval = ApprovalId("pending".into());
1498 table.register_wait_set(
1499 "root",
1500 WaitSet {
1501 mode: WaitMode::Any,
1502 conditions: vec![WaitCondition::Approval(approval.clone())],
1503 },
1504 );
1505 assert!(table.get("root").unwrap().wait_set.is_some());
1506 let key = crate::scheduler::wait_index::WaitKey::Approval(approval);
1507 assert_eq!(table.wait_index().lookup(&key), &[TaskId::from("root")]);
1508
1509 table.clear_wait("root");
1510 assert!(table.get("root").unwrap().wait_set.is_none());
1511 assert!(table.wait_index().lookup(&key).is_empty());
1512 }
1513
1514 #[test]
1515 fn wait_condition_and_wait_mode_variants_compile() {
1516 let conditions = vec![
1517 WaitCondition::Effect(crate::runtime::kernel::wire::EffectId::new("e1").unwrap()),
1518 WaitCondition::Child(TaskId::from("child-1")),
1519 WaitCondition::Children(vec![TaskId::from("child-1"), TaskId::from("child-2")]),
1520 WaitCondition::Approval(ApprovalId("approval-1".into())),
1521 WaitCondition::Signal(SignalFilter("topic:*".into())),
1522 WaitCondition::Timer(LogicalDeadline(1_000)),
1523 WaitCondition::Channel(ChannelId("chan-1".into())),
1524 WaitCondition::Resource(ResourceKey("res-1".into())),
1525 WaitCondition::External(SubscriptionId("sub-1".into())),
1526 ];
1527 for condition in &conditions {
1528 match condition {
1529 WaitCondition::Effect(_)
1530 | WaitCondition::Child(_)
1531 | WaitCondition::Children(_)
1532 | WaitCondition::Approval(_)
1533 | WaitCondition::Signal(_)
1534 | WaitCondition::Timer(_)
1535 | WaitCondition::Channel(_)
1536 | WaitCondition::Resource(_)
1537 | WaitCondition::External(_) => {}
1538 }
1539 }
1540
1541 let wait_set = WaitSet {
1542 mode: WaitMode::Any,
1543 conditions,
1544 };
1545 assert_eq!(wait_set.mode, WaitMode::Any);
1546 assert_eq!(wait_set.conditions.len(), 9);
1547 }
1548
1549 #[test]
1550 fn rebuild_children_recovers_lineage_after_out_of_order_restore() {
1551 let mut table = TaskTable::new();
1556 let mut child = Tcb::root("child", SchedulerBudget::default());
1557 child.parent = Some(TaskId::from("root"));
1558 table.insert(child); table.insert(Tcb::root("root", SchedulerBudget::default()));
1560
1561 assert!(
1562 !table
1563 .get("root")
1564 .unwrap()
1565 .children
1566 .contains(&TaskId::from("child")),
1567 "sanity: out-of-order insert must NOT have already fixed itself"
1568 );
1569
1570 table.rebuild_children();
1571
1572 assert!(
1573 table
1574 .get("root")
1575 .unwrap()
1576 .children
1577 .contains(&TaskId::from("child"))
1578 );
1579 }
1580
1581 #[test]
1582 fn rebuild_wait_index_recovers_a_restored_waits_index_entry() {
1583 let mut table = TaskTable::new();
1584 let mut root = Tcb::root("root", SchedulerBudget::default());
1585 root.wait_set = Some(DurableWaitSet {
1586 mode: WaitMode::Any,
1587 conditions: vec![WaitCondition::Approval(ApprovalId("pending".into()))],
1588 satisfied: BTreeSet::new(),
1589 });
1590 table.insert(root);
1591
1592 let key = crate::scheduler::wait_index::WaitKey::Approval(ApprovalId("pending".into()));
1593 assert!(
1594 table.wait_index().lookup(&key).is_empty(),
1595 "inserting a Tcb does not mutate the derived WaitIndex"
1596 );
1597
1598 table.rebuild_wait_index();
1599
1600 assert_eq!(table.wait_index().lookup(&key), &[TaskId::from("root")]);
1601 }
1602
1603 #[test]
1604 fn rebuild_wait_index_is_idempotent() {
1605 let mut table = TaskTable::new();
1606 let mut root = Tcb::root("root", SchedulerBudget::default());
1607 root.wait_set = Some(DurableWaitSet {
1608 mode: WaitMode::Any,
1609 conditions: vec![WaitCondition::Approval(ApprovalId("pending".into()))],
1610 satisfied: BTreeSet::new(),
1611 });
1612 table.insert(root);
1613
1614 table.rebuild_wait_index();
1615 table.rebuild_wait_index();
1616
1617 let key = crate::scheduler::wait_index::WaitKey::Approval(ApprovalId("pending".into()));
1618 assert_eq!(
1619 table.wait_index().lookup(&key),
1620 &[TaskId::from("root")],
1621 "calling rebuild_wait_index twice must not duplicate the entry"
1622 );
1623 }
1624
1625 #[test]
1626 fn has_cycle_is_false_on_a_normal_tree() {
1627 let mut table = TaskTable::new();
1628 table.insert(Tcb::root("A", SchedulerBudget::default()));
1629 let b = Tcb::spawned_in(
1630 &manifest_for("B"),
1631 SchedulerBudget::default(),
1632 TaskLifecycle::Running,
1633 table.root_id(),
1634 );
1635 table.insert(b);
1636 let c = Tcb::spawned_in(
1637 &manifest_for("C"),
1638 SchedulerBudget::default(),
1639 TaskLifecycle::Running,
1640 Some(TaskId::from("B")),
1641 );
1642 table.insert(c);
1643
1644 assert!(!table.has_cycle());
1645 }
1646
1647 #[test]
1648 fn has_cycle_detects_a_manually_constructed_cycle() {
1649 let mut table = TaskTable::new();
1653 table.insert(Tcb::root("A", SchedulerBudget::default()));
1654 let b = Tcb::spawned_in(
1655 &manifest_for("B"),
1656 SchedulerBudget::default(),
1657 TaskLifecycle::Running,
1658 table.root_id(),
1659 );
1660 table.insert(b);
1661 table.get_mut("A").unwrap().parent = Some(TaskId::from("B"));
1663
1664 assert!(table.has_cycle());
1665 }
1666
1667 #[test]
1668 fn supervision_policy_defaults_match_spec_section_4() {
1669 let root = Tcb::root("root", SchedulerBudget::default());
1670 assert_eq!(
1671 root.supervision.child_failure,
1672 ChildFailurePolicy::Propagate
1673 );
1674 assert_eq!(root.supervision.max_restarts, None);
1675 assert!(root.supervision.cancel_children_on_exit);
1676
1677 let manifest = manifest_for("child");
1678 let child = Tcb::spawned_in(
1679 &manifest,
1680 SchedulerBudget::default(),
1681 TaskLifecycle::Running,
1682 Some(TaskId::from("root")),
1683 );
1684 assert_eq!(child.supervision, SupervisionPolicy::default());
1685 }
1686
1687 #[test]
1688 fn spc_008_05_cancel_subtree_exempts_a_detached_child_and_its_own_descendants() {
1689 let mut table = TaskTable::new();
1690 table.insert(Tcb::root("A", SchedulerBudget::default()));
1691 let mut b = Tcb::spawned_in(
1692 &manifest_for("B"),
1693 SchedulerBudget::default(),
1694 TaskLifecycle::Running,
1695 table.root_id(),
1696 );
1697 b.detached = true;
1698 table.insert(b);
1699 let c = Tcb::spawned_in(
1700 &manifest_for("C"),
1701 SchedulerBudget::default(),
1702 TaskLifecycle::Running,
1703 Some(TaskId::from("B")),
1704 );
1705 table.insert(c);
1706
1707 table.cancel_subtree("A");
1708
1709 assert_eq!(
1710 table.get("A").unwrap().state,
1711 TaskLifecycle::Done(TerminationReason::UserAbort),
1712 "A itself must still be cancelled"
1713 );
1714 assert_eq!(
1715 table.get("B").unwrap().state,
1716 TaskLifecycle::Running,
1717 "detached B must not be cancelled by A's cancellation"
1718 );
1719 assert_eq!(
1720 table.get("C").unwrap().state,
1721 TaskLifecycle::Running,
1722 "C (B's own child) must not be reached either — detaching exempts the whole subtree"
1723 );
1724 }
1725
1726 #[test]
1727 fn cancel_subtree_terminates_the_whole_tree() {
1728 let mut table = TaskTable::new();
1729 table.insert(Tcb::root("A", SchedulerBudget::default()));
1730 let b = Tcb::spawned_in(
1731 &manifest_for("B"),
1732 SchedulerBudget::default(),
1733 TaskLifecycle::Running,
1734 table.root_id(),
1735 );
1736 table.insert(b);
1737 let c = Tcb::spawned_in(
1738 &manifest_for("C"),
1739 SchedulerBudget::default(),
1740 TaskLifecycle::Running,
1741 Some(TaskId::from("B")),
1742 );
1743 table.insert(c);
1744
1745 table.cancel_subtree("A");
1746
1747 for id in ["A", "B", "C"] {
1748 assert_eq!(
1749 table.get(id).unwrap().state,
1750 TaskLifecycle::Done(TerminationReason::UserAbort),
1751 "task {id} should be cancelled"
1752 );
1753 }
1754 }
1755
1756 #[test]
1757 fn cancel_subtree_on_a_leaf_only_cancels_itself() {
1758 let mut table = TaskTable::new();
1759 table.insert(Tcb::root("root", SchedulerBudget::default()));
1760 table.cancel_subtree("root");
1761 assert_eq!(
1762 table.get("root").unwrap().state,
1763 TaskLifecycle::Done(TerminationReason::UserAbort)
1764 );
1765 }
1766
1767 #[test]
1768 fn spawn_registers_child_in_parent_children() {
1769 let mut table = TaskTable::new();
1770 table.insert(Tcb::root("root", SchedulerBudget::default()));
1771 let manifest = manifest_for("child-1");
1772 let child = Tcb::spawned_in(
1773 &manifest,
1774 SchedulerBudget::default(),
1775 TaskLifecycle::Running,
1776 table.root_id(),
1777 );
1778 table.insert(child.clone());
1779
1780 assert!(
1781 table
1782 .get("root")
1783 .unwrap()
1784 .children
1785 .contains(&TaskId::from("child-1"))
1786 );
1787
1788 let manifest_2 = manifest_for("child-2");
1790 let child_2 = Tcb::spawned_in(
1791 &manifest_2,
1792 SchedulerBudget::default(),
1793 TaskLifecycle::Running,
1794 table.root_id(),
1795 );
1796 table.insert(child_2);
1797 let root_children = &table.get("root").unwrap().children;
1798 assert!(root_children.contains(&TaskId::from("child-1")));
1799 assert!(root_children.contains(&TaskId::from("child-2")));
1800 }
1801
1802 #[test]
1803 fn tcb_children_default_empty() {
1804 let tcb = Tcb::root("root", SchedulerBudget::default());
1805 assert!(tcb.children.is_empty());
1806 }
1807
1808 #[test]
1809 fn process_state_maps_to_lifecycle() {
1810 assert_eq!(
1811 TaskLifecycle::from(ProcessState::Running),
1812 TaskLifecycle::Running
1813 );
1814 assert_eq!(
1815 TaskLifecycle::from(ProcessState::Joined),
1816 TaskLifecycle::Done(TerminationReason::Completed)
1817 );
1818 assert_eq!(
1819 TaskLifecycle::from(ProcessState::Failed),
1820 TaskLifecycle::Done(TerminationReason::Error)
1821 );
1822 }
1823
1824 #[test]
1825 fn budget_ledger_delegates_to_scheduler_budget() {
1826 let mut ledger = BudgetLedger::new(SchedulerBudget {
1827 max_turns: 2,
1828 ..SchedulerBudget::default()
1829 });
1830 assert_eq!(ledger.exceeded(None), None);
1831 ledger.turns = 2;
1832 assert_eq!(ledger.exceeded(None), Some("max_turns"));
1833 }
1834
1835 #[test]
1836 fn task_table_insert_and_lineage() {
1837 let mut table = TaskTable::new();
1838 table.insert(Tcb::root("root", SchedulerBudget::default()));
1839 let mut child = Tcb::root("child", SchedulerBudget::default());
1840 child.parent = Some("root".into());
1841 table.insert(child);
1842
1843 assert_eq!(table.children_of("root").len(), 1);
1844 assert!(table.get("root").is_some());
1845 }
1846
1847 #[test]
1848 fn task_table_insert_is_idempotent_by_id() {
1849 let mut table = TaskTable::new();
1850 table.insert(Tcb::root("root", SchedulerBudget::default()));
1851 let mut updated = Tcb::root("root", SchedulerBudget::default());
1852 updated.state = TaskLifecycle::Running;
1853 table.insert(updated);
1854
1855 assert_eq!(table.all().len(), 1);
1856 assert_eq!(table.get("root").unwrap().state, TaskLifecycle::Running);
1857 }
1858
1859 #[test]
1860 fn budget_verdict_none_within_budget() {
1861 let tcb = Tcb::root(
1862 "root",
1863 SchedulerBudget {
1864 max_turns: 5,
1865 ..SchedulerBudget::default()
1866 },
1867 );
1868 assert_eq!(budget_verdict(&tcb, None), None);
1869 }
1870
1871 #[test]
1872 fn budget_verdict_matches_should_terminate_axis() {
1873 let limits = SchedulerBudget {
1874 max_turns: 2,
1875 ..SchedulerBudget::default()
1876 };
1877 let mut tcb = Tcb::root("root", limits.clone());
1878 tcb.budget.turns = 2;
1879 assert_eq!(limits.should_terminate(2, 0, None, None), Some("max_turns"));
1881 assert_eq!(
1882 budget_verdict(&tcb, None),
1883 Some(TerminationReason::MaxTurns)
1884 );
1885 }
1886
1887 #[test]
1888 fn budget_verdict_wall_time_maps_to_timeout() {
1889 let limits = SchedulerBudget {
1890 max_wall_ms: Some(1_000),
1891 ..SchedulerBudget::default()
1892 };
1893 let mut tcb = Tcb::root("root", limits);
1894 tcb.budget.started_at_ms = Some(0);
1895 assert_eq!(
1896 budget_verdict(&tcb, Some(2_000)),
1897 Some(TerminationReason::Timeout)
1898 );
1899 }
1900
1901 #[test]
1902 fn baseline_token_budget_terminates() {
1903 let limits = SchedulerBudget {
1904 max_total_tokens: 100,
1905 ..SchedulerBudget::default()
1906 };
1907 let mut tcb = Tcb::root("root", limits);
1908 tcb.budget.total_tokens = 200;
1909 assert_eq!(
1910 budget_verdict(&tcb, None),
1911 Some(TerminationReason::TokenBudget)
1912 );
1913 }
1914
1915 #[test]
1916 fn spc_006_03_send_message_delivers_into_the_recipients_mailbox() {
1917 use crate::scheduler::mailbox::{LogicalTime, MailboxMessage, MessageId};
1918 use crate::types::signal::Urgency;
1919
1920 let mut table = TaskTable::new();
1921 table.insert(Tcb::root("a", SchedulerBudget::default()));
1922 table.insert(Tcb::root("b", SchedulerBudget::default()));
1923
1924 table.send_message(MailboxMessage {
1925 id: MessageId::from("msg-1"),
1926 from: TaskId::from("a"),
1927 to: TaskId::from("b"),
1928 kind: CompactString::from("research_result"),
1929 payload_handle: 1,
1930 priority: Urgency::Normal,
1931 timestamp: LogicalTime(0),
1932 expires_at: None,
1933 });
1934
1935 let received = table
1936 .get_mut("b")
1937 .unwrap()
1938 .mailbox
1939 .receive()
1940 .expect("B must have received A's message");
1941 assert_eq!(received.from, TaskId::from("a"));
1942 assert_eq!(received.id, MessageId::from("msg-1"));
1943 assert!(table.get_mut("a").unwrap().mailbox.receive().is_none());
1944 }
1945
1946 #[test]
1947 fn spc_006_03_send_message_to_an_unknown_task_is_dropped_not_panicking() {
1948 use crate::scheduler::mailbox::{LogicalTime, MailboxMessage, MessageId};
1949 use crate::types::signal::Urgency;
1950
1951 let mut table = TaskTable::new();
1952 table.insert(Tcb::root("a", SchedulerBudget::default()));
1953
1954 table.send_message(MailboxMessage {
1955 id: MessageId::from("msg-1"),
1956 from: TaskId::from("a"),
1957 to: TaskId::from("nobody"),
1958 kind: CompactString::from("kind"),
1959 payload_handle: 1,
1960 priority: Urgency::Normal,
1961 timestamp: LogicalTime(0),
1962 expires_at: None,
1963 });
1964 }
1966
1967 #[test]
1968 fn spc_019_08_local_ipc_derives_sender_and_wakes_mailbox_and_channel_waiters() {
1969 use crate::scheduler::mailbox::{LogicalTime, MailboxMessage};
1970 use crate::scheduler::wait_index::WaitKey;
1971 use crate::types::signal::Urgency;
1972
1973 let mut table = TaskTable::new();
1974 table.insert(Tcb::root("a", SchedulerBudget::default()));
1975 table.insert(Tcb::root("b", SchedulerBudget::default()));
1976 let mailbox_subscription = SubscriptionId("mailbox:b".into());
1977 table.wait_for_condition("b", &WaitCondition::External(mailbox_subscription.clone()));
1978 let message = MailboxMessage {
1979 id: "m1".into(),
1980 from: "forged".into(),
1981 to: "b".into(),
1982 kind: "result".into(),
1983 payload_handle: 7,
1984 priority: Urgency::Normal,
1985 timestamp: LogicalTime(1),
1986 expires_at: None,
1987 };
1988 assert_eq!(
1989 table.send_message_from("a", message.clone(), LogicalTime(1)),
1990 Ok(true)
1991 );
1992 assert_eq!(table.get("b").unwrap().state, TaskLifecycle::Ready);
1993 assert!(
1994 table
1995 .wait_index()
1996 .lookup(&WaitKey::External(mailbox_subscription))
1997 .is_empty()
1998 );
1999 let received = table.receive_mailbox("b", LogicalTime(1), 1).unwrap();
2000 assert_eq!(received[0].from.as_str(), "a");
2001 assert_eq!(
2002 table.send_message_from("a", message, LogicalTime(1)),
2003 Ok(false),
2004 "redelivery is a durable no-op"
2005 );
2006
2007 let channel_id = ChannelId("results".into());
2008 table.wait_for_condition("b", &WaitCondition::Channel(channel_id.clone()));
2009 assert_eq!(
2010 table.publish_channel(
2011 "a",
2012 channel_id.clone(),
2013 vec!["b".into()],
2014 MailboxMessage {
2015 id: "cm1".into(),
2016 to: "ignored".into(),
2017 ..received[0].clone()
2018 },
2019 LogicalTime(1),
2020 ),
2021 Ok(true)
2022 );
2023 assert_eq!(table.get("b").unwrap().state, TaskLifecycle::Ready);
2024 assert_eq!(
2025 table
2026 .receive_channel("b", &channel_id, LogicalTime(1))
2027 .unwrap()
2028 .len(),
2029 1
2030 );
2031 }
2032
2033 #[test]
2034 fn spc_019_09_object_registration_wakes_resource_wait_and_receiver_capability_gates_read() {
2035 use crate::mm::handle::{Handle, HandleKind, ObjectDescriptor, object_access_allowed};
2036 use crate::types::capability::{
2037 ActionSet, Capability, CapabilityId, CapabilityKind, ConstraintSet, Principal,
2038 ResourceSelector,
2039 };
2040
2041 let mut table = TaskTable::new();
2042 table.insert(Tcb::root("a", SchedulerBudget::default()));
2043 table.insert(Tcb::root("b", SchedulerBudget::default()));
2044 let resource = ResourceKey("object:a/7".into());
2045 table.wait_for_condition("b", &WaitCondition::Resource(resource));
2046 let descriptor = ObjectDescriptor::from_handle(
2047 "a".into(),
2048 &Handle::resident(7, HandleKind::ToolResult, 10),
2049 1,
2050 );
2051 assert_eq!(table.register_object("a", descriptor.clone()), Ok(true));
2052 assert_eq!(table.get("b").unwrap().state, TaskLifecycle::Ready);
2053 assert!(!object_access_allowed(
2054 &table.get("b").unwrap().capabilities,
2055 "read",
2056 &descriptor
2057 ));
2058
2059 table.get_mut("b").unwrap().capabilities = vec![Capability {
2060 id: CapabilityId("read-a-7".into()),
2061 kind: CapabilityKind::Tool,
2062 resource: ResourceSelector("object:a/7".into()),
2063 actions: ActionSet(["read".into()].into_iter().collect()),
2064 constraints: ConstraintSet::default(),
2065 lease: None,
2066 delegatable: false,
2067 issuer: Principal("a".into()),
2068 }];
2069 assert!(object_access_allowed(
2070 &table.get("b").unwrap().capabilities,
2071 "read",
2072 &descriptor
2073 ));
2074 assert_eq!(table.register_object("a", descriptor), Ok(false));
2075 }
2076
2077 #[test]
2078 fn spc_019_11_workflow_nested_timer_and_message_work_share_one_stable_trace() {
2079 use crate::scheduler::runnable::{LocalRunnable, LocalRunnableKind, order_runnables};
2080
2081 let mut table = TaskTable::new();
2082 table.insert(Tcb::root("root", SchedulerBudget::default()));
2083 table.insert(Tcb::spawned_in(
2084 &manifest_for("nested"),
2085 SchedulerBudget::default(),
2086 TaskLifecycle::Ready,
2087 Some("root".into()),
2088 ));
2089 table.insert(Tcb::spawned_in(
2090 &manifest_for("timer"),
2091 SchedulerBudget::default(),
2092 TaskLifecycle::Running,
2093 Some("root".into()),
2094 ));
2095 table.insert(Tcb::spawned_in(
2096 &manifest_for("mail"),
2097 SchedulerBudget::default(),
2098 TaskLifecycle::Running,
2099 Some("root".into()),
2100 ));
2101 table.wait_for_timer("timer", LogicalDeadline(10));
2102 table.wait_for_condition(
2103 "mail",
2104 &WaitCondition::External(SubscriptionId("mailbox:mail".into())),
2105 );
2106 table.wake_expired_timers(10);
2107 table.notify(&super::super::wait_index::WaitKey::External(
2108 SubscriptionId("mailbox:mail".into()),
2109 ));
2110
2111 let mut candidates = table.runnable_candidates();
2112 candidates.push(LocalRunnable::workflow("wf-node2", 0));
2113 let trace = order_runnables(candidates);
2114 assert_eq!(
2115 trace
2116 .iter()
2117 .map(|entry| (entry.id.as_str(), entry.kind))
2118 .collect::<Vec<_>>(),
2119 vec![
2120 ("mail", LocalRunnableKind::MessageWaiter),
2121 ("nested", LocalRunnableKind::NestedTask),
2122 ("timer", LocalRunnableKind::TimerWaiter),
2123 ("wf-node2", LocalRunnableKind::WorkflowNode),
2124 ]
2125 );
2126
2127 let mut restored = table.clone();
2128 restored.rebuild_wait_index();
2129 let mut restored_candidates = restored.runnable_candidates();
2130 restored_candidates.push(LocalRunnable::workflow("wf-node2", 0));
2131 assert_eq!(order_runnables(restored_candidates), trace);
2132 }
2133}