omni_cli/
receive.rs

1//! Receive omni casts from a relay as an NDJSON stream.
2//!
3//! `omni listen` and `omni get` are lightweight local daemon subscribers. The
4//! always-on `omni daemon` owns the remote relay registration and OS injection,
5//! then mirrors each received [`InjectionCommand`] to every subscriber on its
6//! existing local port. Subscribers do not inject keys or open another listening
7//! socket, which makes the stream suitable for shell pipelines and Pi extensions.
8
9use std::io::{BufRead, BufReader, Write};
10use std::thread;
11use std::time::Duration;
12
13use crate::AppError;
14use crate::transport::relay::ControlFrame;
15use crate::transport::{CONNECT_TIMEOUT, connect_bounded};
16use crate::wire::{InjectionCommand, WIRE_VERSION};
17
18/// Why one connected receive session ended.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum SessionEnd {
21    /// The relay closed the connection before the requested command limit.
22    Closed,
23    /// The relay rejected registration, normally because the token is wrong.
24    Rejected,
25    /// The requested number of commands was emitted (`omni get` uses one).
26    Complete,
27}
28
29/// Subscribe to a local daemon and copy received commands to `output` as NDJSON.
30///
31/// `max_commands = None` drains until the connection closes. Control frames are
32/// not written to output; daemon pings are answered with pongs. This function is
33/// public so applications can consume the same low-latency stream directly
34/// without spawning `omni` or implementing the subscription handshake.
35///
36/// # Errors
37/// Returns [`AppError::internal`] on relay or output I/O failure.
38pub fn receive_session<R: BufRead, W: Write, O: Write>(
39    reader: R,
40    daemon_writer: W,
41    output: O,
42    token: &str,
43    max_commands: Option<usize>,
44) -> Result<SessionEnd, AppError> {
45    receive_session_kind(reader, daemon_writer, output, token, max_commands, false)
46}
47
48/// Receive only non-injecting application events, never the command tap.
49pub fn receive_events_session<R: BufRead, W: Write, O: Write>(
50    reader: R,
51    daemon_writer: W,
52    output: O,
53    token: &str,
54    max_events: Option<usize>,
55) -> Result<SessionEnd, AppError> {
56    receive_session_kind(reader, daemon_writer, output, token, max_events, true)
57}
58
59fn receive_session_kind<R: BufRead, W: Write, O: Write>(
60    reader: R,
61    mut daemon_writer: W,
62    mut output: O,
63    token: &str,
64    max_commands: Option<usize>,
65    events: bool,
66) -> Result<SessionEnd, AppError> {
67    if max_commands == Some(0) {
68        return Ok(SessionEnd::Complete);
69    }
70
71    let subscribe = if events {
72        ControlFrame::SubscribeEvents {
73            v: WIRE_VERSION,
74            token: token.to_owned(),
75        }
76    } else {
77        ControlFrame::Subscribe {
78            v: WIRE_VERSION,
79            token: token.to_owned(),
80        }
81    };
82    let line = subscribe.json_line()?;
83    daemon_writer
84        .write_all(line.as_bytes())
85        .and_then(|()| daemon_writer.write_all(b"\n"))
86        .and_then(|()| daemon_writer.flush())
87        .map_err(|e| AppError::internal("receive_subscribe", e.to_string()))?;
88
89    let mut received = 0usize;
90    for line in reader.lines() {
91        let line = line.map_err(|e| AppError::internal("receive_read", e.to_string()))?;
92        let trimmed = line.trim();
93        if trimmed.is_empty() {
94            continue;
95        }
96
97        let encoded = if events {
98            crate::events::SemanticEvent::decode_line(trimmed).and_then(|e| e.json_line())
99        } else {
100            InjectionCommand::decode_line(trimmed).and_then(|c| c.json_line())
101        };
102        if let Ok(encoded) = encoded {
103            output
104                .write_all(encoded.as_bytes())
105                .and_then(|()| output.write_all(b"\n"))
106                .and_then(|()| output.flush())
107                .map_err(|e| AppError::internal("receive_output", e.to_string()))?;
108            received += 1;
109            if max_commands.is_some_and(|max| received >= max) {
110                return Ok(SessionEnd::Complete);
111            }
112            continue;
113        }
114
115        if let Ok(frame) = serde_json::from_str::<ControlFrame>(trimmed) {
116            match frame {
117                ControlFrame::Error { .. } => return Ok(SessionEnd::Rejected),
118                ControlFrame::Ping { .. } => {
119                    // Keep control traffic off stdout: answer only on the daemon socket.
120                    daemon_writer
121                        .write_all(br#"{"v":1,"op":"pong"}"#)
122                        .and_then(|()| daemon_writer.write_all(b"\n"))
123                        .and_then(|()| daemon_writer.flush())
124                        .map_err(|e| AppError::internal("receive_pong", e.to_string()))?;
125                }
126                _ => {}
127            }
128        }
129    }
130    Ok(SessionEnd::Closed)
131}
132
133/// Options for [`receive_from_daemon`].
134#[derive(Debug, Clone)]
135pub struct ReceiveOptions {
136    /// Shared daemon token. Empty is valid for a loopback/local-trust daemon.
137    pub token: String,
138    /// Number of commands to emit, or `None` to listen continuously.
139    pub max_commands: Option<usize>,
140    /// Reconnect after disconnects and transient failures.
141    pub reconnect: bool,
142    /// Delay between reconnect attempts.
143    pub backoff: Duration,
144    /// Optional attempt bound, primarily for embedded callers and tests.
145    pub max_attempts: Option<usize>,
146}
147
148/// Connect to a daemon, subscribe to its command tap, and stream to output.
149///
150/// With reconnect enabled, transient connect/read failures and clean daemon
151/// disconnects retry after `backoff`. Authentication rejection is terminal.
152///
153/// # Errors
154/// Returns an error on authentication rejection, a terminal connection/session
155/// failure, or output failure.
156pub fn receive_from_daemon<O: Write>(
157    daemon_addr: &str,
158    output: &mut O,
159    opts: &ReceiveOptions,
160) -> Result<(), AppError> {
161    receive_from_daemon_kind(daemon_addr, output, opts, false)
162}
163
164/// Event-only subscriber. A stream discontinuity emits a local `omni_stream_reset`
165/// cancel event before reconnecting, so held-state consumers can reset safely.
166pub fn receive_events_from_daemon<O: Write>(
167    daemon_addr: &str,
168    output: &mut O,
169    opts: &ReceiveOptions,
170) -> Result<(), AppError> {
171    receive_from_daemon_kind(daemon_addr, output, opts, true)
172}
173
174fn receive_from_daemon_kind<O: Write>(
175    daemon_addr: &str,
176    output: &mut O,
177    opts: &ReceiveOptions,
178    events: bool,
179) -> Result<(), AppError> {
180    let mut attempts = 0usize;
181    loop {
182        attempts += 1;
183        let result = match connect_bounded(daemon_addr, CONNECT_TIMEOUT) {
184            Ok(stream) => {
185                let reader = stream
186                    .try_clone()
187                    .map(BufReader::new)
188                    .map_err(|e| AppError::internal("receive_clone", e.to_string()))?;
189                receive_session_kind(
190                    reader,
191                    stream,
192                    &mut *output,
193                    &opts.token,
194                    opts.max_commands,
195                    events,
196                )
197            }
198            Err(e) => Err(AppError::internal(
199                "receive_connect",
200                format!("{daemon_addr}: {e}"),
201            )),
202        };
203
204        match result {
205            Ok(SessionEnd::Complete) => return Ok(()),
206            Ok(SessionEnd::Rejected) => {
207                return Err(AppError::validation(
208                    "receive_auth",
209                    "daemon rejected subscription; check the token",
210                ));
211            }
212            Ok(SessionEnd::Closed) => {
213                if !opts.reconnect {
214                    return Err(AppError::internal(
215                        "receive_closed",
216                        "daemon closed the connection",
217                    ));
218                }
219            }
220            Err(error) => {
221                // Output errors are local and retrying the daemon cannot fix them.
222                if error.code() == "receive_output" || !opts.reconnect {
223                    return Err(error);
224                }
225                eprintln!("omni receive: {error}; reconnecting");
226            }
227        }
228
229        if events {
230            let mut reset = crate::events::SemanticEvent::new(
231                "omni_stream_reset",
232                crate::events::Phase::Cancel,
233            )?;
234            reset.data = Some(serde_json::json!({"reason":"stream_discontinuity"}));
235            let line = reset.json_line()?;
236            output
237                .write_all(line.as_bytes())
238                .and_then(|()| output.write_all(b"\n"))
239                .and_then(|()| output.flush())
240                .map_err(|e| AppError::internal("receive_output", e.to_string()))?;
241        }
242        if opts.max_attempts.is_some_and(|max| attempts >= max) {
243            return Err(AppError::internal(
244                "receive_attempts",
245                format!("daemon receive stopped after {attempts} attempt(s)"),
246            ));
247        }
248        thread::sleep(opts.backoff);
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255    use crate::daemon::{CaptureInjector, CommandBus, DaemonOptions, run_daemon_with_bus};
256    use crate::transport::relay::{bind, cast_via_relay};
257    use std::io::Cursor;
258    use std::sync::{Arc, Mutex};
259
260    #[test]
261    fn session_subscribes_and_emits_only_commands() {
262        let input = concat!(
263            "{\"v\":1,\"op\":\"ping\"}\n",
264            "{\"v\":1,\"type\":\"key\",\"key\":\"enter\"}\n",
265            "{\"v\":1,\"type\":\"scroll\",\"amount\":3}\n",
266        );
267        let mut daemon_out = Vec::new();
268        let mut stdout = Vec::new();
269        let end = receive_session(
270            Cursor::new(input),
271            &mut daemon_out,
272            &mut stdout,
273            "tok",
274            Some(2),
275        )
276        .unwrap();
277
278        assert_eq!(end, SessionEnd::Complete);
279        let daemon_text = String::from_utf8(daemon_out).unwrap();
280        assert!(daemon_text.contains("\"op\":\"subscribe\""));
281        assert!(daemon_text.contains("\"token\":\"tok\""));
282        assert!(daemon_text.contains("\"op\":\"pong\""));
283        assert_eq!(
284            String::from_utf8(stdout).unwrap(),
285            concat!(
286                "{\"v\":1,\"type\":\"key\",\"key\":\"enter\"}\n",
287                "{\"v\":1,\"type\":\"scroll\",\"amount\":3}\n",
288            )
289        );
290    }
291
292    #[test]
293    fn session_reports_rejection_without_polluting_stdout() {
294        let input = "{\"op\":\"error\",\"v\":1,\"reason\":\"unauthorized\"}\n";
295        let mut daemon_out = Vec::new();
296        let mut stdout = Vec::new();
297        let end = receive_session(
298            Cursor::new(input),
299            &mut daemon_out,
300            &mut stdout,
301            "bad",
302            None,
303        )
304        .unwrap();
305        assert_eq!(end, SessionEnd::Rejected);
306        assert!(stdout.is_empty());
307    }
308
309    #[test]
310    fn zero_limit_completes_without_registering() {
311        let mut daemon_out = Vec::new();
312        let mut stdout = Vec::new();
313        let end = receive_session(
314            Cursor::new(Vec::<u8>::new()),
315            &mut daemon_out,
316            &mut stdout,
317            "tok",
318            Some(0),
319        )
320        .unwrap();
321        assert_eq!(end, SessionEnd::Complete);
322        assert!(daemon_out.is_empty());
323    }
324
325    #[derive(Clone, Default)]
326    struct SharedWriter(Arc<Mutex<Vec<u8>>>);
327
328    impl Write for SharedWriter {
329        fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
330            self.0.lock().unwrap().extend_from_slice(bytes);
331            Ok(bytes.len())
332        }
333
334        fn flush(&mut self) -> std::io::Result<()> {
335            Ok(())
336        }
337    }
338
339    #[test]
340    fn daemon_cast_reaches_local_subscriber_end_to_end() {
341        let listener = bind("127.0.0.1:0").unwrap();
342        let addr = listener.local_addr().unwrap().to_string();
343        let bus = CommandBus::new();
344        let server_bus = bus.clone();
345        let injector = Arc::new(CaptureInjector::new());
346        let server = thread::spawn(move || {
347            run_daemon_with_bus(
348                &listener,
349                injector,
350                &DaemonOptions {
351                    token: Some("tok".to_owned()),
352                    max_conns: Some(2),
353                },
354                &server_bus,
355            )
356            .unwrap();
357        });
358
359        let output = SharedWriter::default();
360        let captured = output.0.clone();
361        let receive_addr = addr.clone();
362        let receiver = thread::spawn(move || {
363            let mut output = output;
364            receive_from_daemon(
365                &receive_addr,
366                &mut output,
367                &ReceiveOptions {
368                    // Loopback subscribers use local trust even though the
369                    // daemon requires "tok" for inbound casts.
370                    token: String::new(),
371                    max_commands: Some(1),
372                    reconnect: true,
373                    backoff: Duration::from_millis(5),
374                    max_attempts: Some(5),
375                },
376            )
377            .unwrap();
378        });
379
380        for _ in 0..100 {
381            if bus.subscriber_count() == 1 {
382                break;
383            }
384            thread::sleep(Duration::from_millis(5));
385        }
386        assert_eq!(bus.subscriber_count(), 1);
387        let delivery =
388            cast_via_relay(&addr, "tok", "local", &[InjectionCommand::key("enter")]).unwrap();
389        assert!(
390            !delivery.acked,
391            "a daemon does not send relay delivery acks"
392        );
393
394        receiver.join().unwrap();
395        server.join().unwrap();
396        assert_eq!(
397            String::from_utf8(captured.lock().unwrap().clone()).unwrap(),
398            "{\"v\":1,\"type\":\"key\",\"key\":\"enter\"}\n"
399        );
400    }
401}
omni_cli/receive.rs

omni_cli/
receive.rs
