Skip to main content

robonix_pilot/
state_context.rs

1// SPDX-License-Identifier: MulanPSL-2.0
2
3use crate::pb::contracts::robonix_system_executor_execute_client::RobonixSystemExecutorExecuteClient;
4use crate::pb::executor::rtdl_event::RtdlEventEnum;
5use crate::pb::pilot::rtdl_node_state::RtdlNodeStateEnum;
6use crate::pb::pilot::{CapabilityCall, Plan, RtdlNode};
7use crate::planner::ExecutorConn;
8use robonix_atlas::client::AtlasClient;
9use robonix_atlas::pb as atlas_pb;
10use serde_json::{Value, json};
11use std::time::Duration;
12use tonic::Request;
13use tonic::transport::Channel;
14use uuid::Uuid;
15
16const RTDL_DO: u32 = 2;
17const SCENE_CONTEXT: &str = "robonix/system/scene/get_robot_context";
18
19pub async fn collect(
20    executor: &ExecutorConn,
21    atlas: &mut AtlasClient,
22    caps: &[(String, atlas_pb::Capability)],
23) -> String {
24    let scene_target = caps
25        .iter()
26        .find(|(_, capability)| capability.contract_id == SCENE_CONTEXT)
27        .map(|(provider_id, _)| provider_id.clone());
28    let scene = match scene_target {
29        Some(provider_id) => query_scene(executor.graph.clone(), provider_id).await,
30        None => json!({"available": false, "error": "Scene context contract is not registered"}),
31    };
32    let providers = match atlas
33        .query_capabilities("", "", atlas_pb::Transport::Unspecified)
34        .await
35    {
36        Ok(rows) => rows
37            .into_iter()
38            .filter(|provider| {
39                provider.capabilities.iter().any(|capability| {
40                    let id = capability.contract_id.as_str();
41                    id.contains("/navigation/")
42                        || id.contains("/scene/")
43                        || id.contains("/arm/")
44                        || id.contains("/chassis/")
45                        || id.contains("/skill/")
46                })
47            })
48            .map(|provider| {
49                json!({
50                    "provider_id": provider.id,
51                    "state": lifecycle_name(provider.state),
52                    "detail": provider.state_detail,
53                })
54            })
55            .collect::<Vec<_>>(),
56        Err(error) => vec![json!({"available": false, "error": error.to_string()})],
57    };
58    format!(
59        "\n\n## Current environment and provider state\n\
60         Refreshed immediately before this planning round. Scene owns map pose, \
61         room membership, areas, and nearby objects. Atlas state only indicates \
62         provider availability; current task progress remains in the in-flight RTDL \
63         tree block. Missing or stale state means unknown.\n\n{}\n",
64        serde_json::to_string_pretty(&json!({
65            "scene": scene,
66            "provider_availability": providers,
67        }))
68        .unwrap_or_else(|_| "{}".into())
69    )
70}
71
72async fn query_scene(
73    mut graph: RobonixSystemExecutorExecuteClient<Channel>,
74    provider_id: String,
75) -> Value {
76    let plan_id = format!("state-prefetch-{}", Uuid::new_v4());
77    let plan = Plan {
78        plan_id: plan_id.clone(),
79        session_id: "pilot-state-prefetch".into(),
80        round: 0,
81        nodes: vec![RtdlNode {
82            node_kind: RTDL_DO,
83            children: Vec::new(),
84            call: Some(CapabilityCall {
85                call_id: format!("{plan_id}:0"),
86                provider_id,
87                contract_id: SCENE_CONTEXT.into(),
88                args_json: "{}".into(),
89            }),
90            op_id: "scene_snapshot".into(),
91            description: "Read Scene spatial context before planning".into(),
92        }],
93        root_index: 0,
94    };
95    let query = async {
96        let mut stream = graph
97            .execute(Request::new(plan))
98            .await
99            .map_err(|error| error.to_string())?
100            .into_inner();
101        while let Some(event) = stream.message().await.map_err(|error| error.to_string())? {
102            if event.event_kind != RtdlEventEnum::NodeState as u32 {
103                continue;
104            }
105            let Some(state) = event.node_state else {
106                continue;
107            };
108            if state.state == RtdlNodeStateEnum::Succeeded as u32 {
109                let output = state
110                    .leaf_result
111                    .map(|result| result.output)
112                    .unwrap_or(state.operator_detail);
113                return Ok(serde_json::from_str(&output).unwrap_or_else(|_| json!({"raw": output})));
114            }
115            if matches!(
116                RtdlNodeStateEnum::try_from(state.state as i32),
117                Ok(RtdlNodeStateEnum::Failed
118                    | RtdlNodeStateEnum::Canceled
119                    | RtdlNodeStateEnum::Timeout)
120            ) {
121                return Err(state.operator_detail);
122            }
123        }
124        Err("Scene query ended without a terminal result".into())
125    };
126    match tokio::time::timeout(Duration::from_secs(3), query).await {
127        Ok(Ok(value)) => json!({"available": true, "state": value}),
128        Ok(Err(error)) => json!({"available": false, "error": error}),
129        Err(_) => json!({"available": false, "error": "Scene query timed out after 3 seconds"}),
130    }
131}
132
133fn lifecycle_name(state: i32) -> &'static str {
134    match atlas_pb::LifecycleState::try_from(state) {
135        Ok(atlas_pb::LifecycleState::StateRegistered) => "registered",
136        Ok(atlas_pb::LifecycleState::StateInactive) => "inactive",
137        Ok(atlas_pb::LifecycleState::StateActive) => "active",
138        Ok(atlas_pb::LifecycleState::StateError) => "error",
139        Ok(atlas_pb::LifecycleState::StateTerminated) => "terminated",
140        _ => "unknown",
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::lifecycle_name;
147    use robonix_atlas::pb::LifecycleState;
148
149    #[test]
150    fn lifecycle_labels_are_stable() {
151        assert_eq!(lifecycle_name(LifecycleState::StateActive as i32), "active");
152        assert_eq!(lifecycle_name(LifecycleState::StateError as i32), "error");
153    }
154}