Skip to main content

robonix_pilot/
memory.rs

1// SPDX-License-Identifier: MulanPSL-2.0
2// Long-term-memory dispatch (best-effort, fire-and-forget on errors).
3//
4// `prefetch` runs before the first VLM call, `compact` on session_end.
5// Both build a one-call Plan and hand it to executor; the corresponding
6// memory provider is looked up via atlas the same way every other capability
7// is. Missing providers are silently tolerated — memory is never load-bearing.
8
9use 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
47/// Dispatch one `search_memory` call and return the result text. Returns
48/// `None` if the provider is not registered, the index is empty, or any error
49/// occurs.
50pub 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
92/// Best-effort `compact_memory` on session teardown. Logs failures, never
93/// propagates errors (the provider may be absent entirely).
94pub 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
157/// Convert memory helper node events into the shared result record. `do` nodes
158/// carry a concrete capability result; operator detail is only a fallback.
159fn 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}