1use crate::discovery;
10use crate::pb::executor::rtdl_event::RtdlEventEnum;
11use crate::pb::pilot::rtdl_node_state::RtdlNodeStateEnum;
12use crate::pb::pilot::{CapabilityCall, CapabilityCallResult, Plan, RtdlNode};
13use crate::planner::ExecutorConn;
14use robonix_atlas::client::AtlasClient;
15use robonix_scribe::debug;
16use tonic::Request;
17use uuid::Uuid;
18
19const RTDL_SEQUENCE: u32 = 0;
20const RTDL_DO: u32 = 2;
21
22fn single_call_plan(plan_id: String, session_id: String, round: u32, call: CapabilityCall) -> Plan {
23 Plan {
24 plan_id,
25 session_id,
26 round,
27 nodes: vec![
28 RtdlNode {
29 node_kind: RTDL_SEQUENCE,
30 children: vec![1],
31 call: None,
32 op_id: "memory_prefetch_sequence".to_string(),
33 description: "Run memory prefetch before planning".to_string(),
34 },
35 RtdlNode {
36 node_kind: RTDL_DO,
37 children: Vec::new(),
38 call: Some(call),
39 op_id: "memory_prefetch_search".to_string(),
40 description: "Search memory for context relevant to the user request".to_string(),
41 },
42 ],
43 root_index: 0,
44 }
45}
46
47pub async fn prefetch(
51 query: &str,
52 executor: &mut ExecutorConn,
53 target: Option<(String, String)>,
54) -> Option<String> {
55 let (provider_id, contract_id) = target?;
56 let plan = single_call_plan(
57 Uuid::new_v4().to_string(),
58 "memory-prefetch".to_string(),
59 0,
60 CapabilityCall {
61 call_id: Uuid::new_v4().to_string(),
62 provider_id,
63 contract_id,
64 args_json: serde_json::json!({ "data": query }).to_string(),
65 },
66 );
67
68 let submitted_plan = plan.clone();
69 let mut stream = executor
70 .graph
71 .execute(Request::new(plan))
72 .await
73 .ok()?
74 .into_inner();
75 while let Ok(Some(event)) = stream.message().await {
76 if event.event_kind == RtdlEventEnum::NodeState as u32
77 && let Some(ns) = event.node_state
78 && is_terminal_executor_state(ns.state)
79 {
80 let r = executor_node_state_to_result(&submitted_plan, ns);
81 let out = r.output;
82 if r.success && !out.contains("No relevant memories") && !out.is_empty() {
83 debug!("[pilot] memory prefetch: {out}");
84 return Some(out);
85 }
86 return None;
87 }
88 }
89 None
90}
91
92pub async fn try_compact(executor: &mut ExecutorConn, atlas: &mut AtlasClient, _consumer_id: &str) {
95 let providers = match discovery::discover(atlas).await {
96 Ok(c) => c,
97 Err(e) => {
98 debug!("[pilot] compact_memory: discovery failed: {e}");
99 return;
100 }
101 };
102 let Some((provider_id, cap)) = providers
103 .iter()
104 .find(|(_, cap)| cap.contract_id == "robonix/service/memory/compact")
105 else {
106 return;
107 };
108
109 let plan = single_call_plan(
110 Uuid::new_v4().to_string(),
111 "memory-compact".to_string(),
112 0,
113 CapabilityCall {
114 call_id: Uuid::new_v4().to_string(),
115 provider_id: provider_id.clone(),
116 contract_id: cap.contract_id.clone(),
117 args_json: "{}".to_string(),
118 },
119 );
120
121 let submitted_plan = plan.clone();
122 let Ok(mut stream) = executor
123 .graph
124 .execute(Request::new(plan))
125 .await
126 .map(|r| r.into_inner())
127 else {
128 return;
129 };
130 while let Ok(Some(event)) = stream.message().await {
131 if event.event_kind == RtdlEventEnum::NodeState as u32
132 && let Some(ns) = event.node_state
133 && is_terminal_executor_state(ns.state)
134 {
135 let r = executor_node_state_to_result(&submitted_plan, ns);
136 let out = r.output;
137 if r.success {
138 debug!("[pilot] compact_memory: {out}");
139 } else {
140 debug!("[pilot] compact_memory failed: {out}");
141 }
142 return;
143 }
144 }
145}
146
147fn is_terminal_executor_state(state: u32) -> bool {
148 matches!(
149 RtdlNodeStateEnum::try_from(state as i32),
150 Ok(RtdlNodeStateEnum::Succeeded
151 | RtdlNodeStateEnum::Failed
152 | RtdlNodeStateEnum::Canceled
153 | RtdlNodeStateEnum::Timeout)
154 )
155}
156
157fn executor_node_state_to_result(
160 plan: &Plan,
161 ns: crate::pb::pilot::RtdlNodeState,
162) -> CapabilityCallResult {
163 if let Some(result) = ns.leaf_result {
164 return result;
165 }
166 let call = plan
167 .nodes
168 .get(ns.node_index as usize)
169 .and_then(|node| node.call.as_ref());
170 let success = ns.state == RtdlNodeStateEnum::Succeeded as u32;
171 CapabilityCallResult {
172 call_id: call.map(|c| c.call_id.clone()).unwrap_or_default(),
173 provider_id: call.map(|c| c.provider_id.clone()).unwrap_or_default(),
174 contract_id: call.map(|c| c.contract_id.clone()).unwrap_or_default(),
175 success,
176 output: ns.operator_detail.clone(),
177 error: if success {
178 String::new()
179 } else {
180 ns.operator_detail
181 },
182 }
183}