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

robonix_executor/dispatch/
async_poll.rs
