robonix_executor/dispatch/
async_poll.rs

1// SPDX-License-Identifier: MulanPSL-2.0
2// Poll async MCP capabilities via `<contract_id>/status` until terminal state.
3
4use robonix_scribe::warn;
5use std::time::Duration;
6
7use crate::dispatch::{self, async_registry::AsyncGroup};
8use crate::pb::executor::RtdlEvent;
9use crate::pb::pilot::rtdl_node_state::RtdlNodeStateEnum;
10use crate::pb::pilot::{CapabilityCall, CapabilityCallResult};
11use crate::plan_runtime::{PlanRuntime, RunningAsyncCall};
12use crate::rtdl_wire::{self, NodeEventContext};
13use robonix_atlas::client::AtlasClient;
14use robonix_atlas::pb as atlas_pb;
15use tokio::sync::mpsc::Sender;
16use tokio::time;
17
18const POLL_INTERVAL: Duration = Duration::from_secs(2);
19
20/// Dispatch an async cap, poll status every 2s, stream node_state changes, return terminal result.
21pub async fn run_until_terminal(
22    call: &CapabilityCall,
23    group: &AsyncGroup,
24    self_provider_id: &str,
25    atlas: &mut AtlasClient,
26    tx: &Sender<Result<RtdlEvent, tonic::Status>>,
27    node: &NodeEventContext,
28    runtime: &PlanRuntime,
29) -> CapabilityCallResult {
30    let initial = dispatch::dispatch(call, self_provider_id, atlas, runtime, &node.plan_id).await;
31    if !initial.success {
32        let _ = tx
33            .send(Ok(rtdl_wire::node_state_from_result(
34                node,
35                initial.clone(),
36                RtdlNodeStateEnum::Failed as u32,
37            )))
38            .await;
39        return initial;
40    }
41
42    let run_id = extract_run_id(&initial.output);
43    let accepted = runtime
44        .register_or_cancel_async_call(
45            &node.plan_id,
46            RunningAsyncCall {
47                call_id: call.call_id.clone(),
48                provider_id: call.provider_id.clone(),
49                cancel_contract: group.cancel_contract.clone(),
50                run_id: run_id.clone(),
51            },
52            self_provider_id,
53            atlas,
54        )
55        .await;
56    if !accepted {
57        let result = canceled_result(call, "plan was cancelled before async polling began");
58        let _ = tx
59            .send(Ok(rtdl_wire::node_state_from_result(
60                node,
61                result.clone(),
62                RtdlNodeStateEnum::Canceled as u32,
63            )))
64            .await;
65        return result;
66    }
67
68    let mut interval = time::interval(POLL_INTERVAL);
69    interval.tick().await;
70
71    loop {
72        interval.tick().await;
73        if runtime.is_cancelled(&node.plan_id).await {
74            runtime
75                .cancel_async_call_for_plan(&node.plan_id, &call.call_id, self_provider_id, atlas)
76                .await;
77            let result = canceled_result(call, "plan was cancelled");
78            let _ = tx
79                .send(Ok(rtdl_wire::node_state_from_result(
80                    node,
81                    result.clone(),
82                    RtdlNodeStateEnum::Canceled as u32,
83                )))
84                .await;
85            return result;
86        }
87        let status_out = match poll_status(
88            self_provider_id,
89            &call.provider_id,
90            &group.status_contract,
91            &run_id,
92            atlas,
93        )
94        .await
95        {
96            Ok(s) => s,
97            Err(e) => {
98                let error = format!("status poll failed for {}: {e:#}", call.contract_id);
99                warn!("[executor] {error}");
100                let result = failed_result(call, &error);
101                runtime
102                    .unregister_async_call(&node.plan_id, &call.call_id)
103                    .await;
104                let _ = tx
105                    .send(Ok(rtdl_wire::node_state_from_result(
106                        node,
107                        result.clone(),
108                        RtdlNodeStateEnum::Failed as u32,
109                    )))
110                    .await;
111                return result;
112            }
113        };
114
115        let (state, detail) = parse_status_json(&status_out);
116        // Record live state (RUNNING/PENDING/PAUSED or terminal) so
117        // get_plan_status reflects async progress between polls.
118        runtime
119            .record_op_state(&node.plan_id, &node.op_id, state)
120            .await;
121        if rtdl_wire::is_terminal_state(state) {
122            let result = terminal_result(call, state, &detail, &status_out);
123            runtime
124                .unregister_async_call(&node.plan_id, &call.call_id)
125                .await;
126            let _ = tx
127                .send(Ok(rtdl_wire::node_state_from_result(
128                    node,
129                    result.clone(),
130                    state,
131                )))
132                .await;
133            return result;
134        }
135    }
136}
137
138async fn poll_status(
139    consumer_id: &str,
140    provider_id: &str,
141    status_contract: &str,
142    run_id: &str,
143    atlas: &mut AtlasClient,
144) -> anyhow::Result<String> {
145    let args = if run_id.is_empty() {
146        "{}".to_string()
147    } else {
148        serde_json::json!({ "run_id": run_id }).to_string()
149    };
150    let status_call = CapabilityCall {
151        call_id: format!("status-{}", uuid::Uuid::new_v4()),
152        provider_id: provider_id.to_string(),
153        contract_id: status_contract.to_string(),
154        args_json: args,
155    };
156    let (channel_id, endpoint, _) = atlas
157        .connect_capability(
158            consumer_id,
159            provider_id,
160            status_contract,
161            atlas_pb::Transport::Mcp,
162        )
163        .await?;
164    let result = crate::dispatch::mcp::execute(&status_call, &endpoint).await;
165    let _ = atlas.disconnect_capability(&channel_id).await;
166    if result.success {
167        Ok(result.output)
168    } else {
169        anyhow::bail!("{}", result.error)
170    }
171}
172
173/// Read `run_id` from the async cap's initial MCP response JSON.
174pub fn extract_run_id(output: &str) -> String {
175    let Ok(v) = serde_json::from_str::<serde_json::Value>(output) else {
176        return String::new();
177    };
178    v.get("run_id")
179        .and_then(|x| x.as_str())
180        .unwrap_or_default()
181        .to_string()
182}
183
184/// Parse status MCP JSON: requires uppercase `state` enum; optional `detail` string.
185pub fn parse_status_json(output: &str) -> (u32, String) {
186    let Ok(v) = serde_json::from_str::<serde_json::Value>(output) else {
187        warn!("[executor] status response is not valid JSON: {output}");
188        return (RtdlNodeStateEnum::Running as u32, output.to_string());
189    };
190
191    let Some(state_str) = v.get("state").and_then(|s| s.as_str()) else {
192        let error = format!("status response missing required 'state' field: {output}");
193        warn!("[executor] {error}");
194        return (RtdlNodeStateEnum::Failed as u32, error);
195    };
196
197    let detail = v
198        .get("detail")
199        .and_then(|x| x.as_str())
200        .unwrap_or_default()
201        .to_string();
202
203    match parse_state_name(state_str) {
204        Some(state) => (state, detail),
205        None => {
206            let error = format!("status response has unknown state '{state_str}': {output}");
207            warn!("[executor] {error}");
208            (RtdlNodeStateEnum::Failed as u32, error)
209        }
210    }
211}
212
213/// Convert status response state names into RTDL node state constants.
214pub fn parse_state_name(s: &str) -> Option<u32> {
215    match s.to_uppercase().as_str() {
216        "PENDING" => Some(RtdlNodeStateEnum::Pending as u32),
217        "RUNNING" => Some(RtdlNodeStateEnum::Running as u32),
218        "SUCCEEDED" => Some(RtdlNodeStateEnum::Succeeded as u32),
219        "FAILED" => Some(RtdlNodeStateEnum::Failed as u32),
220        "CANCELED" | "CANCELLED" => Some(RtdlNodeStateEnum::Canceled as u32),
221        "TIMEOUT" => Some(RtdlNodeStateEnum::Timeout as u32),
222        "PAUSED" => Some(RtdlNodeStateEnum::Paused as u32),
223        _ => None,
224    }
225}
226
227fn terminal_result(
228    call: &CapabilityCall,
229    state: u32,
230    detail: &str,
231    raw: &str,
232) -> CapabilityCallResult {
233    let success = state == RtdlNodeStateEnum::Succeeded as u32;
234    CapabilityCallResult {
235        call_id: call.call_id.clone(),
236        provider_id: call.provider_id.clone(),
237        contract_id: call.contract_id.clone(),
238        success,
239        output: if success {
240            raw.to_string()
241        } else {
242            String::new()
243        },
244        error: if success {
245            String::new()
246        } else {
247            detail.to_string()
248        },
249    }
250}
251
252fn canceled_result(call: &CapabilityCall, error: &str) -> CapabilityCallResult {
253    failed_result(call, error)
254}
255
256/// Build a failed capability result for async control-plane failures.
257/// The original call identity is preserved so the terminal node event still
258/// correlates with the user-requested async capability, not the status poll.
259fn failed_result(call: &CapabilityCall, error: &str) -> CapabilityCallResult {
260    CapabilityCallResult {
261        call_id: call.call_id.clone(),
262        provider_id: call.provider_id.clone(),
263        contract_id: call.contract_id.clone(),
264        success: false,
265        output: String::new(),
266        error: error.to_string(),
267    }
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273
274    #[test]
275    fn extract_run_id_from_response() {
276        assert_eq!(extract_run_id(r#"{"run_id":"r1","accepted":true}"#), "r1");
277        assert_eq!(extract_run_id(r#"{"goal_id":"g1"}"#), "");
278    }
279
280    #[test]
281    fn parse_status_requires_state_field() {
282        let (s, d) = parse_status_json(r#"{"state":"SUCCEEDED","detail":"done"}"#);
283        assert_eq!(s, RtdlNodeStateEnum::Succeeded as u32);
284        assert_eq!(d, "done");
285    }
286
287    #[test]
288    fn parse_status_missing_state_fails() {
289        let (s, d) = parse_status_json(r#"{"known":true,"terminal":false}"#);
290        assert_eq!(s, RtdlNodeStateEnum::Failed as u32);
291        assert!(d.contains("missing required 'state' field"));
292    }
293}
robonix_executor/dispatch/async_poll.rs

robonix_executor/dispatch/
async_poll.rs
