Skip to main content

agentmux_srv/backend/viewer/
mod.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! The viewer: a read-only window onto this channel's agents for a paired
5//! device (agentmux-mobile's SPEC_AGENT_STATUS_AND_LIVE_PANE_FEED_2026_10_07
6//! §4, §5, §13.2, §13.3).
7//!
8//! What it adds next to the LAN routes, and why each part exists:
9//! - a credential that is not public: a per-device viewer token, handed out
10//!   once against a pairing code shown on this screen, stored hashed, and
11//!   revocable (`pairing`, `storage::viewer_devices`). Never the `lan_key`
12//!   (anyone on the network has it) and never the instance `auth_key`.
13//! - encryption on the wire: the viewer routes are served only on a TLS
14//!   listener whose self-signed certificate the device pins by the
15//!   fingerprint in the QR (`cert`). The listener follows LAN discovery like
16//!   the other LAN listeners (`lan_listeners`), on its own port of the LAN
17//!   port range.
18//! - the feed itself: an agent's transcript, a bounded snapshot then live
19//!   appends (`feed`, the route in `server::http_viewer`).
20
21pub mod cert;
22pub mod client;
23pub mod feed;
24pub mod pairing;
25pub mod panes;
26pub mod terminal;
27
28use std::collections::HashMap;
29use std::net::IpAddr;
30use std::path::PathBuf;
31use std::sync::atomic::{AtomicU16, Ordering};
32use std::sync::{Arc, OnceLock};
33
34use tokio_util::sync::CancellationToken;
35
36/// Open feeds one device may hold; the next gets 503.
37pub const MAX_FEEDS_PER_DEVICE: usize = 4;
38/// Open feeds this srv serves in all; the next gets 503.
39pub const MAX_FEEDS: usize = 16;
40/// A device's `last_seen_ms` is written at most this often.
41pub const LAST_SEEN_EVERY_MS: i64 = 60_000;
42
43/// Why a feed was refused.
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum FeedRefused {
46    PerDevice,
47    Total,
48}
49
50#[derive(Default)]
51struct FeedRegistry {
52    by_device: HashMap<String, Vec<(u64, CancellationToken)>>,
53    next_id: u64,
54}
55
56impl FeedRegistry {
57    fn total(&self) -> usize {
58        self.by_device.values().map(Vec::len).sum()
59    }
60}
61
62/// One open feed's place under the caps. Dropping it frees the place;
63/// `token` is cancelled when the device is revoked.
64pub struct FeedLease {
65    registry: Arc<parking_lot::Mutex<FeedRegistry>>,
66    device_id: String,
67    id: u64,
68    pub token: CancellationToken,
69}
70
71impl Drop for FeedLease {
72    fn drop(&mut self) {
73        let mut reg = self.registry.lock();
74        if let Some(list) = reg.by_device.get_mut(&self.device_id) {
75            list.retain(|(id, _)| *id != self.id);
76            if list.is_empty() {
77                reg.by_device.remove(&self.device_id);
78            }
79        }
80    }
81}
82
83/// The port the viewer listener is advertised on, process-wide, for the UDP
84/// identity replies (which have no `AppState` in hand). Installed once by
85/// `bootstrap`; absent in tests.
86static ADVERTISED_PORT: OnceLock<Arc<AtomicU16>> = OnceLock::new();
87
88/// The viewer listener's port while at least one viewer listener is bound.
89pub fn advertised_viewer_port() -> Option<u16> {
90    ADVERTISED_PORT.get().map(|p| p.load(Ordering::Relaxed)).filter(|p| *p != 0)
91}
92
93pub struct ViewerService {
94    /// Where the certificate is kept; `None` keeps it in memory only.
95    data_dir: Option<PathBuf>,
96    tls: parking_lot::Mutex<Option<Arc<cert::ViewerTls>>>,
97    pub pairing: pairing::Pairing,
98    pub hub: Arc<feed::FeedHub>,
99    feeds: Arc<parking_lot::Mutex<FeedRegistry>>,
100    max_per_device: usize,
101    max_total: usize,
102    /// 0 while no viewer listener is bound.
103    advertised_port: Arc<AtomicU16>,
104    /// The addresses viewer listeners are bound on, for the QR's `host`.
105    bound: parking_lot::Mutex<Vec<IpAddr>>,
106    /// The viewer port, held on loopback (bound, never listening) from the
107    /// first time the listener comes up until exit, so another instance on
108    /// this host never takes it from the range and the port a device saw
109    /// stays ours across LAN off/on.
110    reservation: parking_lot::Mutex<Option<(u16, socket2::Socket)>>,
111    /// When each device's `last_seen_ms` was last written.
112    last_seen_written: parking_lot::Mutex<HashMap<String, i64>>,
113}
114
115impl ViewerService {
116    /// A viewer whose feeds watch `broker`.
117    pub fn new(data_dir: Option<PathBuf>, broker: &crate::backend::mps::Broker) -> Arc<Self> {
118        Self::with_limits(data_dir, broker, Arc::new(feed::FeedHub::default()), MAX_FEEDS_PER_DEVICE, MAX_FEEDS)
119    }
120
121    pub fn with_limits(
122        data_dir: Option<PathBuf>,
123        broker: &crate::backend::mps::Broker,
124        hub: Arc<feed::FeedHub>,
125        max_per_device: usize,
126        max_total: usize,
127    ) -> Arc<Self> {
128        hub.install(broker);
129        Arc::new(Self {
130            data_dir,
131            tls: parking_lot::Mutex::new(None),
132            pairing: pairing::Pairing::new(),
133            hub,
134            feeds: Arc::new(parking_lot::Mutex::new(FeedRegistry::default())),
135            max_per_device,
136            max_total,
137            advertised_port: Arc::new(AtomicU16::new(0)),
138            bound: parking_lot::Mutex::new(Vec::new()),
139            reservation: parking_lot::Mutex::new(None),
140            last_seen_written: parking_lot::Mutex::new(HashMap::new()),
141        })
142    }
143
144    /// Make this the process's advertised viewer port (`advertised_viewer_port`).
145    pub fn install_global(&self) {
146        let _ = ADVERTISED_PORT.set(Arc::clone(&self.advertised_port));
147    }
148
149    /// The handle the fleet feed reads `viewer_port` from.
150    pub fn advertised_port_handle(&self) -> Arc<AtomicU16> {
151        Arc::clone(&self.advertised_port)
152    }
153
154    pub fn advertised_port(&self) -> Option<u16> {
155        Some(self.advertised_port.load(Ordering::Relaxed)).filter(|p| *p != 0)
156    }
157
158    /// The certificate, made and kept on first need.
159    pub fn tls(&self) -> Result<Arc<cert::ViewerTls>, String> {
160        let mut slot = self.tls.lock();
161        if let Some(tls) = slot.as_ref() {
162            return Ok(Arc::clone(tls));
163        }
164        let tls = Arc::new(match &self.data_dir {
165            Some(dir) => cert::load_or_create(dir)?,
166            None => return Err("the viewer has no data directory for its certificate".to_string()),
167        });
168        *slot = Some(Arc::clone(&tls));
169        Ok(tls)
170    }
171
172    /// The viewer port: the one already held, else the first port of
173    /// `candidates` free on loopback, now held. `None` when none is free.
174    pub fn reserve_port(&self, candidates: impl IntoIterator<Item = u16>) -> Option<u16> {
175        let mut slot = self.reservation.lock();
176        if let Some((port, _)) = slot.as_ref() {
177            return Some(*port);
178        }
179        for port in candidates {
180            let Ok(socket) = socket2::Socket::new(socket2::Domain::IPV4, socket2::Type::STREAM, None) else {
181                continue;
182            };
183            let addr = std::net::SocketAddr::from(([127, 0, 0, 1], port));
184            if socket.bind(&addr.into()).is_ok() {
185                *slot = Some((port, socket));
186                return Some(port);
187            }
188        }
189        None
190    }
191
192    /// Record which addresses viewer listeners are bound on, and advertise
193    /// the port while there is at least one.
194    pub fn set_bound(&self, addrs: Vec<IpAddr>, port: u16) {
195        let up = !addrs.is_empty();
196        *self.bound.lock() = addrs;
197        self.advertised_port.store(if up { port } else { 0 }, Ordering::Relaxed);
198    }
199
200    /// The IPv4 address the QR names: the primary outbound one when a viewer
201    /// listener is bound there, else the first bound IPv4 address.
202    pub fn pair_host(&self) -> Option<IpAddr> {
203        let bound = self.bound.lock().clone();
204        let primary = crate::backend::lan_listeners::primary_lan_ipv4();
205        primary
206            .filter(|p| bound.contains(p))
207            .or_else(|| bound.iter().copied().find(IpAddr::is_ipv4))
208    }
209
210    /// Take a place for one more feed of `device_id`, or say why not.
211    pub fn open_feed(&self, device_id: &str) -> Result<FeedLease, FeedRefused> {
212        let mut reg = self.feeds.lock();
213        if reg.total() >= self.max_total {
214            return Err(FeedRefused::Total);
215        }
216        if reg.by_device.get(device_id).map_or(0, Vec::len) >= self.max_per_device {
217            return Err(FeedRefused::PerDevice);
218        }
219        reg.next_id += 1;
220        let id = reg.next_id;
221        let token = CancellationToken::new();
222        reg.by_device.entry(device_id.to_string()).or_default().push((id, token.clone()));
223        Ok(FeedLease { registry: Arc::clone(&self.feeds), device_id: device_id.to_string(), id, token })
224    }
225
226    /// Close every open feed of `device_id` (it was revoked). Returns how many.
227    pub fn close_device_feeds(&self, device_id: &str) -> usize {
228        let reg = self.feeds.lock();
229        let Some(list) = reg.by_device.get(device_id) else { return 0 };
230        for (_, token) in list {
231            token.cancel();
232        }
233        list.len()
234    }
235
236    #[cfg(test)]
237    pub fn open_feeds(&self) -> usize {
238        self.feeds.lock().total()
239    }
240
241    /// Whether `device_id`'s `last_seen_ms` is due to be written at `now_ms`;
242    /// records it when it is.
243    pub fn last_seen_due(&self, device_id: &str, now_ms: i64) -> bool {
244        let mut written = self.last_seen_written.lock();
245        match written.get(device_id) {
246            Some(at) if now_ms - at < LAST_SEEN_EVERY_MS => false,
247            _ => {
248                written.insert(device_id.to_string(), now_ms);
249                true
250            }
251        }
252    }
253
254    /// Forget a revoked device's throttle entry.
255    pub fn forget_device(&self, device_id: &str) {
256        self.last_seen_written.lock().remove(device_id);
257    }
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263
264    fn service(per_device: usize, total: usize) -> Arc<ViewerService> {
265        let broker = crate::backend::mps::Broker::new();
266        ViewerService::with_limits(None, &broker, Arc::new(feed::FeedHub::default()), per_device, total)
267    }
268
269    #[test]
270    fn feeds_are_capped_per_device_and_in_all() {
271        let v = service(2, 3);
272        let a1 = v.open_feed("a").unwrap();
273        let _a2 = v.open_feed("a").unwrap();
274        assert_eq!(v.open_feed("a").err(), Some(FeedRefused::PerDevice));
275        let _b1 = v.open_feed("b").unwrap();
276        assert_eq!(v.open_feed("c").err(), Some(FeedRefused::Total));
277        drop(a1);
278        assert_eq!(v.open_feeds(), 2);
279        assert!(v.open_feed("c").is_ok(), "a closed feed frees its place");
280    }
281
282    #[test]
283    fn closing_a_devices_feeds_cancels_only_its_own() {
284        let v = service(4, 16);
285        let a1 = v.open_feed("a").unwrap();
286        let a2 = v.open_feed("a").unwrap();
287        let b1 = v.open_feed("b").unwrap();
288        assert_eq!(v.close_device_feeds("a"), 2);
289        assert!(a1.token.is_cancelled() && a2.token.is_cancelled());
290        assert!(!b1.token.is_cancelled());
291        assert_eq!(v.close_device_feeds("nobody"), 0);
292    }
293
294    #[test]
295    fn last_seen_is_written_at_most_once_a_minute() {
296        let v = service(4, 16);
297        assert!(v.last_seen_due("a", 1_000));
298        assert!(!v.last_seen_due("a", 1_000 + LAST_SEEN_EVERY_MS - 1));
299        assert!(v.last_seen_due("b", 1_000), "per device");
300        assert!(v.last_seen_due("a", 1_000 + LAST_SEEN_EVERY_MS));
301    }
302
303    #[test]
304    fn the_port_is_reserved_once_and_skips_a_taken_one() {
305        let taken = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
306        let taken_port = taken.local_addr().unwrap().port();
307        let free = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port();
308        let v = service(4, 16);
309        assert_eq!(v.reserve_port([taken_port, free]), Some(free));
310        assert!(
311            std::net::TcpListener::bind(("127.0.0.1", free)).is_err(),
312            "the reservation holds the port against other instances"
313        );
314        assert_eq!(v.reserve_port([taken_port]), Some(free), "kept, not chosen again");
315    }
316
317    #[test]
318    fn the_port_is_advertised_only_while_bound() {
319        let v = service(4, 16);
320        assert_eq!(v.advertised_port(), None);
321        v.set_bound(vec![IpAddr::from([198, 51, 100, 7])], 29702);
322        assert_eq!(v.advertised_port(), Some(29702));
323        assert_eq!(v.pair_host(), Some(IpAddr::from([198, 51, 100, 7])));
324        v.set_bound(Vec::new(), 29702);
325        assert_eq!(v.advertised_port(), None);
326        assert_eq!(v.pair_host(), None);
327    }
328
329    #[test]
330    fn without_a_data_directory_there_is_no_certificate() {
331        assert!(service(4, 16).tls().is_err());
332    }
333}