Skip to main content

helper_mas_learning_loop/
helper_mas_learning_loop.rs

1#[path = "support/utils.rs"]
2mod _utils;
3
4use _utils::{boxed_error, cleanup_run, create_client, new_run_id, print_summary, require};
5use mubit_sdk::{
6    CheckpointOptions, ListAgentsOptions, RecordOutcomeOptions, RegisterAgentOptions,
7    RememberOptions, SurfaceStrategiesOptions, TransportMode,
8};
9use serde_json::json;
10use std::error::Error;
11use std::time::Instant;
12
13#[tokio::main(flavor = "current_thread")]
14async fn main() -> Result<(), Box<dyn Error>> {
15    let name = "helper_mas_learning_loop";
16    let started = Instant::now();
17    let client = create_client().await?;
18    let run_id = new_run_id("helper_mas_learning_loop");
19    client.set_run_id(Some(run_id.clone()));
20    client.set_transport(TransportMode::Http);
21
22    let mut passed = true;
23    let mut detail = "validated helper MAS checkpoint/outcome/strategies flow".to_string();
24    let mut metrics = json!({});
25
26    let scenario = async {
27        let mut register = RegisterAgentOptions::new("planner");
28        register.run_id = Some(run_id.clone());
29        register.role = "planner".to_string();
30        register.read_scopes = vec!["rule".to_string(), "lesson".to_string(), "fact".to_string()];
31        register.write_scopes = vec!["lesson".to_string(), "trace".to_string()];
32        register.shared_memory_lanes = vec!["knowledge".to_string(), "history".to_string()];
33        client.register_agent(register).await?;
34
35        let mut remember = RememberOptions::new(
36            "If a retry succeeds after token signer rotation, record it as a reusable recovery pattern.",
37        );
38        remember.run_id = Some(run_id.clone());
39        remember.agent_id = Some("planner".to_string());
40        remember.intent = Some("lesson".to_string());
41        remember.lesson_type = Some("success".to_string());
42        remember.lesson_scope = Some("session".to_string());
43        remember.lesson_importance = Some("high".to_string());
44        client.remember(remember).await?;
45
46        let mut checkpoint = CheckpointOptions::new(
47            "Planner isolated the recovery path to token signer rotation.",
48        );
49        checkpoint.run_id = Some(run_id.clone());
50        checkpoint.label = Some("pre-compaction-1".to_string());
51        checkpoint.metadata = Some(json!({ "stage": "analysis" }));
52        checkpoint.agent_id = Some("planner".to_string());
53        let checkpoint_response = client.checkpoint(checkpoint).await?;
54
55        let mut list_agents = ListAgentsOptions::default();
56        list_agents.run_id = Some(run_id.clone());
57        let agents = client.list_agents(list_agents).await?;
58
59        let lessons = client.control.lessons(json!({ "run_id": run_id, "limit": 10 })).await?;
60        let lesson_id = lessons
61            .get("lessons")
62            .and_then(|value| value.as_array())
63            .and_then(|items| items.first())
64            .and_then(|item| item.get("id"))
65            .and_then(|value| value.as_str())
66            .ok_or_else(|| boxed_error(format!("expected at least one lesson to record outcome against: {lessons}")))?
67            .to_string();
68
69        let mut outcome = RecordOutcomeOptions::new(lesson_id, "success");
70        outcome.run_id = Some(run_id.clone());
71        outcome.signal = 0.75;
72        outcome.rationale = "Recovery succeeded after following the stored lesson.".to_string();
73        outcome.agent_id = Some("planner".to_string());
74        let outcome_response = client.record_outcome(outcome).await?;
75
76        let mut strategies = SurfaceStrategiesOptions::default();
77        strategies.run_id = Some(run_id.clone());
78        strategies.lesson_types = vec!["success".to_string(), "failure".to_string()];
79        strategies.max_strategies = 5;
80        let strategy_response = client.surface_strategies(strategies).await?;
81
82        require(
83            checkpoint_response
84                .get("success")
85                .and_then(|value| value.as_bool())
86                .unwrap_or(false),
87            format!("checkpoint failed: {checkpoint_response}"),
88        )?;
89        let agent_count = agents
90            .get("agents")
91            .and_then(|value| value.as_array())
92            .map(|items| items.len())
93            .unwrap_or(0);
94        require(agent_count == 1, format!("agent registration missing: {agents}"))?;
95        require(
96            outcome_response
97                .get("success")
98                .and_then(|value| value.as_bool())
99                .unwrap_or(false),
100            format!("record_outcome failed: {outcome_response}"),
101        )?;
102        let strategy_count = strategy_response
103            .get("strategies")
104            .and_then(|value| value.as_array())
105            .map(|items| items.len())
106            .unwrap_or(0);
107
108        metrics = json!({
109            "run_id": run_id,
110            "checkpoint_id": checkpoint_response.get("checkpoint_id").cloned().unwrap_or(serde_json::Value::Null),
111            "agent_count": agent_count,
112            "strategy_count": strategy_count,
113        });
114
115        Ok::<(), Box<dyn Error>>(())
116    }
117    .await;
118
119    if let Err(err) = scenario {
120        passed = false;
121        detail = err.to_string();
122    }
123
124    let cleanup_ok = cleanup_run(&client, &run_id).await;
125    if !cleanup_ok {
126        passed = false;
127        detail = format!("{detail} | cleanup failures");
128    }
129
130    print_summary(
131        name,
132        passed,
133        &detail,
134        &metrics,
135        started.elapsed().as_secs_f64(),
136        cleanup_ok,
137    );
138
139    if passed {
140        Ok(())
141    } else {
142        Err(boxed_error(detail))
143    }
144}