deepstrike_core/scheduler/
tcb.rs

1//! Primitive P2: Task Control Block + unified scheduling entity.
2//!
3//! See `.local-docs/specs/agent-os-three-primitives.md`. The root loop and every
4//! sub-agent are a single [`Tcb`]; the [`TaskTable`] is the sole source of truth for
5//! schedulability and lineage, and the `AgentProcess` view is derived from it.
6//! [`budget_verdict`] is the single budget decision point (turn/token/wall axes),
7//! delegating to [`SchedulerBudget::should_terminate`] via [`BudgetLedger`].
8
9use compact_str::CompactString;
10use serde::{Deserialize, Serialize};
11
12use crate::proc::ProcessState;
13use crate::scheduler::policy::SchedulerBudget;
14use crate::types::agent::{AgentIsolation, AgentRole, ContextInheritance, IsolationManifest};
15use crate::types::result::{SubAgentResult, TerminationReason};
16
17/// Identity of a schedulable task. Task 0 is the root loop; children are sub-agents.
18/// Aligns with `AgentProcess.agent_id` so process rows map onto TCBs 1:1.
19pub type TaskId = CompactString;
20
21/// Schedulability lifecycle of a task — orthogonal to the *intra-turn* step,
22/// which stays on [`crate::scheduler::state_machine::LoopPhase`], and distinct from
23/// the task-goal blackboard [`crate::context::task_state::TaskState`].
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
25#[serde(rename_all = "snake_case")]
26pub enum TaskLifecycle {
27    /// Eligible to run, not yet picked by the scheduler.
28    Ready,
29    /// Currently executing a turn (`ProcessState::Running`).
30    Running,
31    /// Suspended awaiting external resolution (human approval / sub-agent join).
32    Suspended,
33    /// Finished. Carries the termination reason (`ProcessState::{Joined,Failed}`).
34    Done(TerminationReason),
35}
36
37impl TaskLifecycle {
38    pub fn label(self) -> &'static str {
39        match self {
40            Self::Ready => "ready",
41            Self::Running => "running",
42            Self::Suspended => "suspended",
43            Self::Done(_) => "done",
44        }
45    }
46
47    pub fn is_terminal(self) -> bool {
48        matches!(self, Self::Done(_))
49    }
50}
51
52/// A successful join maps to `Done(Completed)`; any other termination is `Done(<reason>)`.
53impl From<ProcessState> for TaskLifecycle {
54    fn from(state: ProcessState) -> Self {
55        match state {
56            ProcessState::Running => TaskLifecycle::Running,
57            ProcessState::Joined => TaskLifecycle::Done(TerminationReason::Completed),
58            // Failed has no single reason at the process level; the real reason travels
59            // in `SubAgentResult`. This projection maps to a generic error.
60            ProcessState::Failed => TaskLifecycle::Done(TerminationReason::Error),
61        }
62    }
63}
64
65/// Why a suspended task is not runnable. Only the reasons production actually
66/// constructs exist; new wait states earn a variant when they earn a producer.
67#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
68#[serde(rename_all = "snake_case")]
69pub enum WaitReason {
70    /// Governance `AskUser` — waiting for SDK to resolve human approval.
71    Approval,
72    /// Parent blocked on child tasks' join results. Tracks pending child IDs.
73    SubAgentJoin(Vec<TaskId>),
74}
75
76impl WaitReason {
77    pub fn label(&self) -> &'static str {
78        match self {
79            Self::Approval => "approval",
80            Self::SubAgentJoin(_) => "sub_agent_join",
81        }
82    }
83}
84
85/// Running budget counters + limits for a task. Wraps the existing [`SchedulerBudget`]
86/// limits so budget evaluation lives here without changing the axes.
87#[derive(Debug, Clone)]
88pub struct BudgetLedger {
89    pub limits: SchedulerBudget,
90    pub turns: u32,
91    pub total_tokens: u64,
92    pub started_at_ms: Option<u64>,
93}
94
95impl BudgetLedger {
96    pub fn new(limits: SchedulerBudget) -> Self {
97        Self {
98            limits,
99            turns: 0,
100            total_tokens: 0,
101            started_at_ms: None,
102        }
103    }
104
105    /// Delegates to the existing budget logic — single source of truth, no axis drift.
106    pub fn exceeded(&self, now_ms: Option<u64>) -> Option<&'static str> {
107        self.limits
108            .should_terminate(self.turns, self.total_tokens, now_ms, self.started_at_ms)
109    }
110}
111
112impl Default for BudgetLedger {
113    fn default() -> Self {
114        Self::new(SchedulerBudget::default())
115    }
116}
117
118/// Sub-agent-specific identity carried by a child [`Tcb`]; `None` on the root task.
119///
120/// This is what makes the `AgentProcess` view *derived* from the [`TaskTable`]: every child task
121/// whose `proc` is `Some` reconstructs exactly one [`crate::proc::AgentProcess`] (see
122/// [`crate::proc::AgentProcess::from_tcb`]).
123#[derive(Debug, Clone)]
124pub struct ProcInfo {
125    pub parent_session_id: CompactString,
126    pub role: AgentRole,
127    pub isolation: AgentIsolation,
128    pub context_inheritance: ContextInheritance,
129    /// The join result once the sub-agent has completed; `None` while running.
130    pub result: Option<SubAgentResult>,
131}
132
133/// One schedulable entity. The root loop and every sub-agent are uniform `Tcb`s.
134#[derive(Debug, Clone)]
135pub struct Tcb {
136    pub id: TaskId,
137    pub parent: Option<TaskId>,
138    pub state: TaskLifecycle,
139    pub budget: BudgetLedger,
140    pub wait: Option<WaitReason>,
141    /// Capability ids permitted to this task (mirrors `AgentProcess.permitted_capability_ids`).
142    pub caps: Vec<CompactString>,
143    /// Sub-agent identity for child tasks; `None` for the root loop.
144    pub proc: Option<ProcInfo>,
145}
146
147impl Tcb {
148    /// The root loop task. Constructed from the runtime task at `Start`.
149    pub fn root(id: impl Into<TaskId>, budget: SchedulerBudget) -> Self {
150        Self {
151            id: id.into(),
152            parent: None,
153            state: TaskLifecycle::Ready,
154            budget: BudgetLedger::new(budget),
155            wait: None,
156            caps: Vec::new(),
157            proc: None,
158        }
159    }
160
161    /// A sub-agent task spawned under the root, seeded `Running`, carrying the manifest's
162    /// process identity. The single source of truth for what the `AgentProcess` view exposes.
163    pub fn spawned(manifest: &IsolationManifest, budget: SchedulerBudget) -> Self {
164        Self {
165            id: manifest.agent_id.clone(),
166            parent: Some("root".into()),
167            state: TaskLifecycle::Running,
168            budget: BudgetLedger::new(budget),
169            wait: None,
170            caps: manifest.permitted_capability_ids.clone(),
171            proc: Some(ProcInfo {
172                parent_session_id: manifest.parent_session_id.clone(),
173                role: manifest.role,
174                isolation: manifest.isolation,
175                context_inheritance: manifest.context_inheritance,
176                result: None,
177            }),
178        }
179    }
180}
181
182/// Unified registry of all tasks: the root loop plus one child per sub-agent. The sole source of
183/// truth for schedulability and lineage; the `AgentProcess` view is derived from it.
184#[derive(Debug, Clone, Default)]
185pub struct TaskTable {
186    tasks: Vec<Tcb>,
187}
188
189impl TaskTable {
190    pub fn new() -> Self {
191        Self::default()
192    }
193
194    pub fn insert(&mut self, tcb: Tcb) {
195        if let Some(existing) = self.tasks.iter_mut().find(|t| t.id == tcb.id) {
196            *existing = tcb;
197        } else {
198            self.tasks.push(tcb);
199        }
200    }
201
202    pub fn get(&self, id: &str) -> Option<&Tcb> {
203        self.tasks.iter().find(|t| t.id.as_str() == id)
204    }
205
206    pub fn get_mut(&mut self, id: &str) -> Option<&mut Tcb> {
207        self.tasks.iter_mut().find(|t| t.id.as_str() == id)
208    }
209
210    pub fn all(&self) -> &[Tcb] {
211        &self.tasks
212    }
213
214    pub fn children_of(&self, parent: &str) -> Vec<&Tcb> {
215        self.tasks
216            .iter()
217            .filter(|t| t.parent.as_deref() == Some(parent))
218            .collect()
219    }
220}
221
222/// Pure budget verdict for one task: `Some(reason)` when a budget axis (turn/token/wall)
223/// is exhausted, mapped to the same `TerminationReason` the state machine applies.
224/// The single budget decision point — evaluated at each turn boundary.
225pub fn budget_verdict(task: &Tcb, now_ms: Option<u64>) -> Option<TerminationReason> {
226    task.budget.exceeded(now_ms).map(|axis| match axis {
227        "max_turns" => TerminationReason::MaxTurns,
228        "wall_time" => TerminationReason::Timeout,
229        _ => TerminationReason::TokenBudget,
230    })
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    #[test]
238    fn process_state_maps_to_lifecycle() {
239        assert_eq!(
240            TaskLifecycle::from(ProcessState::Running),
241            TaskLifecycle::Running
242        );
243        assert_eq!(
244            TaskLifecycle::from(ProcessState::Joined),
245            TaskLifecycle::Done(TerminationReason::Completed)
246        );
247        assert_eq!(
248            TaskLifecycle::from(ProcessState::Failed),
249            TaskLifecycle::Done(TerminationReason::Error)
250        );
251    }
252
253    #[test]
254    fn budget_ledger_delegates_to_scheduler_budget() {
255        let mut ledger = BudgetLedger::new(SchedulerBudget {
256            max_turns: 2,
257            ..SchedulerBudget::default()
258        });
259        assert_eq!(ledger.exceeded(None), None);
260        ledger.turns = 2;
261        assert_eq!(ledger.exceeded(None), Some("max_turns"));
262    }
263
264    #[test]
265    fn task_table_insert_and_lineage() {
266        let mut table = TaskTable::new();
267        table.insert(Tcb::root("root", SchedulerBudget::default()));
268        let mut child = Tcb::root("child", SchedulerBudget::default());
269        child.parent = Some("root".into());
270        table.insert(child);
271
272        assert_eq!(table.children_of("root").len(), 1);
273        assert!(table.get("root").is_some());
274    }
275
276    #[test]
277    fn task_table_insert_is_idempotent_by_id() {
278        let mut table = TaskTable::new();
279        table.insert(Tcb::root("root", SchedulerBudget::default()));
280        let mut updated = Tcb::root("root", SchedulerBudget::default());
281        updated.state = TaskLifecycle::Running;
282        table.insert(updated);
283
284        assert_eq!(table.all().len(), 1);
285        assert_eq!(table.get("root").unwrap().state, TaskLifecycle::Running);
286    }
287
288    #[test]
289    fn budget_verdict_none_within_budget() {
290        let tcb = Tcb::root(
291            "root",
292            SchedulerBudget {
293                max_turns: 5,
294                ..SchedulerBudget::default()
295            },
296        );
297        assert_eq!(budget_verdict(&tcb, None), None);
298    }
299
300    #[test]
301    fn budget_verdict_matches_should_terminate_axis() {
302        let limits = SchedulerBudget {
303            max_turns: 2,
304            ..SchedulerBudget::default()
305        };
306        let mut tcb = Tcb::root("root", limits.clone());
307        tcb.budget.turns = 2;
308        // budget_verdict and the underlying budget check must agree on verdict and reason.
309        assert_eq!(limits.should_terminate(2, 0, None, None), Some("max_turns"));
310        assert_eq!(
311            budget_verdict(&tcb, None),
312            Some(TerminationReason::MaxTurns)
313        );
314    }
315
316    #[test]
317    fn budget_verdict_wall_time_maps_to_timeout() {
318        let limits = SchedulerBudget {
319            max_wall_ms: Some(1_000),
320            ..SchedulerBudget::default()
321        };
322        let mut tcb = Tcb::root("root", limits);
323        tcb.budget.started_at_ms = Some(0);
324        assert_eq!(
325            budget_verdict(&tcb, Some(2_000)),
326            Some(TerminationReason::Timeout)
327        );
328    }
329
330    #[test]
331    fn baseline_token_budget_terminates() {
332        let limits = SchedulerBudget {
333            max_total_tokens: 100,
334            ..SchedulerBudget::default()
335        };
336        let mut tcb = Tcb::root("root", limits);
337        tcb.budget.total_tokens = 200;
338        assert_eq!(
339            budget_verdict(&tcb, None),
340            Some(TerminationReason::TokenBudget)
341        );
342    }
343}
deepstrike_core/scheduler/tcb.rs

deepstrike_core/scheduler/
tcb.rs
