ai_agents_runtime/optimization/
branch.rs1use 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#[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#[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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
85pub enum RuntimeTaskPriority {
86 Low,
87 Normal,
88 High,
89 Critical,
90}
91
92#[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#[derive(Debug)]
107pub struct RuntimeBranchOutcome {
108 pub branch: RuntimeBranch,
109 pub result: RuntimeBranchResult,
110}
111
112#[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}