1use std::collections::BTreeMap;
14
15use serde::Serialize;
16
17use crate::residual::{CorrectionDirection, ResidualEvent};
18use crate::stability::StabilityParameters;
19
20pub type Evidence = BTreeMap<String, String>;
22
23pub type CorrectionDirectionSet = Vec<CorrectionDirection>;
25
26#[derive(Debug, Clone, PartialEq)]
28pub struct AgentBarrierResult {
29 pub ok: bool,
30 pub score: f64,
32 pub residuals: Vec<ResidualEvent>,
33 pub feedback: CorrectionDirectionSet,
34 pub evidence: Evidence,
35}
36
37impl AgentBarrierResult {
38 pub fn new(ok: bool, score: f64) -> Self {
39 Self {
40 ok,
41 score,
42 residuals: Vec::new(),
43 feedback: Vec::new(),
44 evidence: Evidence::new(),
45 }
46 }
47
48 pub fn with_residuals(mut self, residuals: Vec<ResidualEvent>) -> Self {
49 self.residuals = residuals;
50 self
51 }
52
53 pub fn with_feedback(mut self, feedback: CorrectionDirectionSet) -> Self {
54 self.feedback = feedback;
55 self
56 }
57
58 pub fn into_srbn(
62 self,
63 name: impl Into<String>,
64 ) -> srbn::SrbnResult<srbn::BarrierResult<CorrectionDirectionSet, Evidence>> {
65 let feedback = self
66 .feedback
67 .first()
68 .map(|d| d.instruction.clone())
69 .unwrap_or_default();
70 let correction = if self.feedback.is_empty() {
71 None
72 } else {
73 Some(self.feedback)
74 };
75 srbn::BarrierResult::new(
76 name,
77 self.ok,
78 self.score,
79 feedback,
80 correction,
81 self.evidence,
82 )
83 }
84}
85
86#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
89#[serde(rename_all = "snake_case")]
90pub enum AgentStabilizationStatus {
91 Stable,
92 Descended,
93 Stopped,
94 Exhausted,
95 Degraded,
96 ReplanRequired,
97}
98
99impl From<srbn::Status> for AgentStabilizationStatus {
100 fn from(status: srbn::Status) -> Self {
101 match status {
102 srbn::Status::Stable => AgentStabilizationStatus::Stable,
103 srbn::Status::Stopped => AgentStabilizationStatus::Stopped,
104 srbn::Status::Exhausted => AgentStabilizationStatus::Exhausted,
105 }
106 }
107}
108
109pub fn policy(params: &StabilityParameters, max_attempts: usize) -> srbn::Policy {
116 srbn::Policy {
117 max_attempts: max_attempts.max(1),
118 score_tolerance: params.energy_tolerance,
119 require_descent: true,
120 on_no_descent: srbn::OnNoDescent::Stop,
121 min_descent: params.rho_gate,
122 }
123}
124
125pub fn stabilize<State, B, U>(
131 initial: State,
132 mut barrier: B,
133 updater: U,
134 params: &StabilityParameters,
135 max_attempts: usize,
136) -> crate::error::Result<srbn::StabilizationResult<State, CorrectionDirectionSet, Evidence>>
137where
138 State: Clone,
139 B: FnMut(&State) -> AgentBarrierResult,
140 U: FnMut(
141 State,
142 &srbn::BarrierResult<CorrectionDirectionSet, Evidence>,
143 ) -> srbn::SrbnResult<State>,
144{
145 let srbn_barrier = move |state: &State| barrier(state).into_srbn("agent-barrier");
146 let result = srbn::stabilize(initial, srbn_barrier, updater, policy(params, max_attempts))?;
147 Ok(result)
148}
149
150pub fn trace_to_json<State>(
153 result: &srbn::StabilizationResult<State, CorrectionDirectionSet, Evidence>,
154) -> crate::error::Result<String>
155where
156 State: Serialize,
157{
158 srbn_serde::stabilization_result_json(result)
159 .map_err(|e| crate::error::SdkError::Kernel(format!("trace serialization failed: {e}")))
160}
161
162#[cfg(test)]
163mod tests {
164 use super::*;
165
166 #[test]
167 fn maps_kernel_status() {
168 assert_eq!(
169 AgentStabilizationStatus::from(srbn::Status::Stable),
170 AgentStabilizationStatus::Stable
171 );
172 assert_eq!(
173 AgentStabilizationStatus::from(srbn::Status::Stopped),
174 AgentStabilizationStatus::Stopped
175 );
176 assert_eq!(
177 AgentStabilizationStatus::from(srbn::Status::Exhausted),
178 AgentStabilizationStatus::Exhausted
179 );
180 }
181
182 #[test]
183 fn stabilizes_descending_energy_through_kernel() {
184 let params = StabilityParameters::measured(0.5, 0.0);
187 let barrier = |state: &i64| {
188 let v = (*state as f64) * (*state as f64);
189 AgentBarrierResult::new(*state == 0, v)
190 };
191 let updater = |state: i64, _b: &srbn::BarrierResult<CorrectionDirectionSet, Evidence>| {
192 Ok(state - state.signum())
193 };
194 let result = stabilize(3, barrier, updater, ¶ms, 10).unwrap();
195 assert_eq!(result.status, srbn::Status::Stable);
196 assert_eq!(result.state, 0);
197
198 let json = trace_to_json(&result).unwrap();
200 assert!(json.contains("\"status\":\"stable\""));
201 assert!(json.contains("\"attempts\""));
202 }
203
204 #[test]
205 fn rejects_non_finite_score_at_barrier() {
206 let params = StabilityParameters::measured(0.5, 0.0);
207 let barrier = |_state: &i64| AgentBarrierResult::new(false, f64::NAN);
208 let updater =
209 |state: i64, _b: &srbn::BarrierResult<CorrectionDirectionSet, Evidence>| Ok(state);
210 assert!(stabilize(1, barrier, updater, ¶ms, 3).is_err());
212 }
213}