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 { limits, turns: 0, total_tokens: 0, started_at_ms: None }
98    }
99
100    /// Delegates to the existing budget logic — single source of truth, no axis drift.
101    pub fn exceeded(&self, now_ms: Option<u64>) -> Option<&'static str> {
102        self.limits
103            .should_terminate(self.turns, self.total_tokens, now_ms, self.started_at_ms)
104    }
105}
106
107impl Default for BudgetLedger {
108    fn default() -> Self {
109        Self::new(SchedulerBudget::default())
110    }
111}
112
113/// Sub-agent-specific identity carried by a child [`Tcb`]; `None` on the root task.
114///
115/// This is what makes the `AgentProcess` view *derived* from the [`TaskTable`]: every child task
116/// whose `proc` is `Some` reconstructs exactly one [`crate::proc::AgentProcess`] (see
117/// [`crate::proc::AgentProcess::from_tcb`]).
118#[derive(Debug, Clone)]
119pub struct ProcInfo {
120    pub parent_session_id: CompactString,
121    pub role: AgentRole,
122    pub isolation: AgentIsolation,
123    pub context_inheritance: ContextInheritance,
124    /// The join result once the sub-agent has completed; `None` while running.
125    pub result: Option<SubAgentResult>,
126}
127
128/// One schedulable entity. The root loop and every sub-agent are uniform `Tcb`s.
129#[derive(Debug, Clone)]
130pub struct Tcb {
131    pub id: TaskId,
132    pub parent: Option<TaskId>,
133    pub state: TaskLifecycle,
134    pub budget: BudgetLedger,
135    pub wait: Option<WaitReason>,
136    /// Capability ids permitted to this task (mirrors `AgentProcess.permitted_capability_ids`).
137    pub caps: Vec<CompactString>,
138    /// Sub-agent identity for child tasks; `None` for the root loop.
139    pub proc: Option<ProcInfo>,
140}
141
142impl Tcb {
143    /// The root loop task. Constructed from the runtime task at `Start`.
144    pub fn root(id: impl Into<TaskId>, budget: SchedulerBudget) -> Self {
145        Self {
146            id: id.into(),
147            parent: None,
148            state: TaskLifecycle::Ready,
149            budget: BudgetLedger::new(budget),
150            wait: None,
151            caps: Vec::new(),
152            proc: None,
153        }
154    }
155
156    /// A sub-agent task spawned under the root, seeded `Running`, carrying the manifest's
157    /// process identity. The single source of truth for what the `AgentProcess` view exposes.
158    pub fn spawned(manifest: &IsolationManifest, budget: SchedulerBudget) -> Self {
159        Self {
160            id: manifest.agent_id.clone(),
161            parent: Some("root".into()),
162            state: TaskLifecycle::Running,
163            budget: BudgetLedger::new(budget),
164            wait: None,
165            caps: manifest.permitted_capability_ids.clone(),
166            proc: Some(ProcInfo {
167                parent_session_id: manifest.parent_session_id.clone(),
168                role: manifest.role,
169                isolation: manifest.isolation,
170                context_inheritance: manifest.context_inheritance,
171                result: None,
172            }),
173        }
174    }
175}
176
177/// Unified registry of all tasks: the root loop plus one child per sub-agent. The sole source of
178/// truth for schedulability and lineage; the `AgentProcess` view is derived from it.
179#[derive(Debug, Clone, Default)]
180pub struct TaskTable {
181    tasks: Vec<Tcb>,
182}
183
184impl TaskTable {
185    pub fn new() -> Self {
186        Self::default()
187    }
188
189    pub fn insert(&mut self, tcb: Tcb) {
190        if let Some(existing) = self.tasks.iter_mut().find(|t| t.id == tcb.id) {
191            *existing = tcb;
192        } else {
193            self.tasks.push(tcb);
194        }
195    }
196
197    pub fn get(&self, id: &str) -> Option<&Tcb> {
198        self.tasks.iter().find(|t| t.id.as_str() == id)
199    }
200
201    pub fn get_mut(&mut self, id: &str) -> Option<&mut Tcb> {
202        self.tasks.iter_mut().find(|t| t.id.as_str() == id)
203    }
204
205    pub fn all(&self) -> &[Tcb] {
206        &self.tasks
207    }
208
209    pub fn children_of(&self, parent: &str) -> Vec<&Tcb> {
210        self.tasks
211            .iter()
212            .filter(|t| t.parent.as_deref() == Some(parent))
213            .collect()
214    }
215}
216
217/// Pure budget verdict for one task: `Some(reason)` when a budget axis (turn/token/wall)
218/// is exhausted, mapped to the same `TerminationReason` the state machine applies.
219/// The single budget decision point — evaluated at each turn boundary.
220pub fn budget_verdict(task: &Tcb, now_ms: Option<u64>) -> Option<TerminationReason> {
221    task.budget.exceeded(now_ms).map(|axis| match axis {
222        "max_turns" => TerminationReason::MaxTurns,
223        "wall_time" => TerminationReason::Timeout,
224        _ => TerminationReason::TokenBudget,
225    })
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231
232    #[test]
233    fn process_state_maps_to_lifecycle() {
234        assert_eq!(TaskLifecycle::from(ProcessState::Running), TaskLifecycle::Running);
235        assert_eq!(
236            TaskLifecycle::from(ProcessState::Joined),
237            TaskLifecycle::Done(TerminationReason::Completed)
238        );
239        assert_eq!(
240            TaskLifecycle::from(ProcessState::Failed),
241            TaskLifecycle::Done(TerminationReason::Error)
242        );
243    }
244
245    #[test]
246    fn budget_ledger_delegates_to_scheduler_budget() {
247        let mut ledger = BudgetLedger::new(SchedulerBudget {
248            max_turns: 2,
249            ..SchedulerBudget::default()
250        });
251        assert_eq!(ledger.exceeded(None), None);
252        ledger.turns = 2;
253        assert_eq!(ledger.exceeded(None), Some("max_turns"));
254    }
255
256    #[test]
257    fn task_table_insert_and_lineage() {
258        let mut table = TaskTable::new();
259        table.insert(Tcb::root("root", SchedulerBudget::default()));
260        let mut child = Tcb::root("child", SchedulerBudget::default());
261        child.parent = Some("root".into());
262        table.insert(child);
263
264        assert_eq!(table.children_of("root").len(), 1);
265        assert!(table.get("root").is_some());
266    }
267
268    #[test]
269    fn task_table_insert_is_idempotent_by_id() {
270        let mut table = TaskTable::new();
271        table.insert(Tcb::root("root", SchedulerBudget::default()));
272        let mut updated = Tcb::root("root", SchedulerBudget::default());
273        updated.state = TaskLifecycle::Running;
274        table.insert(updated);
275
276        assert_eq!(table.all().len(), 1);
277        assert_eq!(table.get("root").unwrap().state, TaskLifecycle::Running);
278    }
279
280    #[test]
281    fn budget_verdict_none_within_budget() {
282        let tcb = Tcb::root("root", SchedulerBudget { max_turns: 5, ..SchedulerBudget::default() });
283        assert_eq!(budget_verdict(&tcb, None), None);
284    }
285
286    #[test]
287    fn budget_verdict_matches_should_terminate_axis() {
288        let limits = SchedulerBudget { max_turns: 2, ..SchedulerBudget::default() };
289        let mut tcb = Tcb::root("root", limits.clone());
290        tcb.budget.turns = 2;
291        // budget_verdict and the underlying budget check must agree on verdict and reason.
292        assert_eq!(limits.should_terminate(2, 0, None, None), Some("max_turns"));
293        assert_eq!(budget_verdict(&tcb, None), Some(TerminationReason::MaxTurns));
294    }
295
296    #[test]
297    fn budget_verdict_wall_time_maps_to_timeout() {
298        let limits = SchedulerBudget { max_wall_ms: Some(1_000), ..SchedulerBudget::default() };
299        let mut tcb = Tcb::root("root", limits);
300        tcb.budget.started_at_ms = Some(0);
301        assert_eq!(budget_verdict(&tcb, Some(2_000)), Some(TerminationReason::Timeout));
302    }
303
304    #[test]
305    fn baseline_token_budget_terminates() {
306        let limits = SchedulerBudget { max_total_tokens: 100, ..SchedulerBudget::default() };
307        let mut tcb = Tcb::root("root", limits);
308        tcb.budget.total_tokens = 200;
309        assert_eq!(budget_verdict(&tcb, None), Some(TerminationReason::TokenBudget));
310    }
311}
deepstrike_core/scheduler/tcb.rs

deepstrike_core/scheduler/
tcb.rs
