Skip to main content

omni_cli/
dispatch.rs

1//! Target dispatch (bd-bfb732): cast [`InjectionCommand`]s to the endpoints of
2//! a resolved config target, mapping each [`Transport`] to a concrete sender,
3//! with group fan-out and PER-ENDPOINT ERROR ISOLATION (one failed endpoint
4//! never aborts the rest) plus an aggregate report.
5
6use crate::AppError;
7use crate::config::{ResolvedEndpoint, Transport};
8use crate::transport::{android, relay, tcp};
9use crate::wire::InjectionCommand;
10
11/// Outcome for one endpoint in a fan-out.
12#[derive(Debug, Clone)]
13pub struct EndpointOutcome {
14    pub name: String,
15    pub kind: &'static str,
16    pub detail: String,
17    /// `Ok(lines)` on success (sent-summary or, for dry-run, the preview lines);
18    /// `Err(message)` on failure.
19    pub result: Result<Vec<String>, String>,
20}
21
22/// Aggregate report for a fan-out dispatch.
23#[derive(Debug, Clone)]
24pub struct DispatchReport {
25    pub outcomes: Vec<EndpointOutcome>,
26}
27
28impl DispatchReport {
29    #[must_use]
30    pub fn ok_count(&self) -> usize {
31        self.outcomes.iter().filter(|o| o.result.is_ok()).count()
32    }
33    #[must_use]
34    pub fn fail_count(&self) -> usize {
35        self.outcomes.iter().filter(|o| o.result.is_err()).count()
36    }
37    #[must_use]
38    pub fn all_ok(&self) -> bool {
39        self.outcomes.iter().all(|o| o.result.is_ok())
40    }
41    /// True when there were no endpoints at all (an empty/unresolvable target).
42    #[must_use]
43    pub fn is_empty(&self) -> bool {
44        self.outcomes.is_empty()
45    }
46}
47
48/// Dispatch `cmds` to every endpoint, isolating per-endpoint failures. When
49/// `dry_run` is set, render what WOULD be sent without executing.
50///
51/// A multi-endpoint target (a group, bd-bfc5a2) fans out CONCURRENTLY — one
52/// scoped thread per endpoint — so a broadcast completes in ~the time of its
53/// slowest member instead of the sum, while the report stays in endpoint order.
54#[must_use]
55pub fn dispatch(
56    endpoints: &[ResolvedEndpoint],
57    cmds: &[InjectionCommand],
58    dry_run: bool,
59) -> DispatchReport {
60    let outcomes = fan_out(endpoints.iter().collect::<Vec<_>>(), |ep| {
61        outcome_for(ep, cmds, dry_run)
62    });
63    DispatchReport { outcomes }
64}
65
66/// Build the [`EndpointOutcome`] for a single one-shot endpoint send.
67fn outcome_for(ep: &ResolvedEndpoint, cmds: &[InjectionCommand], dry_run: bool) -> EndpointOutcome {
68    EndpointOutcome {
69        name: ep.name.clone(),
70        kind: ep.transport.kind(),
71        detail: ep.transport.detail(),
72        result: send_to_endpoint(ep, cmds, dry_run).map_err(|e| e.to_string()),
73    }
74}
75
76/// Run `f` on every item CONCURRENTLY (one scoped thread each), collecting the
77/// results in INPUT ORDER so a fan-out completes in ~its slowest member's time
78/// instead of the sum while staying deterministically ordered (bd-bfc5a2). A
79/// single item runs inline (no thread) to keep the common unicast path cheap. A
80/// panicking worker is re-raised on the caller, matching the sequential
81/// behaviour it replaced.
82fn fan_out<T, R, F>(items: Vec<T>, f: F) -> Vec<R>
83where
84    T: Send,
85    R: Send,
86    F: Fn(T) -> R + Sync,
87{
88    if items.len() <= 1 {
89        return items.into_iter().map(f).collect();
90    }
91    std::thread::scope(|scope| {
92        let f = &f;
93        let handles: Vec<_> = items
94            .into_iter()
95            .map(|item| scope.spawn(move || f(item)))
96            .collect();
97        handles
98            .into_iter()
99            .map(|h| h.join().unwrap_or_else(|p| std::panic::resume_unwind(p)))
100            .collect()
101    })
102}
103
104fn send_to_endpoint(
105    ep: &ResolvedEndpoint,
106    cmds: &[InjectionCommand],
107    dry_run: bool,
108) -> Result<Vec<String>, AppError> {
109    match &ep.transport {
110        Transport::AdbBroadcast {
111            serial,
112            sender_package,
113        } => {
114            let sp = sender_package.as_deref();
115            let mut rendered = Vec::new();
116            for cmd in cmds {
117                let casts = android::to_am_broadcasts(cmd)?;
118                let lines = android::dispatch(&casts, serial.as_deref(), sp, !dry_run)?;
119                rendered.extend(lines);
120            }
121            Ok(rendered)
122        }
123        Transport::Termux {} => {
124            // On-device local `am broadcast` (no adb). Reuse the android command
125            // rendering and run it via the local shell; OFF-device this fails
126            // (am not found) and is isolated + reported per endpoint.
127            let mut rendered = Vec::new();
128            for cmd in cmds {
129                for cast in &android::to_am_broadcasts(cmd)? {
130                    let device_cmd = cast.device_command(android::ACTION, None);
131                    rendered.push(format!("sh -c '{device_cmd}'"));
132                    if !dry_run {
133                        run_local_shell(&device_cmd)?;
134                    }
135                }
136            }
137            Ok(rendered)
138        }
139        Transport::Tcp { host, port } => {
140            if dry_run {
141                tcp::render_bare(host, *port, cmds)
142            } else {
143                tcp::cast_bare(host, *port, cmds)?;
144                Ok(vec![format!(
145                    "tcp {host}:{port} <- {} command(s)",
146                    cmds.len()
147                )])
148            }
149        }
150        Transport::Relay {
151            url,
152            device,
153            token_env,
154        } => {
155            let env = token_env.as_deref().unwrap_or(relay::TOKEN_ENV);
156            let addr = relay_addr(url);
157            if dry_run {
158                Ok(vec![format!(
159                    "relay {addr} device={device} <- {} command(s)",
160                    cmds.len()
161                )])
162            } else {
163                let token = std::env::var(env).map_err(|_| {
164                    AppError::validation(
165                        "dispatch_relay_token",
166                        format!("relay token env {env} is not set for this endpoint"),
167                    )
168                })?;
169                let delivery = relay::cast_via_relay(&addr, &token, device, cmds)?;
170                let mut lines = vec![format!(
171                    "relay {addr} device={device} <- {} command(s)",
172                    cmds.len()
173                )];
174                if let Some(warning) = relay_delivery_warning(&addr, device, &delivery) {
175                    lines.push(warning);
176                }
177                Ok(lines)
178            }
179        }
180    }
181}
182
183/// Strip an optional `tcp://` (or `tcps://`) scheme from a relay url so it is a
184/// bare `host:port` for [`std::net::TcpStream::connect`].
185fn relay_addr(url: &str) -> String {
186    url.strip_prefix("tcp://")
187        .or_else(|| url.strip_prefix("tcps://"))
188        .unwrap_or(url)
189        .to_owned()
190}
191
192/// A per-relay-endpoint delivery warning for the `--target` dispatch path
193/// (bd-b9c6b5): `Some(line)` when the relay CONFIRMS the cast reached 0 devices
194/// (the endpoint's device is offline), mirroring the direct `omni send --relay`
195/// warning. `None` when delivered, or when delivery is unconfirmed (a legacy
196/// relay that sends no result).
197fn relay_delivery_warning(
198    addr: &str,
199    device: &str,
200    delivery: &relay::RelayDelivery,
201) -> Option<String> {
202    if delivery.acked && delivery.delivered == 0 {
203        Some(format!(
204            "relay {addr} warning: device '{device}' is not connected (delivered to 0 devices)"
205        ))
206    } else {
207        None
208    }
209}
210
211fn run_local_shell(device_cmd: &str) -> Result<(), AppError> {
212    let status = std::process::Command::new("sh")
213        .arg("-c")
214        .arg(device_cmd)
215        .status()
216        .map_err(|e| {
217            AppError::internal(
218                "termux_spawn",
219                format!("failed to run am broadcast locally (termux is on-device only): {e}"),
220            )
221        })?;
222    if !status.success() {
223        return Err(AppError::internal(
224            "termux_broadcast",
225            format!("local am broadcast exited with status {status}"),
226        ));
227    }
228    Ok(())
229}
230
231// ---------------------------------------------------------------------------
232// streaming dispatch: hold-open per-endpoint senders (bd-772cb2)
233// ---------------------------------------------------------------------------
234
235/// A per-endpoint persistent sender for a streaming cast.
236enum LiveSender {
237    /// Stateless: each command is a separate adb am broadcast.
238    Adb {
239        serial: Option<String>,
240        sender_package: Option<String>,
241    },
242    /// Stateless: each command is a separate local am broadcast.
243    Termux,
244    /// Hold-open bare-NDJSON-over-TCP.
245    Tcp(tcp::TcpSender),
246    /// Hold-open relay sender.
247    Relay(relay::RelaySender),
248    /// A build-time failure (e.g. relay token env unset), reported per cast so
249    /// it stays isolated to this endpoint.
250    Broken(String),
251}
252
253impl LiveSender {
254    fn cast_one(&mut self, cmd: &InjectionCommand) -> Result<(), AppError> {
255        match self {
256            LiveSender::Adb {
257                serial,
258                sender_package,
259            } => {
260                let casts = android::to_am_broadcasts(cmd)?;
261                android::dispatch(&casts, serial.as_deref(), sender_package.as_deref(), true)?;
262                Ok(())
263            }
264            LiveSender::Termux => {
265                for cast in &android::to_am_broadcasts(cmd)? {
266                    run_local_shell(&cast.device_command(android::ACTION, None))?;
267                }
268                Ok(())
269            }
270            LiveSender::Tcp(sender) => sender.cast_one(cmd),
271            LiveSender::Relay(sender) => sender.cast_one(cmd),
272            LiveSender::Broken(msg) => {
273                Err(AppError::validation("dispatch_live_sender", msg.clone()))
274            }
275        }
276    }
277}
278
279struct LiveEndpoint {
280    name: String,
281    kind: &'static str,
282    detail: String,
283    sender: LiveSender,
284}
285
286/// Persistent fan-out dispatcher for a streaming cast (bd-772cb2): builds ONE
287/// hold-open sender per resolved endpoint, then routes each command to all of
288/// them with PER-ENDPOINT ERROR ISOLATION (a failed endpoint logs but never
289/// aborts the stream; tcp/relay reconnect on transient errors). Unlike the
290/// one-shot [`dispatch`], tcp/relay endpoints keep their connection open across
291/// the whole stream instead of reconnecting per command.
292pub struct LiveDispatcher {
293    senders: Vec<LiveEndpoint>,
294}
295
296impl LiveDispatcher {
297    /// Build hold-open senders for every endpoint. Relay tokens are resolved
298    /// from their `token_env` now; an unset token becomes a sender that reports
299    /// the error per cast (preserving per-endpoint isolation).
300    #[must_use]
301    pub fn build(endpoints: &[ResolvedEndpoint]) -> Self {
302        let senders = endpoints
303            .iter()
304            .map(|ep| {
305                let sender = match &ep.transport {
306                    Transport::AdbBroadcast {
307                        serial,
308                        sender_package,
309                    } => LiveSender::Adb {
310                        serial: serial.clone(),
311                        sender_package: sender_package.clone(),
312                    },
313                    Transport::Termux {} => LiveSender::Termux,
314                    Transport::Tcp { host, port } => {
315                        LiveSender::Tcp(tcp::TcpSender::new(host, *port))
316                    }
317                    Transport::Relay {
318                        url,
319                        device,
320                        token_env,
321                    } => {
322                        let env = token_env.as_deref().unwrap_or(relay::TOKEN_ENV);
323                        match std::env::var(env) {
324                            Ok(token) => LiveSender::Relay(relay::RelaySender::new(
325                                relay_addr(url),
326                                device.clone(),
327                                token,
328                            )),
329                            Err(_) => LiveSender::Broken(format!(
330                                "relay token env {env} is not set for this endpoint"
331                            )),
332                        }
333                    }
334                };
335                LiveEndpoint {
336                    name: ep.name.clone(),
337                    kind: ep.transport.kind(),
338                    detail: ep.transport.detail(),
339                    sender,
340                }
341            })
342            .collect();
343        Self { senders }
344    }
345
346    /// Send one command to every endpoint, returning a per-endpoint outcome
347    /// (errors are isolated, never propagated). Endpoints are cast to
348    /// CONCURRENTLY (one scoped thread each, bd-bfc5a2) so a broadcast command
349    /// reaches every device in ~the slowest endpoint's time, with outcomes kept
350    /// in endpoint order; a single endpoint runs inline.
351    pub fn cast(&mut self, cmd: &InjectionCommand) -> Vec<EndpointOutcome> {
352        fan_out(self.senders.iter_mut().collect::<Vec<_>>(), |ep| {
353            live_outcome(ep, cmd)
354        })
355    }
356
357    /// Number of endpoints this dispatcher fans out to.
358    #[must_use]
359    pub fn endpoint_count(&self) -> usize {
360        self.senders.len()
361    }
362}
363
364/// Build the [`EndpointOutcome`] for one streaming endpoint cast (used by the
365/// concurrent fan-out in [`LiveDispatcher::cast`]).
366fn live_outcome(ep: &mut LiveEndpoint, cmd: &InjectionCommand) -> EndpointOutcome {
367    EndpointOutcome {
368        name: ep.name.clone(),
369        kind: ep.kind,
370        detail: ep.detail.clone(),
371        result: ep
372            .sender
373            .cast_one(cmd)
374            .map(|()| Vec::new())
375            .map_err(|e| e.to_string()),
376    }
377}
378
379#[cfg(test)]
380mod tests {
381    use super::*;
382
383    #[test]
384    fn fan_out_preserves_input_order() {
385        let out = fan_out((0..5).collect::<Vec<_>>(), |i| i * 2);
386        assert_eq!(out, vec![0, 2, 4, 6, 8]);
387    }
388
389    #[test]
390    fn fan_out_single_and_empty_use_inline_path() {
391        assert_eq!(fan_out(vec![21], |i| i + 1), vec![22]);
392        let empty: Vec<i32> = fan_out(Vec::new(), |i| i);
393        assert!(empty.is_empty());
394    }
395
396    #[test]
397    fn fan_out_runs_items_concurrently() {
398        use std::time::{Duration, Instant};
399        // Four 100ms jobs run sequentially would take ~400ms; concurrently they
400        // finish in ~100ms. Assert well under the sequential sum (tolerant of CI
401        // scheduling jitter) to prove the fan-out is concurrent, and that the
402        // results still come back in input order.
403        let start = Instant::now();
404        let out = fan_out((0..4).collect::<Vec<_>>(), |i: u64| {
405            std::thread::sleep(Duration::from_millis(100));
406            i
407        });
408        let elapsed = start.elapsed();
409        assert_eq!(out, vec![0, 1, 2, 3]);
410        assert!(
411            elapsed < Duration::from_millis(350),
412            "fan_out should run concurrently, took {elapsed:?}"
413        );
414    }
415
416    fn ep(name: &str, transport: Transport) -> ResolvedEndpoint {
417        ResolvedEndpoint {
418            name: name.to_owned(),
419            transport,
420        }
421    }
422
423    #[test]
424    fn relay_addr_strips_scheme() {
425        assert_eq!(relay_addr("tcp://relay.example:9000"), "relay.example:9000");
426        assert_eq!(relay_addr("relay.example:9000"), "relay.example:9000");
427    }
428
429    #[test]
430    fn dry_run_renders_each_transport_without_executing() {
431        let endpoints = vec![
432            ep(
433                "phone",
434                Transport::AdbBroadcast {
435                    serial: Some("ABC123".into()),
436                    sender_package: None,
437                },
438            ),
439            ep(
440                "watch",
441                Transport::Tcp {
442                    host: "1.2.3.4".into(),
443                    port: 8099,
444                },
445            ),
446            ep(
447                "ipad",
448                Transport::Relay {
449                    url: "tcp://relay:9000".into(),
450                    device: "harry-ipad".into(),
451                    token_env: None,
452                },
453            ),
454        ];
455        let report = dispatch(&endpoints, &[InjectionCommand::text("hi")], true);
456        assert!(
457            report.all_ok(),
458            "dry-run should succeed for all: {report:?}"
459        );
460        assert_eq!(report.ok_count(), 3);
461        // adb line carries the serial; tcp line carries the bare NDJSON; relay names the device.
462        assert!(report.outcomes[0].result.as_ref().unwrap()[0].contains("adb -s ABC123"));
463        assert!(
464            report.outcomes[1].result.as_ref().unwrap()[0]
465                .contains(r#"{"v":1,"type":"text","text":"hi"}"#)
466        );
467        assert!(report.outcomes[2].result.as_ref().unwrap()[0].contains("device=harry-ipad"));
468    }
469
470    #[test]
471    fn per_endpoint_error_isolation() {
472        // dry-run with an unmappable key: the adb endpoint fails (no keycode),
473        // but the tcp endpoint sends the bare NDJSON regardless — the failure is
474        // isolated and BOTH are reported. No external processes / sockets needed.
475        let endpoints = vec![
476            ep(
477                "phone",
478                Transport::AdbBroadcast {
479                    serial: None,
480                    sender_package: None,
481                },
482            ),
483            ep(
484                "watch",
485                Transport::Tcp {
486                    host: "1.2.3.4".into(),
487                    port: 8099,
488                },
489            ),
490        ];
491        let report = dispatch(&endpoints, &[InjectionCommand::key("nope")], true);
492        assert_eq!(report.ok_count(), 1);
493        assert_eq!(report.fail_count(), 1);
494        assert!(!report.all_ok());
495        let phone = report.outcomes.iter().find(|o| o.name == "phone").unwrap();
496        assert!(phone.result.as_ref().unwrap_err().contains("keycode"));
497        let watch = report.outcomes.iter().find(|o| o.name == "watch").unwrap();
498        assert!(watch.result.is_ok());
499    }
500
501    #[test]
502    fn relay_endpoint_unset_token_fails_fast() {
503        // execute mode, but an unset token env errors at the env check BEFORE any
504        // socket connect (so the test never hangs / needs a relay).
505        let endpoints = vec![ep(
506            "ipad",
507            Transport::Relay {
508                url: "tcp://relay:9000".into(),
509                device: "x".into(),
510                token_env: Some("OMNI_DEFINITELY_UNSET_TOKEN_ENV_XYZ".into()),
511            },
512        )];
513        let report = dispatch(&endpoints, &[InjectionCommand::text("hi")], false);
514        assert_eq!(report.fail_count(), 1);
515        assert!(
516            report.outcomes[0]
517                .result
518                .as_ref()
519                .unwrap_err()
520                .contains("token env")
521        );
522    }
523
524    #[test]
525    fn live_dispatcher_builds_and_isolates_broken_relay() {
526        // A relay endpoint with an unset token env builds into a Broken sender
527        // that fails per cast; a tcp endpoint to a closed loopback port fails
528        // fast on connect (refused, NOT a slow timeout) — both isolated.
529        let endpoints = vec![
530            ep(
531                "watch",
532                Transport::Tcp {
533                    host: "127.0.0.1".into(),
534                    port: 1,
535                },
536            ),
537            ep(
538                "ipad",
539                Transport::Relay {
540                    url: "tcp://relay:9000".into(),
541                    device: "x".into(),
542                    token_env: Some("OMNI_DEFINITELY_UNSET_TOKEN_ENV_LIVE".into()),
543                },
544            ),
545        ];
546        let mut live = LiveDispatcher::build(&endpoints);
547        assert_eq!(live.endpoint_count(), 2);
548        let outcomes = live.cast(&InjectionCommand::text("hi"));
549        assert_eq!(outcomes.len(), 2);
550        let ipad = outcomes.iter().find(|o| o.name == "ipad").unwrap();
551        assert!(ipad.result.as_ref().unwrap_err().contains("token env"));
552    }
553
554    #[test]
555    fn relay_delivery_warning_only_on_confirmed_zero() {
556        use crate::transport::relay::RelayDelivery;
557        // Confirmed not-delivered -> warning (target endpoint's device offline).
558        let w = relay_delivery_warning(
559            "h:9",
560            "watch",
561            &RelayDelivery {
562                acked: true,
563                delivered: 0,
564            },
565        );
566        assert!(w.unwrap().contains("not connected"));
567        // Delivered -> no warning.
568        assert!(
569            relay_delivery_warning(
570                "h:9",
571                "watch",
572                &RelayDelivery {
573                    acked: true,
574                    delivered: 1,
575                },
576            )
577            .is_none()
578        );
579        // Unconfirmed (legacy relay / no result) -> no warning.
580        assert!(
581            relay_delivery_warning(
582                "h:9",
583                "watch",
584                &RelayDelivery {
585                    acked: false,
586                    delivered: 0,
587                },
588            )
589            .is_none()
590        );
591    }
592}