diff --git a/cmux-tui/crates/cmux-tui-core/src/server/agent_session_attach_tests.rs b/cmux-tui/crates/cmux-tui-core/src/server/agent_session_attach_tests.rs index 587dddb5cf5d..c0945ab03a28 100644 --- a/cmux-tui/crates/cmux-tui-core/src/server/agent_session_attach_tests.rs +++ b/cmux-tui/crates/cmux-tui-core/src/server/agent_session_attach_tests.rs @@ -150,8 +150,14 @@ pub(super) struct Fixture { pub(super) watch: Vec, /// The records an `_acpmux/events` (history) page answers. pub(super) history: Vec, + /// Records every history page answers after [`Fixture::history`], added while the fake runs. + pub(super) history_later: Arc>>, /// When set, the attach answer waits for one message on it (at most [`WAIT`]). pub(super) attach_gate: Option>>>, + /// When set, each history page and each prompt's acceptance waits for one message on it (at + /// most [`WAIT`], or until its sender is gone). + pub(super) events_gate: Option>>>, + pub(super) prompt_gate: Option>>>, } fn serve( @@ -206,7 +212,12 @@ fn serve( ), })), "_acpmux/events" => { - Some(json!({"events": fixture.history, "hasMore": false, "lastSeq": 1})) + if let Some(gate) = &fixture.events_gate { + let _ = gate.lock().unwrap_or_else(|e| e.into_inner()).recv_timeout(WAIT); + } + let mut events = fixture.history.clone(); + events.extend(fixture.history_later.lock().unwrap().iter().cloned()); + Some(json!({"events": events, "hasMore": false, "lastSeq": 1})) } "_acpmux/typing" => { typing_calls += 1; @@ -239,6 +250,9 @@ fn serve( reply } "session/prompt" => { + if let Some(gate) = &fixture.prompt_gate { + let _ = gate.lock().unwrap_or_else(|e| e.into_inner()).recv_timeout(WAIT); + } let accepted = json!({"jsonrpc":"2.0","method":"_acpmux/prompt_accepted","params":{ "sessionId": SESSION, "promptId": params["_meta"]["acpmux"]["promptId"], "turnId": "turn_1", "queued": false}}); diff --git a/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_relay.rs b/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_relay.rs new file mode 100644 index 000000000000..4970a9ee04d3 --- /dev/null +++ b/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_relay.rs @@ -0,0 +1,705 @@ +//! The relay of one guest of a bridged link share (design chat-multiplayer.md 4.3). For each guest +//! whose browser is in the room, the bridge opens a same-user connection to this daemon, served by +//! the real connection path, binds it with `agent-session-proxy-bind {link_guest}` on its first +//! line, and attaches. So every frame is the bound guest's and passes the daemon's per-frame +//! checks (its share and guest live, the guest verbs only, every prompt held for the owner), and +//! revoke, kick and expiry end the connection at the daemon first. +//! +//! What the relay sends the room names this guest only: a `snapshot` first (the newest records +//! that fit in 512 KiB, `has_more` when older ones remain), then each live record, `page` for +//! `older`, and `prompt_state` for the guest's prompts, read from acpmux's records (`queued` +//! awaiting the owner, `queue_approved`, `queue_declined`, `queue_removed`). A record over +//! 256 KiB goes cut: its seq and kind with `truncated: true, full_bytes`. Records the daemon sends +//! before the attach answer wait for the snapshot. A prompt's state only moves forward (a +//! decision is final; held replaces refused), kept in the guest's [`Ledger`]. + +use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::{Arc, Mutex, MutexGuard}; + +use cmux_json::JsonSet as _; +use serde_json::{Value, json}; +use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _, BufReader}; +use tokio::net::unix::OwnedWriteHalf; +use tokio::sync::mpsc; + +use super::super::guest_prompt_id; +use crate::mux::Mux; +use crate::server::{RenderService, claim_connection, handle_connection_with_permit, transport}; + +/// A whole `snapshot` or `page` frame. +const MAX_FRAME: usize = 512 * 1024; +/// A record bigger than this is sent cut. +const MAX_RECORD: usize = 256 * 1024; +/// Records one attach or history page asks for. +const PAGE: u64 = 200; +/// A guest's prompts held at once, and the decided ones a snapshot names. +const MAX_HELD: usize = 16; +const MAX_DECIDED: usize = 16; +/// Prompt states and client ids one guest's ledger remembers. +const MAX_STATES: usize = 512; +const MAX_CLIENTS: usize = 64; +/// The room's asks a relay keeps until its attachment answers. +const MAX_WAITING: usize = 64; +/// The room's asks queued for a relay, its calls awaiting the daemon's answer, and its history +/// pages among them: past these an ask is refused at once (`share.busy`, an empty page). +pub(super) const MAX_ASKS: usize = 64; +const MAX_PENDING: usize = 32; +const MAX_PAGES: usize = 2; +/// The records the daemon sends before the attach answer, kept for after the snapshot. +const MAX_EARLY: usize = 256; + +/// The prompt states that never change again. +fn decided(state: &str) -> bool { + matches!(state, "approved" | "declined" | "removed") +} + +/// Whether a prompt's state may go from `from` to `to`: forward only. A decision is final; +/// `held` (acpmux took it) is never undone by a refused retry, and replaces an earlier refusal. +fn advance(from: Option<&str>, to: &str) -> bool { + match from { + Some(from) if decided(from) => false, + Some("held") => to != "refused", + _ => true, + } +} + +struct State { + state: &'static str, + reason: Option, + at: u64, +} + +/// One guest's prompts: their states, their client ids, and which count against [`MAX_HELD`]. +/// The bridge keeps one per guest of each share, so the guest's successive relays and the +/// share's successive host sockets share it (a reconnect of either keeps it). Keyed by prompt +/// id: a retry, a decision said again (live, in a snapshot, in a history page) changes nothing. +/// The count is the guest's prompts sent and neither refused nor decided, filled only from what +/// the guest sent (never from `prompt_status`), so the cap counts every one. A prompt the daemon +/// has not answered yet stays here with its text, and every relay of the guest sends it again +/// once attached (acpmux runs a prompt id once), so neither a guest nor a host reconnect loses +/// it unsaid. +#[derive(Default)] +pub(super) struct Ledger { + open: HashSet, + /// Prompts taken and not answered yet (prompt id, client id, text), oldest first; at most + /// [`MAX_HELD`], as each is in `open`. + unanswered: VecDeque<(String, String, String)>, + /// Prompt id to its state (the chat's guest prompts the relays saw), oldest first in `order`. + states: HashMap, + order: VecDeque, + /// Prompt id to this guest's `client_id`. + clients: HashMap, + tick: u64, +} + +pub(super) type Held = Arc>; + +impl Ledger { + fn state(&self, prompt: &str) -> Option<&'static str> { + self.states.get(prompt).map(|s| s.state) + } + + /// Whether `prompt` may go to the daemon: it counts already (sent, or decided), or there is + /// room for it under the cap, which it then takes. + fn admit(&mut self, prompt: &str) -> bool { + if self.open.contains(prompt) || self.state(prompt).is_some_and(decided) { + return true; + } + if self.open.len() >= MAX_HELD { + return false; + } + self.open.insert(prompt.to_owned()); + true + } + + /// Takes the guest's `prompt` to send: false at the cap. Kept until the daemon answers it. + fn take(&mut self, prompt: &str, client_id: String, text: String) -> bool { + self.remember(prompt, client_id.clone()); + if !self.admit(prompt) { + return false; + } + if !self.unanswered.iter().any(|(p, _, _)| p == prompt) { + self.unanswered.push_back((prompt.to_owned(), client_id, text)); + } + true + } + + /// Records a prompt's state, forward only; whether it changed. Any state answers it. + fn note(&mut self, prompt: &str, state: &'static str, reason: Option) -> bool { + if !prompt.starts_with("g-") { + return false; + } + self.unanswered.retain(|(p, _, _)| p != prompt); + let from = self.state(prompt); + if !advance(from, state) { + return false; + } + match state { + // Counted when it is this guest's (its client id is known). + "held" if self.clients.contains_key(prompt) => { + self.open.insert(prompt.to_owned()); + } + "held" => {} + _ => { + self.open.remove(prompt); + } + } + let same = self.states.get(prompt).is_some_and(|s| s.state == state && s.reason == reason); + self.tick = self.tick.wrapping_add(1); + if !self.states.contains_key(prompt) { + if self.order.len() >= MAX_STATES + && let Some(oldest) = self.order.pop_front() + { + self.states.remove(&oldest); + } + self.order.push_back(prompt.to_owned()); + } + self.states.insert(prompt.to_owned(), State { state, reason, at: self.tick }); + !same + } + + fn remember(&mut self, prompt: &str, client_id: String) { + if !self.clients.contains_key(prompt) && self.clients.len() >= MAX_CLIENTS { + let oldest = + self.clients.keys().min_by_key(|p| self.states.get(*p).map_or(0, |s| s.at)); + if let Some(oldest) = oldest.cloned() { + self.clients.remove(&oldest); + } + } + self.clients.insert(prompt.to_owned(), client_id); + } + + /// The `prompt_state` of this guest's `prompt`, when its client id and state are known. + fn frame(&self, guest_id: &str, prompt: &str) -> Option { + let (client_id, state) = (self.clients.get(prompt)?, self.states.get(prompt)?); + let mut frame = json!({"v": 1, "t": "prompt_state", "guest_id": guest_id, + "client_id": client_id, "prompt_id": prompt, "state": state.state}); + if let Some(reason) = &state.reason { + frame.set("reason", json!(reason)); + } + Some(frame) + } + + /// This guest's held prompts and its most recently decided ones. + fn prompts(&self) -> Vec { + let mut known: Vec<(&String, &String, &State)> = (self.clients.iter()) + .filter_map(|(prompt, client)| Some((prompt, client, self.states.get(prompt)?))) + .collect(); + known.sort_by_key(|(_, _, state)| std::cmp::Reverse(state.at)); + let (held, decided): (Vec<_>, Vec<_>) = + known.into_iter().partition(|(_, _, s)| s.state == "held"); + let entry = |(prompt, client, state): (&String, &String, &State)| json!({"client_id": client, "prompt_id": prompt, "state": state.state}); + let held = held.into_iter().take(MAX_HELD).map(&entry); + held.chain(decided.into_iter().take(MAX_DECIDED).map(&entry)).collect() + } +} + +fn lock(held: &Held) -> MutexGuard<'_, Ledger> { + held.lock().unwrap_or_else(|e| e.into_inner()) +} + +/// What the room asked of one guest. +pub(super) enum Ask { + Prompt { client_id: String, text: String }, + Typing(bool), + Older(u64), + Status(Vec), +} + +/// The answer to an ask no relay can take now: a prompt is refused with `reason` (its state +/// said; a prompt acpmux took already stays held), `older` gets an empty page that still has +/// more; the rest the guest's page asks again. +pub(super) fn refuse(held: &Held, ids: (&str, &str), ask: Ask, reason: &str) -> Option { + match ask { + Ask::Prompt { client_id, .. } if (1..=128).contains(&client_id.len()) => { + let prompt = guest_prompt_id(ids.0, ids.1, &client_id); + let mut ledger = lock(held); + ledger.remember(&prompt, client_id); + ledger.note(&prompt, "refused", Some(reason.to_owned())); + ledger.frame(ids.1, &prompt) + } + Ask::Older(_) => Some(bounded(json!({"v": 1, "t": "page", "to": ids.1}), Vec::new(), true)), + Ask::Prompt { .. } | Ask::Typing(_) | Ask::Status(_) => None, + } +} + +/// What a relay hands the host socket. +pub(super) enum Out { + /// A frame of this guest's relay of this generation. + Frame(String, u64, Value), + /// The relay of this guest and generation ended. + Ended(String, u64), + /// Try again to start the relay of this connected guest (no connection permit was free, or + /// its last relay ended while the guest stayed connected). + Retry(String), +} + +/// One running relay, as the host socket keeps it. Dropping it ends the relay. +pub(super) struct Relay { + /// Bounded ([`MAX_ASKS`]): a full queue refuses the ask. + pub(super) asks: mpsc::Sender, + pub(super) generation: u64, + /// The guest's browser sockets the room said are open; `None` once their count is unknown + /// (a host `hello` names each guest once). + pub(super) sockets: Option, + /// The starts in a row this guest's relay took (its backoff when it ends again). + pub(super) tries: u32, +} + +/// Starts the relay of `guest_id` (on the bridge's runtime: the caller is one of its tasks). Its +/// connection takes one of the daemon's connection permits like any client's; with none free no +/// relay starts (the guest's next `guest_joined` asks again). +pub(super) fn start( + mux: &Arc, + render: Arc, + ids: (&str, &str), + generation: u64, + out: mpsc::UnboundedSender, + held: Held, +) -> Option { + let permit = claim_connection(mux.connection_stats())?; + let (ours, daemons) = std::os::unix::net::UnixStream::pair().ok()?; + ours.set_nonblocking(true).ok()?; + let ours = tokio::net::UnixStream::from_std(ours).ok()?; + let mux = mux.clone(); + let serve = move || { + let stream: Box = Box::new(daemons); + handle_connection_with_permit(mux, stream, render, Some(permit)); + }; + std::thread::Builder::new().name("mux-share-guest".into()).spawn(serve).ok()?; + let (asks, receiver) = mpsc::channel(MAX_ASKS); + let (read, write) = ours.into_split(); + let link = Link::new(ids, generation, out, write, held); + tokio::spawn(run(BufReader::new(read), link, receiver)); + Some(Relay { asks, generation, sockets: Some(1), tries: 0 }) +} + +async fn run( + read: BufReader, + mut link: Link, + mut asks: mpsc::Receiver, +) { + let mut lines = read.lines(); + let bind = json!({"cmd": "agent-session-proxy-bind", + "link_guest": {"share_id": link.share_id, "guest_id": link.guest_id}}); + let mut open = link.call(bind, Pending::Bind).await; + while open { + open = tokio::select! { + line = lines.next_line() => match line { + Ok(Some(line)) => link.inbound(&line).await, + _ => false, + }, + // Until the attachment answers, the room's asks wait (one sent before it would be + // refused as not attached, and lost): a prompt in the guest's ledger, which every + // relay of the guest sends once attached, the rest here. + ask = asks.recv() => match ask { + Some(ask) if link.attached => link.ask(ask).await, + Some(Ask::Prompt { client_id, text }) => { + link.take(client_id, text); + true + } + Some(ask) if link.waiting.len() < MAX_WAITING => { + link.waiting.push_back(ask); + true + } + // Too many before the attachment: answered as refused, never lost unsaid. + Some(ask) => { + link.busy(ask); + true + } + None => false, + }, + }; + while open + && link.attached + && let Some(ask) = link.waiting.pop_front() + { + open = link.ask(ask).await; + } + } + // The prompts it sent and did not get answered stay in the guest's ledger: its next relay + // sends them again. + let _ = link.out.send(Out::Ended(link.guest_id.clone(), link.generation)); +} + +enum Pending { + Bind, + Attach, + Page, + /// A prompt sent: its id. + Prompt(String), + Ignore, +} + +struct Link { + share_id: String, + guest_id: String, + generation: u64, + out: mpsc::UnboundedSender, + write: OwnedWriteHalf, + next_id: u64, + /// The calls awaiting the daemon's answer (at most [`MAX_PENDING`] the room's asks made). + pending: HashMap, + /// The attachment answered: the room's asks go to the daemon. + attached: bool, + /// The room's asks that came before the attachment answered, sent once it does. + waiting: VecDeque, + /// The records the daemon sent before the attachment answered, sent after the snapshot. + early: VecDeque, + /// This guest's prompts (kept by the bridge across relays and host sockets). + held: Held, +} + +impl Link { + fn new( + ids: (&str, &str), + generation: u64, + out: mpsc::UnboundedSender, + write: OwnedWriteHalf, + held: Held, + ) -> Self { + Self { + share_id: ids.0.to_owned(), + guest_id: ids.1.to_owned(), + generation, + out, + write, + next_id: 1, + pending: HashMap::new(), + attached: false, + waiting: VecDeque::new(), + early: VecDeque::new(), + held, + } + } + + async fn call(&mut self, mut frame: Value, pending: Pending) -> bool { + let id = self.next_id; + self.next_id = self.next_id.wrapping_add(1); + frame.set("id", json!(id)); + self.pending.insert(id, pending); + let mut line = frame.to_string(); + line.push('\n'); + self.write.write_all(line.as_bytes()).await.is_ok() + } + + fn emit(&self, frame: Value) { + let _ = self.out.send(Out::Frame(self.guest_id.clone(), self.generation, frame)); + } + + fn ledger(&self) -> MutexGuard<'_, Ledger> { + lock(&self.held) + } + + /// One line from the daemon: a reply or an event of the guest's attachment. + async fn inbound(&mut self, line: &str) -> bool { + let Ok(frame) = serde_json::from_str::(line) else { return true }; + match frame.get("event").and_then(Value::as_str) { + Some("agent-session-record") => { + let record = frame.get("record").cloned().unwrap_or(Value::Null); + if self.attached { + self.record(record); + return true; + } + // Before the snapshot nothing reaches the room; too many and the relay ends + // (the guest's next one starts from a fresh snapshot). + self.early.push_back(record); + return self.early.len() <= MAX_EARLY; + } + // share_ended, detached, overflow: the guest's page reconnects for a fresh snapshot. + Some("agent-session-closed") => return false, + Some(_) => return true, + None => {} + } + let Some(pending) = + frame.get("id").and_then(Value::as_u64).and_then(|id| self.pending.remove(&id)) + else { + return true; + }; + let ok = frame.get("ok") == Some(&Value::Bool(true)); + match pending { + Pending::Bind | Pending::Attach if !ok => false, + Pending::Bind => { + let attach = json!({"cmd": "agent-session-attach", "surface": 0, "limit": PAGE, + "kinds": ["transcript"]}); + self.call(attach, Pending::Attach).await + } + Pending::Attach => { + self.attached = true; + let watermark = self.snapshot(frame.get("data").unwrap_or(&Value::Null)); + for record in std::mem::take(&mut self.early) { + if record.get("seq").and_then(Value::as_u64).is_none_or(|s| s > watermark) { + self.record(record); + } + } + let view = json!({"cmd": "agent-session-view", "surface": 0, "visible": true}); + if !self.call(view, Pending::Ignore).await { + return false; + } + // The guest's prompts no daemon answered yet (taken before this attachment, or + // by an earlier relay): sent again. + let unanswered = self.ledger().unanswered.clone(); + for (prompt, client_id, text) in unanswered { + if !self.send(prompt, client_id, text).await { + return false; + } + } + true + } + Pending::Page if ok => { + let page = json!({"v": 1, "t": "page", "to": self.guest_id}); + let data = frame.get("data").unwrap_or(&Value::Null); + let more = data.get("hasMore") == Some(&Value::Bool(true)); + let records = ordered(data); + // The states it names (older than the snapshot: decided while the guest was + // away). Forward only, so an older page never undoes a newer state. + let states: Vec<_> = records.iter().filter_map(decision).collect(); + self.emit(bounded(page, records, more)); + for (prompt, state) in states { + self.note(&prompt, state, None); + } + true + } + Pending::Page => { + self.busy(Ask::Older(0)); + true + } + // Held unless acpmux's records said otherwise first (a retried prompt it ran). + Pending::Prompt(prompt) if ok => { + self.note(&prompt, "held", None); + true + } + Pending::Prompt(prompt) => { + let reason = frame + .get("error_code") + .and_then(Value::as_str) + .unwrap_or("share.refused") + .to_owned(); + self.note(&prompt, "refused", Some(reason)); + true + } + Pending::Ignore => true, + } + } + + /// Takes a prompt of the guest into its ledger (held, plus sent and not answered yet: a + /// burst cannot pass the cap before the replies come; a retry of one counted or decided + /// takes no slot again); its id, or `None` when refused at the cap. + fn take(&mut self, client_id: String, text: String) -> Option { + if !(1..=128).contains(&client_id.len()) { + return None; + } + let prompt = guest_prompt_id(&self.share_id, &self.guest_id, &client_id); + if self.ledger().take(&prompt, client_id, text) { + return Some(prompt); + } + self.note(&prompt, "refused", Some("share.too_many_held".into())); + None + } + + /// Sends a prompt the ledger took, unless this relay sent it and awaits its answer. + async fn send(&mut self, prompt: String, client_id: String, text: String) -> bool { + if self.pending.values().any(|p| matches!(p, Pending::Prompt(p) if *p == prompt)) { + return true; + } + let frame = json!({"cmd": "agent-session-prompt", "surface": 0, + "client_id": client_id, "text": text}); + self.call(frame, Pending::Prompt(prompt)).await + } + + /// An ask the relay cannot take now (`share.busy`). + fn busy(&self, ask: Ask) { + if let Some(frame) = refuse(&self.held, (&self.share_id, &self.guest_id), ask, "share.busy") + { + self.emit(frame); + } + } + + async fn ask(&mut self, ask: Ask) -> bool { + let pages = self.pending.values().filter(|p| matches!(p, Pending::Page)).count(); + let full = self.pending.len() >= MAX_PENDING; + match ask { + Ask::Prompt { client_id, text } => { + if !(1..=128).contains(&client_id.len()) { + return true; + } + if full { + self.busy(Ask::Prompt { client_id, text }); + return true; + } + match self.take(client_id.clone(), text.clone()) { + Some(prompt) => self.send(prompt, client_id, text).await, + None => true, + } + } + // Dropped when the daemon is behind: the next one says the same. + Ask::Typing(_) if full => true, + Ask::Typing(on) => { + let frame = json!({"cmd": "agent-session-typing", "surface": 0, "on": on}); + self.call(frame, Pending::Ignore).await + } + Ask::Older(before) if full || pages >= MAX_PAGES => { + self.busy(Ask::Older(before)); + true + } + Ask::Older(before) => { + let frame = json!({"cmd": "agent-session-events", "surface": 0, + "before_seq": before, "limit": PAGE, "kinds": ["transcript"]}); + self.call(frame, Pending::Page).await + } + // One `prompt_state` per prompt the relay knows the state of. + Ask::Status(client_ids) => { + for client_id in client_ids.into_iter().take(MAX_HELD) { + if (1..=128).contains(&client_id.len()) { + let prompt = guest_prompt_id(&self.share_id, &self.guest_id, &client_id); + let frame = { + let mut ledger = self.ledger(); + ledger.remember(&prompt, client_id); + ledger.frame(&self.guest_id, &prompt) + }; + if let Some(frame) = frame { + self.emit(frame); + } + } + } + true + } + } + } + + /// Sends the snapshot of the attach answer; its newest seq. + fn snapshot(&mut self, data: &Value) -> u64 { + let records = ordered(data); + let prompts = { + let mut ledger = self.ledger(); + for (prompt, state) in records.iter().filter_map(decision) { + ledger.note(&prompt, state, None); + } + ledger.prompts() + }; + let watermark = + records.iter().filter_map(|r| r.get("seq").and_then(Value::as_u64)).max().unwrap_or(0); + let session = data.get("session").unwrap_or(&Value::Null); + let chat = json!({"title": session.get("title").or_else(|| session.get("name")), + "harness": session.get("harness"), "status": session.get("status")}); + let snapshot = json!({"v": 1, "t": "snapshot", "to": self.guest_id, "chat": chat, + "watermark": watermark, "prompts": prompts}); + self.emit(bounded(snapshot, records, data.get("hasMore") == Some(&Value::Bool(true)))); + watermark + } + + fn record(&mut self, record: Value) { + let decided = decision(&record); + let records = + json!({"v": 1, "t": "records", "to": [self.guest_id], "records": [cut(record)]}); + self.emit(records); + if let Some((prompt, state)) = decided { + self.note(&prompt, state, None); + } + } + + /// Records a prompt's state; a change of this guest's prompt goes to the room. + fn note(&mut self, prompt: &str, state: &'static str, reason: Option) { + let frame = { + let mut ledger = self.ledger(); + if ledger.note(prompt, state, reason) { + ledger.frame(&self.guest_id, prompt) + } else { + None + } + }; + if let Some(frame) = frame { + self.emit(frame); + } + } +} + +/// The page's records, oldest first. +fn ordered(data: &Value) -> Vec { + let mut records = data.get("events").and_then(Value::as_array).cloned().unwrap_or_default(); + records.sort_by_key(|record| record.get("seq").and_then(Value::as_u64).unwrap_or(0)); + records +} + +/// A guest prompt's state a record says: its prompt id and state. +fn decision(record: &Value) -> Option<(String, &'static str)> { + let msg = record.get("msg").unwrap_or(record); + let prompt = msg.get("promptId").and_then(Value::as_str)?; + let state = match record.get("kind").and_then(Value::as_str)? { + "queued" if msg.get("awaitingOwner") == Some(&Value::Bool(true)) => "held", + "queue_approved" => "approved", + "queue_declined" => "declined", + "queue_removed" | "prompt_retracted" => "removed", + _ => return None, + }; + Some((prompt.to_owned(), state)) +} + +/// A record over [`MAX_RECORD`] as its seq, dir and kind, `truncated` with its size. +fn cut(record: Value) -> Value { + let full = record.to_string().len(); + if full <= MAX_RECORD { + return record; + } + let mut short = json!({"truncated": true, "full_bytes": full}); + for key in ["seq", "dir", "kind", "ts", "sessionId"] { + if let Some(value) = record.get(key) { + short.set(key, value.clone()); + } + } + short +} + +/// `frame` with the newest of `records` (oldest first) that fit in [`MAX_FRAME`] with it, each +/// cut; `has_more` when `more` or any older record was left out. +fn bounded(mut frame: Value, records: Vec, more: bool) -> Value { + let records: Vec = records.into_iter().map(cut).collect(); + frame.set("records", json!([])); + frame.set("has_more", json!(false)); + let mut total = frame.to_string().len(); + let mut from = records.len(); + for (index, record) in records.iter().enumerate().rev() { + let size = record.to_string().len().saturating_add(1); + if total.saturating_add(size) > MAX_FRAME { + break; + } + total = total.saturating_add(size); + from = index; + } + frame.set("has_more", json!(more || from > 0)); + frame.set("records", Value::Array(records.into_iter().skip(from).collect())); + frame +} + +#[cfg(test)] +mod tests { + use super::*; + + fn link(out: mpsc::UnboundedSender, write: OwnedWriteHalf) -> Link { + let mut link = Link::new(("s", "g"), 1, out, write, Held::default()); + link.next_id = 2; + link + } + /// A record the daemon sends before the attach answer reaches the room after the snapshot. + #[test] + fn records_before_the_attach_answer_follow_the_snapshot() { + let runtime = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap(); + runtime.block_on(async { + let (ours, _theirs) = tokio::net::UnixStream::pair().unwrap(); + let (_read, write) = ours.into_split(); + let (out, mut frames) = mpsc::unbounded_channel(); + let mut link = link(out, write); + link.pending.insert(1, Pending::Attach); + let record = json!({"event": "agent-session-record", + "record": {"seq": 5, "kind": "agent_message"}}); + assert!(link.inbound(&record.to_string()).await); + let reply = json!({"id": 1, "ok": true, "data": {"events": [{"seq": 4}]}}); + assert!(link.inbound(&reply.to_string()).await); + let mut kinds = Vec::new(); + while let Ok(Out::Frame(_, _, frame)) = frames.try_recv() { + kinds.push(frame["t"].clone()); + } + assert_eq!(kinds, vec![json!("snapshot"), json!("records")]); + }); + } +} diff --git a/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_room.rs b/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_room.rs index 172812a3ebe4..cf368b32d732 100644 --- a/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_room.rs +++ b/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_room.rs @@ -9,10 +9,12 @@ //! with the change: the kick is resent on every host connect until the room answers `kicked`, and //! the `DELETE` is retried until it is answered (or the share's expiry ends it at the room). //! -//! Everything waits on events: the socket, a command channel, and one-shot sleeps for backoff -//! (1 s doubling to 60 s), run only while something is pending. The relay of each guest's chat -//! (snapshot, records, prompts, presence) comes next; until then the room's guest frames -//! (`guest_joined`, `guest_prompt`, `older`, ...) go unanswered. +//! Each guest whose browser is in the room gets a relay of the chat (`agent_session_share_relay.rs`) +//! from the room's `hello` or `guest_joined` until its `guest_left`, a kick or the share's end; a +//! host or guest reconnect starts it again, with a fresh snapshot. +//! +//! Everything waits on events: the sockets, a command channel, and one-shot sleeps for backoff +//! (1 s doubling to 60 s), run only while something is pending. use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; @@ -30,6 +32,10 @@ use zeroize::Zeroizing; use super::{GuestState, LinkShare, LinkShares, Redeem, ShareState, find, op_of}; use crate::mux::Mux; +use crate::server::RenderService; + +#[path = "agent_session_share_relay.rs"] +mod relay; type Ws = tokio_tungstenite::WebSocketStream>; @@ -83,6 +89,11 @@ pub(in crate::server) struct Bridge { runtime: OnceLock>, hosts: Mutex)>>, next_host: AtomicU64, + /// The guest relays' connections' outbound budget. + render: OnceLock>, + /// Each share's guests' held prompts (share id, then guest id), kept across the share's host + /// sockets and the guests' relays until the guest is kicked or the share ends. + held: Mutex>>, } impl Drop for Bridge { @@ -108,6 +119,27 @@ fn origin(raw: &str) -> Option { } impl Bridge { + /// The held prompts of `guest` of `share_id`. + fn held(&self, share_id: &str, guest: &str) -> relay::Held { + let mut held = self.held.lock().unwrap_or_else(|e| e.into_inner()); + held.entry(share_id.to_owned()).or_default().entry(guest.to_owned()).or_default().clone() + } + + /// Forgets the held prompts of a kicked guest, or (`None`) of an ended share's guests. + fn forget(&self, share_id: &str, guest: Option<&str>) { + let mut held = self.held.lock().unwrap_or_else(|e| e.into_inner()); + match guest { + Some(guest) => { + if let Some(guests) = held.get_mut(share_id) { + guests.remove(guest); + } + } + None => { + held.remove(share_id); + } + } + } + /// Sets the room new shares are provisioned in (`None`: links have no URL). False when `url` /// is no https origin (including an explicit empty string): creation fails with /// `share.room_unreachable`, while an unset room preserves local sharing. @@ -635,6 +667,7 @@ fn host(mux: &Arc, share_id: &str) { pub(super) fn ended(shares: &LinkShares, ended: &[(String, Option)]) { let Some(mux) = shares.mux.lock().unwrap_or_else(|e| e.into_inner()).upgrade() else { return }; for (share_id, guest) in ended { + shares.bridge.forget(share_id, guest.as_deref()); if guest.is_some() { // An ended share's own entry sends `end`; its kicks are moot. if view(&mux, share_id).is_some_and(|v| v.state == ShareState::Active) { @@ -788,6 +821,8 @@ async fn session( let mut pinged = false; // The reason of the room's own `end`, if it sent one. let mut room_end: Option = None; + let (relay_out, mut relayed) = mpsc::unbounded_channel(); + let mut guests = Guests::new(relay_out); loop { let quiet = heard + if pinged { timing.heartbeat * 2 } else { timing.heartbeat }; let out = tokio::select! { @@ -801,14 +836,17 @@ async fn session( frame = stream.next() => match frame { Some(Ok(Message::Text(text))) => { (heard, pinged) = (tokio::time::Instant::now(), false); + let frame: Value = serde_json::from_str(text.as_str()).unwrap_or(Value::Null); // The room says why it ends the share before it closes with 4001. - if let Ok(frame) = serde_json::from_str::(text.as_str()) - && frame.get("t").and_then(Value::as_str) == Some("end") - { + if frame.get("t").and_then(Value::as_str) == Some("end") { room_end = frame.get("reason").and_then(Value::as_str).map(str::to_owned); } let mux = mux.upgrade()?; - answer(&mux, share_id, text.as_str(), &mut ready) + let mut out = answer(&mux, share_id, &frame, &mut ready); + if ready { + out.extend(guests.room(&mux, share_id, &frame)); + } + out } Some(Ok(Message::Close(frame))) => { let code = frame.map(|frame| u16::from(frame.code)); @@ -859,6 +897,25 @@ async fn session( } None => return None, }, + relayed = relayed.recv() => match relayed { + Some(relay::Out::Frame(guest, generation, frame)) => { + if !guests.current(&guest, generation) { + continue; + } + vec![frame] + } + Some(relay::Out::Ended(guest, generation)) => { + let mux = mux.upgrade()?; + guests.ended(&mux, share_id, &guest, generation); + continue; + } + Some(relay::Out::Retry(guest)) => { + let mux = mux.upgrade()?; + guests.retry(&mux, share_id, &guest); + continue; + } + None => continue, + }, }; for frame in out { if !put(&mut sink, Message::text(frame.to_string())).await { @@ -880,6 +937,196 @@ fn reconcile(mux: &Mux, share_id: &str) -> Vec { out } +/// A guest relay that could not start, or ended while its guest stayed connected, starts again +/// after this, doubling up to [`RELAY_RETRY_MAX`]. +const RELAY_RETRY_MIN: Duration = Duration::from_millis(250); +const RELAY_RETRY_MAX: Duration = Duration::from_secs(30); + +/// A connected guest with no relay running: one is started again after a backoff. +struct Waiting { + sockets: Option, + tries: u32, +} + +/// The relays of one host socket's guests (dropped with it). +struct Guests { + relays: HashMap, + /// Connected guests whose relay could not start (no connection permit) or ended. + waiting: HashMap, + next: u64, + out: mpsc::UnboundedSender, +} + +impl Guests { + fn new(out: mpsc::UnboundedSender) -> Self { + Self { relays: HashMap::new(), waiting: HashMap::new(), next: 0, out } + } + + /// A room frame about the guests' browsers, after the daemon's `guests`; what to send the + /// room back (a prompt no relay can take is refused, `share.busy`). + fn room(&mut self, mux: &Arc, share_id: &str, frame: &Value) -> Vec { + let guest = frame.get("guest_id").and_then(Value::as_str).unwrap_or_default(); + let ask = match frame.get("t").and_then(Value::as_str).unwrap_or_default() { + // `hello` names each connected guest once; its socket count when the room gives it + // (`sockets`), else unknown. + "hello" => { + self.relays.clear(); + self.waiting.clear(); + for guest in frame.get("guests").and_then(Value::as_array).into_iter().flatten() { + let sockets = guest.get("sockets").and_then(Value::as_u64); + let sockets = sockets.and_then(|n| usize::try_from(n).ok()).filter(|n| *n > 0); + let guest = guest.get("guest_id").and_then(Value::as_str).unwrap_or_default(); + self.open(mux, share_id, guest, sockets, 0); + } + return Vec::new(); + } + // A guest's (re)connect gets a fresh snapshot; one more socket than before. + "guest_joined" => { + let before = match (self.relays.remove(guest), self.waiting.remove(guest)) { + (Some(relay), _) => Some(relay.sockets), + (None, Some(waiting)) => Some(waiting.sockets), + (None, None) => None, + }; + let sockets = match before { + Some(sockets) => sockets.map(|n| n.saturating_add(1)), + None => Some(1), + }; + self.open(mux, share_id, guest, sockets, 0); + return Vec::new(); + } + // The relay ends with the guest's last socket. With no count (a room whose `hello` + // gives none) the first `guest_left` ends it: a relay never outlives its guest's + // sockets holding a daemon connection; another open tab of the guest reconnects. + "guest_left" => { + let sockets = match (self.relays.get_mut(guest), self.waiting.get_mut(guest)) { + (Some(relay), _) => Some(&mut relay.sockets), + (None, Some(waiting)) => Some(&mut waiting.sockets), + (None, None) => None, + }; + if let Some(sockets) = sockets { + let left = sockets.map_or(0, |n| n.saturating_sub(1)); + *sockets = Some(left); + if left == 0 { + self.relays.remove(guest); + self.waiting.remove(guest); + } + } + return Vec::new(); + } + "guest_prompt" => { + let client_id = frame.get("client_id").and_then(Value::as_str).unwrap_or_default(); + let text = frame.get("text").and_then(Value::as_str).unwrap_or_default(); + let ask = + relay::Ask::Prompt { client_id: client_id.to_owned(), text: text.to_owned() }; + // A connected guest whose relay waits to start: refused, said, not lost. + if self.waiting.contains_key(guest) { + return refused(mux, share_id, guest, ask); + } + ask + } + "guest_typing" => relay::Ask::Typing(frame.get("on") == Some(&Value::Bool(true))), + "older" => match frame.get("before_seq").and_then(Value::as_u64) { + Some(before) => relay::Ask::Older(before), + None => return Vec::new(), + }, + "prompt_status" => { + let ids: Vec = + (frame.get("client_ids").and_then(Value::as_array).into_iter().flatten()) + .filter_map(|id| id.as_str().map(str::to_owned)) + .collect(); + relay::Ask::Status(ids) + } + _ => return Vec::new(), + }; + self.send(mux, share_id, guest, ask) + } + + /// Queues `ask` for `guest`'s relay; with its queue full, refused at once. + fn send(&self, mux: &Mux, share_id: &str, guest: &str, ask: relay::Ask) -> Vec { + let Some(relay) = self.relays.get(guest) else { return Vec::new() }; + match relay.asks.try_send(ask) { + Ok(()) => Vec::new(), + Err(full) => refused(mux, share_id, guest, full.into_inner()), + } + } + + /// Starts the relay of a guest the daemon holds active, with its sockets (`None`: unknown). + /// With no connection permit free the guest waits, and it is tried again after a backoff. + fn open( + &mut self, + mux: &Arc, + share_id: &str, + guest: &str, + sockets: Option, + tries: u32, + ) { + self.waiting.remove(guest); + if !view(mux, share_id).is_some_and(|v| v.active.iter().any(|g| g == guest)) { + return; + } + let bridge = &shares(mux).bridge; + let render = bridge.render.get_or_init(|| Arc::new(RenderService::new())).clone(); + self.next = self.next.wrapping_add(1); + let ids = (share_id, guest); + let held = bridge.held(share_id, guest); + let out = self.out.clone(); + match relay::start(mux, render, ids, self.next, out, held) { + Some(mut relay) => { + relay.sockets = sockets; + relay.tries = tries; + self.relays.insert(guest.to_owned(), relay); + } + None => self.wait(guest, sockets, tries), + } + } + + /// `guest` stays connected with no relay: it is started again after a backoff. + fn wait(&mut self, guest: &str, sockets: Option, tries: u32) { + self.waiting.insert(guest.to_owned(), Waiting { sockets, tries }); + let factor = 2_u32.saturating_pow(tries.min(16)); + let delay = RELAY_RETRY_MIN.saturating_mul(factor).min(RELAY_RETRY_MAX); + let (out, guest) = (self.out.clone(), guest.to_owned()); + tokio::spawn(async move { + tokio::time::sleep(delay).await; + let _ = out.send(relay::Out::Retry(guest)); + }); + } + + /// A waiting guest's backoff ran out: its relay is started again. + fn retry(&mut self, mux: &Arc, share_id: &str, guest: &str) { + if let Some(waiting) = self.waiting.remove(guest) { + self.open(mux, share_id, guest, waiting.sockets, waiting.tries.saturating_add(1)); + } + } + + /// Whether `generation` is `guest`'s current relay: a replaced relay's late frames are dropped. + fn current(&self, guest: &str, generation: u64) -> bool { + self.relays.get(guest).is_some_and(|relay| relay.generation == generation) + } + + /// A relay ended (its attachment closed: detached, overflow, the guest kicked). The current + /// one of a guest still connected starts again after a backoff; a kicked guest's does not + /// (`open` checks the guest is active). + fn ended(&mut self, mux: &Mux, share_id: &str, guest: &str, generation: u64) { + if !self.current(guest, generation) { + return; + } + let active = view(mux, share_id).is_some_and(|v| v.active.iter().any(|g| g == guest)); + if let Some(relay) = self.relays.remove(guest) + && active + && relay.sockets != Some(0) + { + self.wait(guest, relay.sockets, relay.tries.saturating_add(1)); + } + } +} + +/// What the room is told of an ask for `guest` no relay can take (`share.busy`). +fn refused(mux: &Mux, share_id: &str, guest: &str, ask: relay::Ask) -> Vec { + let held = shares(mux).bridge.held(share_id, guest); + relay::refuse(&held, (share_id, guest), ask, "share.busy").into_iter().collect() +} + /// The redeem refusals the room may pass to a guest. const GUEST_REFUSALS: [&str; 8] = [ "share.secret", @@ -893,8 +1140,7 @@ const GUEST_REFUSALS: [&str; 8] = [ ]; /// The daemon's answer to one room frame. -fn answer(mux: &Mux, share_id: &str, text: &str, ready: &mut bool) -> Vec { - let frame: Value = serde_json::from_str(text).unwrap_or(Value::Null); +fn answer(mux: &Mux, share_id: &str, frame: &Value, ready: &mut bool) -> Vec { if frame.get("v").and_then(Value::as_u64) != Some(1) { return Vec::new(); } diff --git a/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_room_tests.rs b/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_room_tests.rs index a782e19b4ce4..036db4fea9a3 100644 --- a/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_room_tests.rs +++ b/cmux-tui/crates/cmux-tui-core/src/server/agent_session_share_room_tests.rs @@ -4,7 +4,7 @@ use super::super::{MAX_SHARES, digest, hex}; use super::*; use crate::server::agent_session_attach_tests::{ - Client, FakeAcpmux, SESSION, WAIT, agent_tab, assert_code, mux_with, + Client, FakeAcpmux, Fixture, SESSION, WAIT, agent_tab, assert_code, mux_with, }; use std::cell::RefCell; use std::collections::{HashSet, VecDeque}; @@ -69,6 +69,8 @@ struct Shared { delete_closes_host: AtomicBool, /// Host sockets neither read nor write while set, as a connection lost without a close. silent: AtomicBool, + /// The guests the next `hello` names as connected. + hello: Mutex>, } struct FakeRoom { @@ -299,7 +301,8 @@ fn serve_host(stream: TcpStream, shared: &Shared, seen: &Sender) { *shared.to_host.lock().unwrap() = Some(outbox); let _ = seen.send(Seen::Host { conn, bearer }); socket.get_mut().set_read_timeout(Some(Duration::from_millis(10))).unwrap(); - let hello = json!({"v": 1, "t": "hello", "guests": []}).to_string(); + let guests = shared.hello.lock().unwrap().clone(); + let hello = json!({"v": 1, "t": "hello", "guests": guests}).to_string(); let _ = socket.send(Message::text(hello)); loop { if shared.silent.load(Ordering::SeqCst) { @@ -365,7 +368,15 @@ fn fast() -> Timing { /// A daemon with one agent chat, bridged to `room`. fn bridged(label: &str, room: &FakeRoom) -> (FakeAcpmux, Arc, u64, Client) { - let fake = FakeAcpmux::start(label); + bridged_with(label, room, Fixture::default()) +} + +fn bridged_with( + label: &str, + room: &FakeRoom, + fixture: Fixture, +) -> (FakeAcpmux, Arc, u64, Client) { + let fake = FakeAcpmux::start_with(label, fixture); let mux = mux_with(&fake, label); assert!(shares(&mux).bridge.configure(Some(&room.url), fast())); let surface = agent_tab(&mux, Some(SESSION)); @@ -1352,3 +1363,561 @@ fn a_room_end_stands_when_its_write_fails() { } shares_of(&mux).refuse_writes.store(false, Ordering::Relaxed); } + +// MARK: The guest relay + +impl FakeRoom { + /// The next frame of type `t` the daemon sent on host socket `conn`. + fn next(&self, conn: u64, t: &str) -> Value { + let wanted = + |s: &Seen| matches!(s, Seen::Frame { conn: c, frame } if *c == conn && frame["t"] == t); + match self.take(t, wanted) { + Seen::Frame { frame, .. } => frame, + _ => unreachable!(), + } + } +} + +/// A transcript record of `bytes` bytes of text. +fn sized(seq: u64, bytes: usize) -> Value { + json!({"sessionId": SESSION, "seq": seq, "dir": "in", "kind": "agent_message", + "msg": {"text": "x".repeat(bytes)}}) +} + +/// A share with two redeemed guests and its host socket past `guests`: (share, [guest ids], conn). +fn two_guests(room: &FakeRoom, owner: &mut Client, surface: u64) -> (Value, [String; 2], u64) { + let reply = owner.request(json!({"cmd":"agent-session-share-create", "surface":surface, + "owner_token":OWNER_TOKEN, "uses":2})); + let share = reply["data"].clone(); + let (conn, _) = room.host(); + assert_eq!(room.next(conn, "guests")["active"], json!([])); + let mut guests = [String::new(), String::new()]; + for (index, name) in ["Ann", "Bob"].into_iter().enumerate() { + let reply = owner.request(json!({"cmd":"agent-session-share-redeem", + "share_id":share["share_id"], "secret":share["secret"], "name":name, + "attempt_id":format!("attempt-{index}")})); + guests[index] = reply["data"]["guest_id"].as_str().unwrap().to_owned(); + } + (share, guests, conn) +} + +#[test] +fn a_guest_gets_a_bounded_snapshot_its_own_records_and_its_prompts_states() { + let room = FakeRoom::start(&[]); + // 200 KiB, 200 KiB, a 300 KiB record (sent cut) and 200 KiB: the newest that fit in 512 KiB. + let page = vec![ + sized(10, 200 << 10), + sized(11, 200 << 10), + sized(12, 300 << 10), + sized(13, 200 << 10), + ]; + let history = vec![sized(3, 10), sized(4, 10)]; + let fixture = Fixture { page, history, ..Fixture::default() }; + let (mut fake, _mux, surface, mut owner) = bridged_with("relay-snap", &room, fixture); + let (_, [ann, bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + let snapshot = room.next(conn, "snapshot"); + assert!(snapshot.to_string().len() <= 512 << 10, "the whole frame fits"); + assert_eq!((snapshot["to"].clone(), snapshot["has_more"].clone()), (json!(ann), json!(true))); + let seqs: Vec = snapshot["records"] + .as_array() + .unwrap() + .iter() + .map(|r| r["seq"].as_u64().unwrap()) + .collect(); + assert_eq!(seqs, vec![11, 12, 13], "the newest records, oldest first"); + let cut = &snapshot["records"][1]; + assert_eq!( + (cut["truncated"].clone(), cut["kind"].clone()), + (json!(true), json!("agent_message")) + ); + assert!(cut["full_bytes"].as_u64().unwrap() > 300 << 10 && cut.get("msg").is_none()); + assert_eq!( + (snapshot["watermark"].clone(), snapshot["chat"]["status"].clone()), + (json!(13), json!("running")) + ); + // The record acpmux sent while the attach was answered follows the snapshot, to Ann only. + let early = room.next(conn, "records"); + assert_eq!((early["to"].clone(), early["records"][0]["seq"].clone()), (json!([ann]), json!(2))); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + // Ann's prompt reaches acpmux as hers and held; its states follow acpmux's records. + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "c1", + "text": "hello"})); + let sent = fake.frame(); + assert_eq!(sent["method"], "session/prompt"); + let prompt = sent["params"]["_meta"]["acpmux"]["promptId"].as_str().unwrap().to_owned(); + assert_eq!(sent["params"]["_meta"]["acpmux"]["author"]["approval"], "owner"); + let state = |state: &str| { + json!({"v": 1, "t": "prompt_state", "guest_id": ann, + "client_id": "c1", "prompt_id": prompt, "state": state}) + }; + assert_eq!(room.next(conn, "prompt_state"), state("held")); + for (seq, kind, extra) in [ + (20, "queued", json!({"awaitingOwner": true, "revision": 0})), + (21, "queue_approved", json!({"revision": 0})), + ] { + let mut msg = json!({"promptId": prompt}); + msg.as_object_mut().unwrap().extend(extra.as_object().unwrap().clone()); + fake.notify( + "_acpmux/event", + json!({"sessionId": SESSION, "seq": seq, "dir": "mux", + "kind": kind, "msg": msg}), + ); + let records = room.next(conn, "records"); + assert_eq!( + (records["to"].clone(), records["records"][0]["seq"].clone()), + (json!([ann]), json!(seq)) + ); + } + assert_eq!(room.next(conn, "prompt_state"), state("approved"), "held was already said"); + room.send(json!({"v": 1, "t": "prompt_status", "guest_id": ann, "client_ids": ["c1", "c9"]})); + assert_eq!(room.next(conn, "prompt_state"), state("approved")); + room.send(json!({"v": 1, "t": "older", "guest_id": ann, "before_seq": 11})); + let page = room.next(conn, "page"); + assert_eq!((page["to"].clone(), page["has_more"].clone()), (json!(ann), json!(false))); + assert_eq!(page["records"].as_array().unwrap().len(), 2); + // Nothing went to Bob, whose browser never joined, nor about a prompt it never sent. + room.none("a frame for Bob", |s| { + matches!(s, Seen::Frame { frame, .. } + if frame.to_string().contains(&bob)) + }); + room.none("a second state", |s| { + matches!(s, Seen::Frame { frame, .. } + if frame["t"] == "prompt_state") + }); +} + +#[test] +fn reconnects_get_fresh_snapshots_and_a_kick_ends_the_guests_relay() { + let room = FakeRoom::start(&[]); + let (fake, _mux, surface, mut owner) = bridged("relay-again", &room); + let (share, [ann, bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + // A guest reconnect: a fresh snapshot. + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + // A host reconnect: `guests` first, then a snapshot for each guest still connected. + room.shared + .hello + .lock() + .unwrap() + .extend([json!({"guest_id": ann, "name": "Ann"}), json!({"guest_id": bob, "name": "Bob"})]); + room.drop_host(); + let (conn, _) = room.host(); + assert_eq!(room.frame(conn), json!({"v": 1, "t": "guests", "active": [ann, bob]})); + let mut to: Vec = (0..2).map(|_| room.next(conn, "snapshot")["to"].clone()).collect(); + to.sort_by_key(Value::to_string); + let mut both = vec![json!(ann), json!(bob)]; + both.sort_by_key(Value::to_string); + assert_eq!(to, both); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + // A kick ends Ann's relay at the daemon: her prompts reach nothing. + let kick = json!({"cmd":"agent-session-share-kick", "share_id":share["share_id"], + "guest_id":ann}); + assert_eq!(owner.request(kick)["ok"], true); + assert_eq!(room.next(conn, "kick")["guest_id"], json!(ann)); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "c1", + "text": "hello"})); + fake.no_frame(); + // Bob's still does. + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": bob, "client_id": "c1", + "text": "hello"})); + assert_eq!(fake.frame()["method"], "session/prompt"); + assert_eq!(room.next(conn, "prompt_state")["guest_id"], json!(bob)); +} + +/// A prompt the room sends right after `guest_joined` (before the relay's attachment answers) +/// still reaches the daemon; after a host reconnect, the relay counts the guest's sockets from +/// `hello`, so one of two leaving keeps it. +#[test] +fn an_early_prompt_waits_for_the_attachment_and_a_hello_guest_outlives_one_left_socket() { + let room = FakeRoom::start(&[]); + let (fake, _mux, surface, mut owner) = bridged("relay-early", &room); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "early", + "text": "hello"})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + let deadline = Instant::now() + WAIT; + loop { + let frame = fake.next_frame(deadline.saturating_duration_since(Instant::now())); + let frame = frame.expect("the early prompt reaches the daemon"); + if frame["method"] == "session/prompt" { + break; + } + } + assert_eq!(room.next(conn, "prompt_state")["client_id"], "early"); + // A host reconnect names Ann with her two open sockets: one leaving keeps her relay. + room.shared.hello.lock().unwrap().push(json!({"guest_id": ann, "name": "Ann", "sockets": 2})); + room.drop_host(); + let (conn, _) = room.host(); + assert_eq!(room.next(conn, "guests")["t"], "guests"); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + room.send(json!({"v": 1, "t": "guest_left", "guest_id": ann})); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "after", + "text": "still here"})); + assert_eq!(fake.frame()["method"], "session/prompt"); +} + +/// A guest relay's connection takes one of the daemon's connection permits like any client's: +/// with none free no relay starts, and the guest's next join gets one once a permit is back. +#[test] +fn a_guest_relay_takes_a_connection_permit() { + let room = FakeRoom::start(&[]); + let (_fake, mux, surface, mut owner) = bridged("relay-permit", &room); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + let mut held = Vec::new(); + while let Some(permit) = crate::server::claim_connection(mux.connection_stats()) { + held.push(permit); + } + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + room.none( + "a snapshot without a free connection permit", + |s| matches!(s, Seen::Frame { frame, .. } if frame["t"] == "snapshot"), + ); + // With a permit back, the waiting guest's relay starts on its own (no new join). + drop(held); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); +} + +/// A burst of prompts cannot pass the 16-held cap before their replies come: prompts sent and +/// not answered count too, and the 17th is refused. +#[test] +fn a_burst_of_prompts_stops_at_the_held_cap() { + let room = FakeRoom::start(&[]); + let (fake, _mux, surface, mut owner) = bridged("relay-burst", &room); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + for n in 0..16 { + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, + "client_id": format!("c{n}"), "text": "hello"})); + } + // Client ids a page names in `prompt_status` do not push the held prompts out of the count. + let fakes: Vec = (0..80).map(|n| format!("fake{n}")).collect(); + room.send(json!({"v": 1, "t": "prompt_status", "guest_id": ann, "client_ids": fakes})); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "c16", + "text": "hello"})); + let deadline = Instant::now() + WAIT; + let refused = loop { + let state = room.next(conn, "prompt_state"); + if state["state"] == "refused" { + break state; + } + assert!(Instant::now() < deadline, "the 17th prompt is refused"); + }; + assert_eq!( + (refused["client_id"].clone(), refused["reason"].clone()), + (json!("c16"), json!("share.too_many_held")) + ); +} + +/// A room whose `hello` gives no socket count: the guest's first `guest_left` ends its relay, so +/// no relay holds a daemon connection past its guest's sockets. +#[test] +fn a_hello_guest_without_a_count_ends_with_its_first_left_socket() { + let room = FakeRoom::start(&[]); + let (fake, _mux, surface, mut owner) = bridged("relay-uncounted", &room); + let (_share, [ann, _bob], _conn) = two_guests(&room, &mut owner, surface); + room.shared.hello.lock().unwrap().push(json!({"guest_id": ann, "name": "Ann"})); + room.drop_host(); + let (conn, _) = room.host(); + assert_eq!(room.next(conn, "guests")["t"], "guests"); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + room.send(json!({"v": 1, "t": "guest_left", "guest_id": ann})); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "gone", + "text": "nobody relays this"})); + fake.no_frame(); +} + +/// Ann's prompts `c{from}..c{to}`, each sent and said held: their acpmux prompt ids. +fn held_prompts( + room: &FakeRoom, + fake: &FakeAcpmux, + conn: u64, + ann: &str, + n: std::ops::Range, +) -> Vec { + let mut prompts = Vec::new(); + for n in n { + let client_id = format!("c{n}"); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": client_id, + "text": "hello"})); + let sent = loop { + let frame = fake.frame(); + if frame["method"] == "session/prompt" { + break frame; + } + }; + prompts.push(sent["params"]["_meta"]["acpmux"]["promptId"].as_str().unwrap().to_owned()); + let state = room.next(conn, "prompt_state"); + let said = (state["client_id"].clone(), state["state"].clone()); + assert_eq!(said, (json!(client_id), json!("held"))); + } + prompts +} + +/// The state the room is next told of Ann's prompt `client_id` (states of others skipped). +fn state_of(room: &FakeRoom, conn: u64, client_id: &str) -> Value { + let deadline = Instant::now() + WAIT; + loop { + let state = room.next(conn, "prompt_state"); + if state["client_id"] == client_id { + return state; + } + assert!(Instant::now() < deadline, "a state of {client_id}"); + } +} + +/// acpmux's record that the owner approved `prompt`. +fn approved(seq: u64, prompt: &str) -> Value { + json!({"sessionId": SESSION, "seq": seq, "dir": "mux", "kind": "queue_approved", + "msg": {"promptId": prompt, "revision": 0}}) +} + +/// A host reconnect keeps the count of the guest's held prompts: 16 held before it, the 17th +/// after it is refused. +#[test] +fn a_host_reconnect_keeps_the_held_cap() { + let room = FakeRoom::start(&[]); + let (fake, _mux, surface, mut owner) = bridged("relay-cap-host", &room); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + held_prompts(&room, &fake, conn, &ann, 0..16); + room.shared.hello.lock().unwrap().push(json!({"guest_id": ann, "name": "Ann", "sockets": 1})); + room.drop_host(); + let (conn, _) = room.host(); + assert_eq!(room.next(conn, "guests")["t"], "guests"); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "c16", + "text": "hello"})); + let state = state_of(&room, conn, "c16"); + assert_eq!( + (state["state"].clone(), state["reason"].clone()), + (json!("refused"), json!("share.too_many_held")) + ); +} + +/// A retried prompt the owner already decided is not counted again: 16 decided and retried +/// leave room for a new one. +#[test] +fn retrying_decided_prompts_takes_no_held_slot() { + let room = FakeRoom::start(&[]); + let (mut fake, _mux, surface, mut owner) = bridged("relay-cap-retry", &room); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + let prompts = held_prompts(&room, &fake, conn, &ann, 0..16); + for (n, prompt) in prompts.iter().enumerate() { + fake.notify("_acpmux/event", approved(40 + n as u64, prompt)); + assert_eq!(state_of(&room, conn, &format!("c{n}"))["state"], "approved"); + } + // Each retried (its page lost the reply): acpmux's ledger runs none of them again. One at a + // time: the attachment takes 8 calls at once (`agent_session.limit`). + for n in 0..16 { + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, + "client_id": format!("c{n}"), "text": "hello"})); + while fake.frame()["method"] != "session/prompt" {} + } + // A new prompt sent while a retry's answer is still out may be refused for that too, never + // for the held cap. + let deadline = Instant::now() + WAIT; + let state = loop { + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "c16", + "text": "hello"})); + let state = state_of(&room, conn, "c16"); + if state["reason"] != "agent_session.limit" || Instant::now() > deadline { + break state; + } + std::thread::sleep(Duration::from_millis(50)); + }; + assert_eq!(state["state"], "held", "{state}"); +} + +/// Decisions a history page names release the held slots of prompts decided while the guest was +/// away (their records older than its reconnect's snapshot), and the guest is told of them. +#[test] +fn a_history_page_applies_the_decisions_it_names() { + let room = FakeRoom::start(&[]); + let fixture = Fixture::default(); + let later = fixture.history_later.clone(); + let (fake, _mux, surface, mut owner) = bridged_with("relay-cap-page", &room, fixture); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + let prompts = held_prompts(&room, &fake, conn, &ann, 0..16); + // Ann leaves; the owner decides all 16; Ann comes back to a snapshot that names none of it. + room.send(json!({"v": 1, "t": "guest_left", "guest_id": ann})); + let decided = prompts.iter().enumerate().map(|(n, p)| approved(40 + n as u64, p)); + later.lock().unwrap().extend(decided); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + room.send(json!({"v": 1, "t": "older", "guest_id": ann, "before_seq": 3})); + assert_eq!(room.next(conn, "page")["to"], json!(ann)); + assert_eq!(state_of(&room, conn, "c0")["state"], "approved"); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "c16", + "text": "hello"})); + assert_eq!(state_of(&room, conn, "c16")["state"], "held"); +} + +/// A retry the daemon takes after an earlier send of the prompt was refused is said held. +#[test] +fn a_retry_taken_after_a_refusal_is_held() { + let room = FakeRoom::start(&[]); + let (mut fake, _mux, surface, mut owner) = bridged("relay-retry-held", &room); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + let prompts = held_prompts(&room, &fake, conn, &ann, 0..16); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "c16", + "text": "hello"})); + assert_eq!(state_of(&room, conn, "c16")["reason"], "share.too_many_held"); + fake.notify("_acpmux/event", approved(40, &prompts[0])); + assert_eq!(state_of(&room, conn, "c0")["state"], "approved"); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "c16", + "text": "hello"})); + let state = state_of(&room, conn, "c16"); + assert_eq!(state["state"], "held", "{state}"); +} + +/// A retry of a decided prompt that the daemon refuses (another send of it still in flight) +/// leaves the owner's decision standing. +#[test] +fn a_refused_retry_never_undoes_a_decision() { + let room = FakeRoom::start(&[]); + let (gate, gated) = channel(); + let fixture = Fixture { prompt_gate: Some(Arc::new(Mutex::new(gated))), ..Fixture::default() }; + let (mut fake, _mux, surface, mut owner) = bridged_with("relay-retry-final", &room, fixture); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + gate.send(()).unwrap(); + let prompts = held_prompts(&room, &fake, conn, &ann, 0..1); + fake.notify("_acpmux/event", approved(40, &prompts[0])); + assert_eq!(state_of(&room, conn, "c0")["state"], "approved"); + // The first retry waits at acpmux; the second is refused while it is out. + for _ in 0..2 { + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "c0", + "text": "hello"})); + } + std::thread::sleep(Duration::from_millis(500)); + drop(gate); + room.send(json!({"v": 1, "t": "prompt_status", "guest_id": ann, "client_ids": ["c0"]})); + let state = state_of(&room, conn, "c0"); + assert_eq!(state["state"], "approved", "{state}"); +} + +/// A guest's reconnect while its relay's attachment is still answering keeps the prompts that +/// relay held: the new relay sends them. +#[test] +fn a_guest_reconnect_keeps_the_prompts_its_old_relay_held() { + let room = FakeRoom::start(&[]); + let (gate, gated) = channel(); + let fixture = Fixture { attach_gate: Some(Arc::new(Mutex::new(gated))), ..Fixture::default() }; + let (fake, _mux, surface, mut owner) = bridged_with("relay-rejoin", &room, fixture); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "early", + "text": "hello"})); + std::thread::sleep(Duration::from_millis(300)); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + std::thread::sleep(Duration::from_millis(300)); + gate.send(()).unwrap(); + gate.send(()).unwrap(); + let deadline = Instant::now() + WAIT; + loop { + let frame = fake.next_frame(deadline.saturating_duration_since(Instant::now())); + let frame = frame.expect("the early prompt reaches acpmux"); + if frame["method"] == "session/prompt" { + break; + } + } + let state = state_of(&room, conn, "early"); + assert_eq!(state["state"], "held", "{state}"); +} + +/// A host reconnect's snapshot names the guest's prompts from before it, though the new attach +/// page has none of their records. +#[test] +fn a_host_reconnect_snapshot_names_the_guests_prompts() { + let room = FakeRoom::start(&[]); + let (fake, _mux, surface, mut owner) = bridged("relay-host-states", &room); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + let prompts = held_prompts(&room, &fake, conn, &ann, 0..1); + room.shared.hello.lock().unwrap().push(json!({"guest_id": ann, "name": "Ann", "sockets": 1})); + room.drop_host(); + let (conn, _) = room.host(); + assert_eq!(room.next(conn, "guests")["t"], "guests"); + let snapshot = room.next(conn, "snapshot"); + assert_eq!( + snapshot["prompts"], + json!([{"client_id": "c0", "prompt_id": prompts[0], "state": "held"}]), + "{snapshot}" + ); +} + +/// History asks past what the relay keeps in flight, while acpmux is slow to answer, are each +/// answered at once with an empty page that has more. +#[test] +fn history_asks_past_the_in_flight_cap_get_a_busy_page() { + let room = FakeRoom::start(&[]); + let (gate, gated) = channel::<()>(); + let fixture = Fixture { events_gate: Some(Arc::new(Mutex::new(gated))), ..Fixture::default() }; + let (fake, _mux, surface, mut owner) = bridged_with("relay-older-flood", &room, fixture); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + for _ in 0..20 { + room.send(json!({"v": 1, "t": "older", "guest_id": ann, "before_seq": 3})); + } + let started = Instant::now(); + for _ in 0..16 { + let page = room.next(conn, "page"); + assert_eq!((page["records"].clone(), page["has_more"].clone()), (json!([]), json!(true))); + } + assert!(started.elapsed() < Duration::from_secs(4), "answered without waiting on acpmux"); + drop(gate); +} + +/// A prompt the daemon has not answered when the host socket drops is sent again by the guest's +/// relay on the new host socket, and its state reaches the guest there. +#[test] +fn a_prompt_in_flight_across_a_host_reconnect_is_sent_again() { + let room = FakeRoom::start(&[]); + let (gate, gated) = channel(); + let fixture = Fixture { prompt_gate: Some(Arc::new(Mutex::new(gated))), ..Fixture::default() }; + let (fake, _mux, surface, mut owner) = bridged_with("relay-host-inflight", &room, fixture); + let (_share, [ann, _bob], conn) = two_guests(&room, &mut owner, surface); + room.send(json!({"v": 1, "t": "guest_joined", "guest_id": ann})); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + while fake.next_frame(Duration::from_millis(300)).is_some() {} + room.send(json!({"v": 1, "t": "guest_prompt", "guest_id": ann, "client_id": "c0", + "text": "hello"})); + while fake.frame()["method"] != "session/prompt" {} + // acpmux has not accepted it when the host socket drops. + room.shared.hello.lock().unwrap().push(json!({"guest_id": ann, "name": "Ann", "sockets": 1})); + room.drop_host(); + let (conn, _) = room.host(); + assert_eq!(room.next(conn, "guests")["t"], "guests"); + assert_eq!(room.next(conn, "snapshot")["to"], json!(ann)); + // The first send's acceptance goes to the gone relay; the second is the new relay's. + gate.send(()).unwrap(); + while fake.frame()["method"] != "session/prompt" {} + gate.send(()).unwrap(); + let state = state_of(&room, conn, "c0"); + assert_eq!(state["state"], "held", "{state}"); +} diff --git a/cmux-tui/crates/cmux-tui-core/src/server/agent_session_shares.rs b/cmux-tui/crates/cmux-tui-core/src/server/agent_session_shares.rs index 037bcdb43db5..dcf6eec67c9a 100644 --- a/cmux-tui/crates/cmux-tui-core/src/server/agent_session_shares.rs +++ b/cmux-tui/crates/cmux-tui-core/src/server/agent_session_shares.rs @@ -913,6 +913,14 @@ pub(super) fn guest_author(mux: &Arc, link: Value) -> Result String { + let hash = digest(&[share_id, guest_id, client_id]); + let prompt_id = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(hash); + format!("g-{}", prompt_id.get(..22).unwrap_or(&prompt_id)) +} + /// A bound guest's frame: its share and guest are still live, it names no other author or chat, /// and it goes to the surface that shows the share's chat (whatever surface it named). A prompt's /// id is `g-` and 22 characters of base64url SHA-256 of the share, the guest and the frame's @@ -946,10 +954,7 @@ pub(super) fn check_guest( let client_id = object.remove("client_id"); let client_id = client_id.as_ref().and_then(Value::as_str).filter(|c| valid_id(c)); let Some(client_id) = client_id else { return Err(Refusal::BadRequest) }; - let hash = digest(&[share_id, guest_id, client_id]); - let prompt_id = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(hash); - let prompt_id = prompt_id.get(..22).unwrap_or(&prompt_id); - object.insert("prompt_id".into(), json!(format!("g-{prompt_id}"))); + object.insert("prompt_id".into(), json!(guest_prompt_id(share_id, guest_id, client_id))); } if cmd == "agent-session-detach" { // This connection's attachment to the shared chat, on whatever surface it resolved to