ai_agents_runtime/optimization/
branch.rs

1use ai_agents_core::{AgentError, Result};
2use ai_agents_reasoning::ReasoningMode;
3use uuid::Uuid;
4
5use super::maintenance::RuntimeTaskPurpose;
6use super::response::MainResponseDraft;
7use super::skill::SkillCandidate;
8use super::turn::TransitionCandidate;
9
10/// Runtime branch category used for observation and winner selection.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12pub enum RuntimeOptimizationKind {
13    ParallelStateTransition,
14    SpeculativeSkillRouting,
15    SpeculativeReasoningAuto,
16    BufferedStreamingRouting,
17}
18
19impl RuntimeOptimizationKind {
20    pub fn as_label(self) -> &'static str {
21        match self {
22            Self::ParallelStateTransition => "parallel_state_transition",
23            Self::SpeculativeSkillRouting => "speculative_skill_routing",
24            Self::SpeculativeReasoningAuto => "speculative_reasoning_auto",
25            Self::BufferedStreamingRouting => "buffered_streaming_routing",
26        }
27    }
28}
29
30/// Describes what can commit if a branch wins.
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
32pub enum RuntimeCommitBehavior {
33    FinalResponse,
34    TransitionDecision,
35    SkillSelection,
36    ReasoningDecision,
37    DiscardOnly,
38}
39
40impl RuntimeCommitBehavior {
41    pub fn as_label(self) -> &'static str {
42        match self {
43            Self::FinalResponse => "final_response",
44            Self::TransitionDecision => "transition_decision",
45            Self::SkillSelection => "skill_selection",
46            Self::ReasoningDecision => "reasoning_decision",
47            Self::DiscardOnly => "discard_only",
48        }
49    }
50}
51
52/// Lifecycle state for one runtime branch.
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
54pub enum RuntimeBranchStatus {
55    Scheduled,
56    Completed,
57    Committed,
58    Discarded,
59    Failed,
60    Cancelled,
61}
62
63impl RuntimeBranchStatus {
64    pub fn as_label(self) -> &'static str {
65        match self {
66            Self::Scheduled => "scheduled",
67            Self::Completed => "completed",
68            Self::Committed => "committed",
69            Self::Discarded => "discarded",
70            Self::Failed => "failed",
71            Self::Cancelled => "cancelled",
72        }
73    }
74
75    pub fn is_terminal(self) -> bool {
76        matches!(
77            self,
78            Self::Committed | Self::Discarded | Self::Failed | Self::Cancelled
79        )
80    }
81}
82
83/// Priority hint used when several branch results are available.
84#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
85pub enum RuntimeTaskPriority {
86    Low,
87    Normal,
88    High,
89    Critical,
90}
91
92/// Result produced by one speculative branch before commit or discard.
93#[derive(Debug)]
94pub enum RuntimeBranchResult {
95    MainDraft(MainResponseDraft),
96    Transition(Option<TransitionCandidate>),
97    Skill(Option<SkillCandidate>),
98    Reasoning(ReasoningMode),
99    Failed(AgentError),
100    Cancelled,
101}
102
103/// Metadata and result for one completed runtime branch.
104#[derive(Debug)]
105pub struct RuntimeBranchOutcome {
106    pub branch: RuntimeBranch,
107    pub result: RuntimeBranchResult,
108}
109
110/// Metadata for one branch scheduled during a turn.
111#[derive(Debug, Clone)]
112pub struct RuntimeBranch {
113    pub id: Uuid,
114    pub purpose: RuntimeTaskPurpose,
115    pub optimization: RuntimeOptimizationKind,
116    pub priority: RuntimeTaskPriority,
117    pub commit_behavior: RuntimeCommitBehavior,
118    pub status: RuntimeBranchStatus,
119}
120
121impl RuntimeBranch {
122    pub fn new(
123        purpose: RuntimeTaskPurpose,
124        optimization: RuntimeOptimizationKind,
125        priority: RuntimeTaskPriority,
126        commit_behavior: RuntimeCommitBehavior,
127    ) -> Self {
128        Self {
129            id: Uuid::new_v4(),
130            purpose,
131            optimization,
132            priority,
133            commit_behavior,
134            status: RuntimeBranchStatus::Scheduled,
135        }
136    }
137
138    pub fn transition_to(&mut self, next: RuntimeBranchStatus) -> Result<()> {
139        if self.status.is_terminal() && self.status != next {
140            return Err(AgentError::Other(format!(
141                "runtime branch {} cannot move from terminal status {} to {}",
142                self.id,
143                self.status.as_label(),
144                next.as_label()
145            )));
146        }
147        self.status = next;
148        Ok(())
149    }
150
151    pub fn complete(&mut self) -> Result<()> {
152        self.transition_to(RuntimeBranchStatus::Completed)
153    }
154
155    pub fn commit(&mut self) -> Result<()> {
156        self.transition_to(RuntimeBranchStatus::Committed)
157    }
158
159    pub fn discard(&mut self) -> Result<()> {
160        self.transition_to(RuntimeBranchStatus::Discarded)
161    }
162
163    pub fn fail(&mut self) -> Result<()> {
164        self.transition_to(RuntimeBranchStatus::Failed)
165    }
166
167    pub fn cancel(&mut self) -> Result<()> {
168        self.transition_to(RuntimeBranchStatus::Cancelled)
169    }
170
171    pub fn branch_id(&self) -> String {
172        self.id.to_string()
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    #[test]
181    fn terminal_branch_status_cannot_change() {
182        let mut branch = RuntimeBranch::new(
183            RuntimeTaskPurpose::MainResponse,
184            RuntimeOptimizationKind::SpeculativeSkillRouting,
185            RuntimeTaskPriority::Normal,
186            RuntimeCommitBehavior::FinalResponse,
187        );
188        branch
189            .transition_to(RuntimeBranchStatus::Committed)
190            .unwrap();
191        assert!(
192            branch
193                .transition_to(RuntimeBranchStatus::Discarded)
194                .is_err()
195        );
196    }
197
198    #[test]
199    fn repeated_terminal_status_is_allowed() {
200        let mut branch = RuntimeBranch::new(
201            RuntimeTaskPurpose::MainResponse,
202            RuntimeOptimizationKind::SpeculativeSkillRouting,
203            RuntimeTaskPriority::Normal,
204            RuntimeCommitBehavior::FinalResponse,
205        );
206        branch
207            .transition_to(RuntimeBranchStatus::Discarded)
208            .unwrap();
209        assert!(branch.transition_to(RuntimeBranchStatus::Discarded).is_ok());
210    }
211}
ai_agents_runtime/optimization/branch.rs

ai_agents_runtime/optimization/
branch.rs
