rbnx/cmd/
ask.rs

1// SPDX-License-Identifier: MulanPSL-2.0
2// `rbnx ask` — non-interactive sibling of `rbnx chat`. Same atlas
3// + SubmitTask gRPC path, but no ratatui shell: prints events to
4// stdout as they arrive and exits when the pilot stream closes.
5//
6// Use cases:
7//   * scripted regression / smoke tests that need pilot in the loop
8//     without a human typing into the TUI
9//   * CI / agent-driven runs where stdout is the artifact
10//   * quick one-shot prompts ("describe what's in front of the robot")
11//
12// Wire format is identical to chat: same Task message, same session
13// semantics, same event stream. So if a prompt works in `rbnx ask`
14// it will work in `rbnx chat`, and vice versa.
15
16use anyhow::{Context, Result};
17use robonix_atlas::client::{self as atlas_client, AtlasClient};
18use std::io::{self, Write};
19use std::time::{SystemTime, UNIX_EPOCH};
20use tokio_stream::StreamExt;
21use tonic::Request;
22use uuid::Uuid;
23
24use crate::pb::contracts::robonix_system_pilot_client::RobonixSystemPilotClient;
25use crate::pb::pilot::rtdl_node_state::RtdlNodeStateEnum;
26use crate::pb::pilot::{CapabilityCall, Plan, Task};
27
28// PilotEvent.event_kind discriminants — must mirror service.rs / .msg.
29const EVT_TEXT_CHUNK: u32 = 0;
30const EVT_PLAN: u32 = 1;
31const EVT_BATCH_RESULT: u32 = 2;
32const EVT_STATUS: u32 = 3;
33const EVT_FINAL_TEXT: u32 = 4;
34
35const STATE_FAILED: u32 = 2;
36const CONSUMER_ID: &str = "rbnx-cli/ask";
37
38pub async fn execute(server: &str, prompt: &str, json: bool) -> Result<()> {
39    let mut atlas = AtlasClient::connect(server)
40        .await
41        .with_context(|| format!("connect to atlas at '{server}'"))?;
42    let (channel_id, pilot_cap_id, channel) =
43        atlas_client::connect_to_capability(&mut atlas, CONSUMER_ID, "robonix/system/pilot")
44            .await
45            .context("locate pilot via atlas")?;
46    if !json {
47        eprintln!("[ask] connected to pilot '{pilot_cap_id}' (channel {channel_id})");
48    }
49    let mut client = RobonixSystemPilotClient::new(channel);
50
51    let session_id = Uuid::new_v4().to_string();
52    let task = Task {
53        task_id: Uuid::new_v4().to_string(),
54        session_id: session_id.clone(),
55        source: 0, // INTENT_SOURCE_TEXT
56        text: prompt.to_string(),
57        audio_data: vec![],
58        context_json: String::new(),
59        timestamp_ms: now_ms(),
60    };
61    let mut stream = client
62        .submit_task(Request::new(task))
63        .await
64        .context("SubmitTask RPC failed")?
65        .into_inner();
66
67    let stdout = io::stdout();
68    let mut out = stdout.lock();
69    let mut last_was_chunk = false;
70    let mut had_failure = false;
71
72    while let Some(event) = stream.next().await {
73        let event = event.context("pilot stream error")?;
74        if json {
75            // One JSON object per event, line-delimited. Lets callers
76            // pipe through `jq` without grappling with the TUI shape.
77            let v = serde_json::json!({
78                "event_kind": event.event_kind,
79                "session_id": event.session_id,
80                "text_chunk": event.text_chunk,
81                "final_text": event.final_text,
82                "plan": event.plan.as_ref().map(|p| serde_json::json!({
83                    "round": p.round,
84                    "root_index": p.root_index,
85                    "calls": plan_calls(p).into_iter().map(|c| serde_json::json!({
86                        "call_id": c.call_id,
87                        "provider_id": c.provider_id,
88                        "contract_id": c.contract_id,
89                        "args_json": c.args_json,
90                    })).collect::<Vec<_>>(),
91                })),
92                "batch_result": event.batch_result.as_ref().map(|b| serde_json::json!({
93                    "round": b.round,
94                    "any_failed": b.any_failed,
95                    // One full per-node record (leaf and non-leaf); leaf nodes
96                    // carry the capability call result, operator nodes the detail.
97                    "results": b.results.iter().map(|r| serde_json::json!({
98                        "node_index": r.node_index,
99                        "node_kind": r.node_kind,
100                        "state": r.state,
101                        "op_id": r.op_id,
102                        "description": r.description,
103                        "operator_detail": r.operator_detail,
104                        "leaf_result": r.leaf_result.as_ref().map(|lr| serde_json::json!({
105                            "call_id": lr.call_id,
106                            "contract_id": lr.contract_id,
107                            "success": lr.success,
108                            "output": lr.output,
109                            "error": lr.error,
110                        })),
111                    })).collect::<Vec<_>>(),
112                })),
113                "status": event.status.as_ref().map(|s| serde_json::json!({
114                    "state": s.state,
115                    "message": s.message,
116                })),
117            });
118            writeln!(out, "{v}")?;
119            out.flush()?;
120            continue;
121        }
122
123        // Plain-text mode: stream agent text inline; surface tool
124        // calls + results on their own lines so a human reading the
125        // log can follow the loop without parsing JSON.
126        match event.event_kind {
127            EVT_TEXT_CHUNK => {
128                let chunk = event.text_chunk;
129                if chunk.is_empty() {
130                    continue;
131                }
132                if !last_was_chunk {
133                    write!(out, "\n[agent] ")?;
134                }
135                write!(out, "{chunk}")?;
136                out.flush()?;
137                last_was_chunk = true;
138            }
139            EVT_PLAN => {
140                if let Some(plan) = event.plan {
141                    if last_was_chunk {
142                        writeln!(out)?;
143                        last_was_chunk = false;
144                    }
145                    let leaves: Vec<String> = plan_calls(&plan)
146                        .into_iter()
147                        .map(|c| {
148                            c.contract_id
149                                .rsplit_once('/')
150                                .map(|(_, l)| l.to_string())
151                                .unwrap_or_else(|| c.contract_id.clone())
152                        })
153                        .collect();
154                    writeln!(out, "[plan r{}] {}", plan.round, leaves.join(", "))?;
155                    out.flush()?;
156                }
157            }
158            EVT_BATCH_RESULT => {
159                if let Some(batch) = event.batch_result {
160                    if last_was_chunk {
161                        writeln!(out)?;
162                        last_was_chunk = false;
163                    }
164                    for r in &batch.results {
165                        let success = r.state == RtdlNodeStateEnum::Succeeded as u32;
166                        // Leaf nodes carry a capability call result; non-leaf
167                        // (sequence/parallel) nodes carry an operator detail.
168                        let (label, body) = match r.leaf_result.as_ref() {
169                            Some(lr) => (
170                                lr.contract_id
171                                    .rsplit_once('/')
172                                    .map(|(_, l)| l.to_string())
173                                    .unwrap_or_else(|| lr.contract_id.clone()),
174                                if lr.success {
175                                    lr.output.clone()
176                                } else {
177                                    lr.error.clone()
178                                },
179                            ),
180                            None if !r.description.is_empty() => {
181                                (r.description.clone(), r.operator_detail.clone())
182                            }
183                            None => (format!("node{}", r.node_index), r.operator_detail.clone()),
184                        };
185                        let summary = compact_one_line(&body, 200);
186                        let mark = if success { "✓" } else { "✗" };
187                        writeln!(out, "  [{mark} {label}] {summary}")?;
188                    }
189                    out.flush()?;
190                }
191            }
192            EVT_STATUS => {
193                if let Some(s) = event.status {
194                    if last_was_chunk {
195                        writeln!(out)?;
196                        last_was_chunk = false;
197                    }
198                    if s.state == STATE_FAILED {
199                        had_failure = true;
200                        writeln!(out, "[status FAILED] {}", s.message)?;
201                    } else if !s.message.is_empty() {
202                        writeln!(out, "[status] {}", s.message)?;
203                    }
204                    out.flush()?;
205                }
206            }
207            EVT_FINAL_TEXT => {
208                // text_chunk events already streamed the full text.
209            }
210            _ => {}
211        }
212    }
213
214    if last_was_chunk {
215        writeln!(out)?;
216    }
217    let _ = atlas.disconnect_capability(&channel_id).await;
218    if had_failure {
219        anyhow::bail!("pilot reported FAILED status");
220    }
221    Ok(())
222}
223
224fn compact_one_line(s: &str, n: usize) -> String {
225    let flat: String = s
226        .chars()
227        .map(|c| if c == '\n' || c == '\r' { ' ' } else { c })
228        .collect();
229    if flat.chars().count() > n {
230        let mut out: String = flat.chars().take(n).collect();
231        out.push('…');
232        out
233    } else {
234        flat
235    }
236}
237
238fn plan_calls(plan: &Plan) -> Vec<&CapabilityCall> {
239    plan.nodes
240        .iter()
241        .filter_map(|node| node.call.as_ref())
242        .collect()
243}
244
245fn now_ms() -> u64 {
246    SystemTime::now()
247        .duration_since(UNIX_EPOCH)
248        .unwrap_or_default()
249        .as_millis() as u64
250}
rbnx/cmd/ask.rs

rbnx/cmd/
ask.rs
