Skip to main content

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/// Boxing the draft would break frozen public v1 variant construction, so this enum keeps its current representation.
94#[derive(Debug)]
95#[allow(clippy::large_enum_variant)]
96pub enum RuntimeBranchResult {
97    MainDraft(MainResponseDraft),
98    Transition(Option<TransitionCandidate>),
99    Skill(Option<SkillCandidate>),
100    Reasoning(ReasoningMode),
101    Failed(AgentError),
102    Cancelled,
103}
104
105/// Metadata and result for one completed runtime branch.
106#[derive(Debug)]
107pub struct RuntimeBranchOutcome {
108    pub branch: RuntimeBranch,
109    pub result: RuntimeBranchResult,
110}
111
112/// Metadata for one branch scheduled during a turn.
113#[derive(Debug, Clone)]
114pub struct RuntimeBranch {
115    pub id: Uuid,
116    pub purpose: RuntimeTaskPurpose,
117    pub optimization: RuntimeOptimizationKind,
118    pub priority: RuntimeTaskPriority,
119    pub commit_behavior: RuntimeCommitBehavior,
120    pub status: RuntimeBranchStatus,
121}
122
123impl RuntimeBranch {
124    pub fn new(
125        purpose: RuntimeTaskPurpose,
126        optimization: RuntimeOptimizationKind,
127        priority: RuntimeTaskPriority,
128        commit_behavior: RuntimeCommitBehavior,
129    ) -> Self {
130        Self {
131            id: Uuid::new_v4(),
132            purpose,
133            optimization,
134            priority,
135            commit_behavior,
136            status: RuntimeBranchStatus::Scheduled,
137        }
138    }
139
140    pub fn transition_to(&mut self, next: RuntimeBranchStatus) -> Result<()> {
141        if self.status.is_terminal() && self.status != next {
142            return Err(AgentError::Other(format!(
143                "runtime branch {} cannot move from terminal status {} to {}",
144                self.id,
145                self.status.as_label(),
146                next.as_label()
147            )));
148        }
149        self.status = next;
150        Ok(())
151    }
152
153    pub fn complete(&mut self) -> Result<()> {
154        self.transition_to(RuntimeBranchStatus::Completed)
155    }
156
157    pub fn commit(&mut self) -> Result<()> {
158        self.transition_to(RuntimeBranchStatus::Committed)
159    }
160
161    pub fn discard(&mut self) -> Result<()> {
162        self.transition_to(RuntimeBranchStatus::Discarded)
163    }
164
165    pub fn fail(&mut self) -> Result<()> {
166        self.transition_to(RuntimeBranchStatus::Failed)
167    }
168
169    pub fn cancel(&mut self) -> Result<()> {
170        self.transition_to(RuntimeBranchStatus::Cancelled)
171    }
172
173    pub fn branch_id(&self) -> String {
174        self.id.to_string()
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181
182    #[test]
183    fn terminal_branch_status_cannot_change() {
184        let mut branch = RuntimeBranch::new(
185            RuntimeTaskPurpose::MainResponse,
186            RuntimeOptimizationKind::SpeculativeSkillRouting,
187            RuntimeTaskPriority::Normal,
188            RuntimeCommitBehavior::FinalResponse,
189        );
190        branch
191            .transition_to(RuntimeBranchStatus::Committed)
192            .unwrap();
193        assert!(
194            branch
195                .transition_to(RuntimeBranchStatus::Discarded)
196                .is_err()
197        );
198    }
199
200    #[test]
201    fn repeated_terminal_status_is_allowed() {
202        let mut branch = RuntimeBranch::new(
203            RuntimeTaskPurpose::MainResponse,
204            RuntimeOptimizationKind::SpeculativeSkillRouting,
205            RuntimeTaskPriority::Normal,
206            RuntimeCommitBehavior::FinalResponse,
207        );
208        branch
209            .transition_to(RuntimeBranchStatus::Discarded)
210            .unwrap();
211        assert!(branch.transition_to(RuntimeBranchStatus::Discarded).is_ok());
212    }
213}