Skip to main content

bamboo_server/handlers/agent/ws_v2/
mod.rs

1//! v2-P1 unified WebSocket multiplex: `GET /v2/stream`.
2//!
3//! One WebSocket replaces the two v1 SSE streams (`GET /events/{id}` +
4//! `GET /stream`) plus a minimal `stop` control uplink, so a mobile client uses
5//! ONE connection instead of two SSE + scattered POSTs. The v1 SSE/REST
6//! endpoints stay UNCHANGED (dual-track per `docs/api-v2-transport.md` §8.1).
7//!
8//! ## Channels (server→client)
9//! - `feed` — the account `ChangeEvent` stream (reuses `plan_replay` + journal).
10//! - `agent.{session_id}` — a per-session `AgentEvent` stream (reuses the v1
11//!   critical-event replay + the v1 `Coalescer` for token batching).
12//!
13//! ## Control (client→server, P1)
14//! - `stop` only — cancels a session (reuses the v1 stop discipline).
15//!
16//! ## Design
17//! Each subscribed channel runs its OWN forwarder task with its OWN broadcast
18//! receiver (per-channel lag independence; §10-Q3) AND its OWN bounded outbound
19//! `mpsc` queue. The driver holds a `StreamMap<channel, ReceiverStream>` and
20//! drains every per-channel queue with a fair (randomized-start) merge, so a burst
21//! on one channel can no longer head-of-line another at the socket (RFC §10-Q3 —
22//! see `forwarders.rs` for the honest fairness guarantee). A per-channel
23//! `JoinHandle` lets `unsubscribe`/teardown abort exactly that forwarder and drop
24//! its queue, leaving no orphaned broadcast reader.
25//!
26//! ## Encoding (v2-P3, #181)
27//! The wire encoding is negotiated ONCE at the upgrade from the offered
28//! `Sec-WebSocket-Protocol`: `bamboo.v2.msgpack` selects binary MessagePack,
29//! anything else (or nothing) stays JSON text (the default — desktop /
30//! debuggability). The SAME envelope schema is carried either way; only the
31//! serialization + WS frame type (Text vs Binary) differs. See `envelope::Encoding`.
32//!
33//! ## DEFERRED (later slices; see PR / §5.3)
34//! - `execute` / `approve` over control (handler refactors) — clients keep REST.
35//!
36//! ## Auth (v2-P2, #181 / #189)
37//! `/v2/stream` is on the public route whitelist, so the upgrade OPENS without
38//! a middleware credential — this is the only way a browser device-token client
39//! (which cannot set `Authorization`/`X-Device-Id` headers on a WS upgrade) can
40//! present its token, via the `hello` frame. The handler is therefore the
41//! AUTHORITATIVE gate:
42//!
43//! - `pre_authorized` is computed from the SAME allow-decision the middleware
44//!   uses for every other route (`request_is_authorized`): a local bypass, a
45//!   verified password cookie (cookies ARE sent on the upgrade), or a header
46//!   device token. These connections stay frictionless — exactly as before.
47//! - While NOT authorized, the ONLY frame that does anything is `hello` carrying
48//!   a VALID `device_id` + `token`. A `subscribe`/`unsubscribe`/`stop` received
49//!   before authorization is IGNORED — no forwarder, no cancel, no channel data.
50//! - A token-less `hello` NEVER authorizes a connection that wasn't already
51//!   pre-authorized.
52//! - An unauthenticated socket that never sends a valid `hello` within
53//!   [`AUTH_DEADLINE`] is CLOSED, so it cannot linger.
54//! - The token is NEVER logged.
55
56mod envelope;
57mod forwarders;
58
59use std::collections::HashMap;
60use std::time::Duration;
61
62use actix_web::{web, HttpRequest, Responder};
63use actix_ws::Message;
64use futures::StreamExt;
65use tokio::sync::mpsc;
66use tokio_stream::wrappers::ReceiverStream;
67use tokio_stream::StreamMap;
68
69use serde::Deserialize;
70
71use self::envelope::{
72    decode_client_frame, pong_frame, sys_keepalive_envelope, Channel, ClientFrame, Encoding,
73    OutFrame, SUBPROTOCOL_JSON, SUBPROTOCOL_MSGPACK,
74};
75use self::forwarders::{
76    spawn_agent_forwarder, spawn_agent_terminal_forwarder, spawn_feed_forwarder, OutboundTx,
77};
78use crate::app_state::AppState;
79use crate::handlers::agent::events::MAX_BATCH_MS;
80use crate::handlers::agent::stop::cancel_session;
81
82/// Bound on EACH per-channel outbound mpsc (RFC §10-Q3). Caps how far ahead one
83/// channel's forwarder may run before the WS writer applies backpressure to that
84/// channel ALONE — a slow socket no longer lets a bursting channel starve the
85/// queue of another, because each channel owns its own bounded buffer and the
86/// driver drains them with a fair merge.
87const OUTBOUND_BUFFER: usize = 64;
88
89/// Reserved connection-level queue for heartbeat/control frames. A capacity of
90/// one intentionally coalesces probes while the socket writer is busy.
91const SYS_CHANNEL: &str = "sys";
92const SYS_OUTBOUND_BUFFER: usize = 1;
93
94/// Best-effort enqueue for connection-level control traffic. Heartbeats are
95/// probes, so retaining an old one is never worth blocking the socket read loop.
96fn try_enqueue_sys(sys_tx: &mpsc::Sender<OutFrame>, frame: Option<OutFrame>) {
97    if let Some(frame) = frame {
98        let _ = sys_tx.try_send(frame);
99    }
100}
101
102/// WS ping interval. Each tick sends BOTH the protocol-level ping (server-side
103/// write probe) and the app-level `sys` keepalive data frame (client-side
104/// liveness signal, #533).
105///
106/// 2s (15s → 5s in #543, then → 2s): the client watchdog cannot detect a dead
107/// socket faster than a few missed keepalives, and a user actively watching a
108/// running session should not stare at a dead screen for long — at 2s the
109/// watchdog resolves in ~6s. The frames are tiny (a ping + a ~50-byte text
110/// frame); the cost is negligible even for remote clients. The lotus watchdog
111/// ADAPTS its threshold to the observed cadence (3×, clamped), so old-server ×
112/// new-client pairings stay safe in both directions.
113const PING_INTERVAL: Duration = Duration::from_secs(2);
114
115/// Env var the live integration test sets to SHORTEN the ping interval so the
116/// `sys` keepalive cadence is asserted in milliseconds, not 15s. Read once per
117/// connection in [`drive`]. Production NEVER sets it.
118const PING_INTERVAL_OVERRIDE_ENV: &str = "BAMBOO_WS_PING_INTERVAL_MS";
119
120/// The effective ping interval: the [`PING_INTERVAL`] default unless
121/// [`PING_INTERVAL_OVERRIDE_ENV`] is set to a valid millisecond count (test-only).
122fn ping_interval() -> Duration {
123    std::env::var(PING_INTERVAL_OVERRIDE_ENV)
124        .ok()
125        .and_then(|v| v.parse::<u64>().ok())
126        .map(Duration::from_millis)
127        .unwrap_or(PING_INTERVAL)
128}
129
130/// How long an UNAUTHORIZED connection may stay open before it must present a
131/// valid `hello` device token. A connection that opens the (now public) upgrade
132/// without a header/cookie credential and never authenticates is CLOSED when
133/// this elapses, so an unauthenticated socket can never linger (#189). Once the
134/// connection is authorized this deadline is disarmed.
135const AUTH_DEADLINE: Duration = Duration::from_secs(10);
136
137/// Env var the live integration test sets to SHORTEN the auth deadline so the
138/// "closed when no hello arrives" path is asserted in milliseconds, not 10s. It
139/// is read once per connection in [`drive`]. Production NEVER sets it, so the
140/// 10s [`AUTH_DEADLINE`] default is unchanged.
141const AUTH_DEADLINE_OVERRIDE_ENV: &str = "BAMBOO_WS_AUTH_DEADLINE_MS";
142
143/// The effective unauthorized deadline: the [`AUTH_DEADLINE`] default unless
144/// [`AUTH_DEADLINE_OVERRIDE_ENV`] is set to a valid millisecond count (test-only).
145fn auth_deadline() -> Duration {
146    std::env::var(AUTH_DEADLINE_OVERRIDE_ENV)
147        .ok()
148        .and_then(|v| v.parse::<u64>().ok())
149        .map(Duration::from_millis)
150        .unwrap_or(AUTH_DEADLINE)
151}
152
153/// The verdict for one client frame while a connection is NOT yet authorized.
154///
155/// Pure decision over the parsed frame (no `AppState`), so the auth-gating
156/// contract is unit-testable without a live WS driver: while unauthorized the
157/// ONLY frame that can change anything is a `hello`, and only a `hello` carrying
158/// a credential is even a candidate to authorize.
159#[derive(Debug, PartialEq, Eq)]
160enum UnauthorizedAction {
161    /// A `hello` carrying `device_id` + `token` — verify it; valid → authorize,
162    /// invalid → close.
163    VerifyHello,
164    /// A token-less `hello` while unauthorized — a harmless no-op that does NOT
165    /// grant access (the deadline still governs).
166    TokenlessHelloNoop,
167    /// Any other frame (`subscribe`/`unsubscribe`/`stop`/unknown) while
168    /// unauthorized — IGNORE it (serve no channel data, cancel nothing).
169    Ignore,
170}
171
172impl UnauthorizedAction {
173    /// Classify a parsed frame for an unauthorized connection.
174    fn classify(frame: &ClientFrame) -> Self {
175        match frame {
176            ClientFrame::Hello {
177                device_id: Some(_),
178                token: Some(_),
179            } => UnauthorizedAction::VerifyHello,
180            ClientFrame::Hello { .. } => UnauthorizedAction::TokenlessHelloNoop,
181            _ => UnauthorizedAction::Ignore,
182        }
183    }
184}
185
186/// What the driver should do with a frame after the auth gate ran.
187#[derive(Debug, PartialEq, Eq)]
188enum GateOutcome {
189    /// The frame was fully handled by the gate (the connection is/was
190    /// unauthorized, or it just authorized): keep the socket open, do NOT fall
191    /// through to channel dispatch.
192    Handled,
193    /// An invalid credential was presented: CLOSE the connection.
194    Close,
195    /// The connection is authorized and this frame should proceed to the normal
196    /// channel dispatch (subscribe/unsubscribe/stop/hello-rebind).
197    Dispatch,
198}
199
200/// The AppState-aware auth gate, factored out of `handle_client_frame` so it is
201/// unit-testable WITHOUT a live WS `Session`. It owns the entire `!authorized`
202/// decision plus the credential verification, and flips `authorized` to `true`
203/// only on a VERIFIED `hello` device token.
204///
205/// Invariants (the security review hammers these):
206/// - While `!*authorized`, a `subscribe`/`unsubscribe`/`stop` returns
207///   [`GateOutcome::Handled`] (ignored — never `Dispatch`), so no channel is
208///   served and no session cancelled before authorization.
209/// - A token-less `hello` NEVER sets `*authorized` when it was `false`.
210/// - A credentialed `hello` with a VALID token sets `*authorized = true`.
211/// - A credentialed `hello` with an INVALID token returns [`GateOutcome::Close`].
212/// - The token is NEVER logged.
213async fn apply_auth_gate(
214    state: &web::Data<AppState>,
215    frame: &ClientFrame,
216    authorized: &mut bool,
217) -> GateOutcome {
218    if *authorized {
219        return GateOutcome::Dispatch;
220    }
221
222    match UnauthorizedAction::classify(frame) {
223        UnauthorizedAction::VerifyHello => {
224            let ClientFrame::Hello {
225                device_id: Some(device_id),
226                token: Some(token),
227            } = frame
228            else {
229                unreachable!("VerifyHello implies a credentialed Hello");
230            };
231            let config = state.config.read().await.clone();
232            if crate::handlers::settings::verify_device_token(&config, device_id, token) {
233                *authorized = true;
234                // Bind device id for logging. NEVER log the token.
235                tracing::debug!("ws_v2: hello verified for device {device_id}; authorized");
236                GateOutcome::Handled
237            } else {
238                tracing::warn!(
239                    "ws_v2: hello rejected — invalid device credential for {device_id}; closing"
240                );
241                GateOutcome::Close
242            }
243        }
244        UnauthorizedAction::TokenlessHelloNoop => {
245            // A token-less hello does NOT authorize a non-pre-authorized
246            // connection. Keep the socket open; the deadline still governs.
247            tracing::debug!("ws_v2: token-less hello while unauthorized — not granting access");
248            GateOutcome::Handled
249        }
250        UnauthorizedAction::Ignore => {
251            // subscribe/unsubscribe/stop/unknown before auth: serve nothing,
252            // cancel nothing. Tolerate hello-after-subscribe ordering.
253            tracing::debug!("ws_v2: ignoring frame on unauthorized connection");
254            GateOutcome::Handled
255        }
256    }
257}
258
259/// Query parameters for the `GET /v2/stream` upgrade.
260#[derive(Debug, Default, Deserialize)]
261pub struct StreamQuery {
262    /// Token-coalescing window in milliseconds for `agent.{sid}` channels.
263    /// `0` (default) = no coalescing (desktop). Mobile passes e.g. `50`.
264    #[serde(default)]
265    pub batch_ms: u64,
266}
267
268/// Negotiate the wire [`Encoding`] from the upgrade's `Sec-WebSocket-Protocol`
269/// header (v2-P3, §5.3 / §7.2). Returns the chosen encoding AND the subprotocol
270/// token the server must ECHO on the upgrade response (per RFC 6455 the server
271/// echoes the single selected subprotocol).
272///
273/// - Offers including `bamboo.v2.msgpack` → `(Msgpack, Some("bamboo.v2.msgpack"))`.
274/// - Offers including only `bamboo.v2` → `(Json, Some("bamboo.v2"))`.
275/// - No (recognized) subprotocol offered → `(Json, None)` — today's behavior is
276///   preserved byte-for-byte for clients that offer nothing.
277///
278/// `bamboo.v2.msgpack` wins if BOTH are offered (the client opted into binary).
279fn negotiate_encoding(req: &HttpRequest) -> (Encoding, Option<&'static str>) {
280    let offered = req
281        .headers()
282        .get(actix_web::http::header::SEC_WEBSOCKET_PROTOCOL)
283        .and_then(|v| v.to_str().ok())
284        .unwrap_or("");
285    // The header is a comma-separated list of subprotocol tokens.
286    let mut has_msgpack = false;
287    let mut has_json = false;
288    for tok in offered.split(',') {
289        match tok.trim() {
290            SUBPROTOCOL_MSGPACK => has_msgpack = true,
291            SUBPROTOCOL_JSON => has_json = true,
292            _ => {}
293        }
294    }
295    if has_msgpack {
296        (Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
297    } else if has_json {
298        (Encoding::Json, Some(SUBPROTOCOL_JSON))
299    } else {
300        (Encoding::Json, None)
301    }
302}
303
304/// `GET /v2/stream` — upgrade to the unified WS multiplex.
305///
306/// The upgrade itself is PUBLIC (whitelisted in `is_public_access_route`), so a
307/// browser device-token client can open the socket and authenticate via `hello`
308/// (it cannot set headers on a WS upgrade). Auth is enforced in `drive`: a
309/// connection that the middleware WOULD have allowed is `pre_authorized` here
310/// (local bypass / verified password cookie / header device token), via the same
311/// `request_is_authorized` allow-decision the middleware uses; everything else
312/// must present a verified `hello` before any channel is served, on a deadline.
313///
314/// The wire [`Encoding`] is negotiated from `Sec-WebSocket-Protocol` (v2-P3):
315/// `bamboo.v2.msgpack` → binary MessagePack; otherwise JSON text (default). The
316/// selected subprotocol is ECHOED on the upgrade response per RFC 6455.
317pub async fn handler(
318    state: web::Data<AppState>,
319    query: web::Query<StreamQuery>,
320    req: HttpRequest,
321    body: web::Payload,
322) -> actix_web::Result<impl Responder> {
323    // Clamp the untrusted `batch_ms` the same way the v1 SSE handler does.
324    let batch_ms = query.batch_ms.min(MAX_BATCH_MS);
325
326    // Negotiate the wire encoding from the offered subprotocols (v2-P3).
327    let (encoding, selected_subprotocol) = negotiate_encoding(&req);
328
329    // Pre-authorize exactly the connections the middleware would have allowed on
330    // a gated route: local bypass, a verified password cookie (sent on the
331    // upgrade), or a header device token. This preserves every existing client
332    // with ZERO change; a remote browser device-token client is NOT pre-auth and
333    // must send a valid `hello`.
334    let pre_authorized = {
335        let config = state.config.read().await.clone();
336        crate::handlers::settings::request_is_authorized(&req, &config)
337    };
338
339    let (mut response, session, msg_stream) = actix_ws::handle(&req, body)?;
340
341    // Echo the selected subprotocol on the upgrade RESPONSE (RFC 6455). A client
342    // that offered no recognized subprotocol gets none, keeping the legacy
343    // handshake byte-for-byte unchanged.
344    if let Some(proto) = selected_subprotocol {
345        response.headers_mut().insert(
346            actix_web::http::header::SEC_WEBSOCKET_PROTOCOL,
347            actix_web::http::header::HeaderValue::from_static(proto),
348        );
349    }
350
351    actix_web::rt::spawn(drive(
352        state,
353        session,
354        msg_stream,
355        batch_ms,
356        pre_authorized,
357        encoding,
358    ));
359
360    Ok(response)
361}
362
363/// The per-connection driver: owns the WS `session` (write) + `msg_stream`
364/// (read), the shared outbound mpsc, the per-channel forwarder handles, and the
365/// keepalive ping timer.
366async fn drive(
367    state: web::Data<AppState>,
368    mut session: actix_ws::Session,
369    mut msg_stream: actix_ws::MessageStream,
370    batch_ms: u64,
371    pre_authorized: bool,
372    encoding: Encoding,
373) {
374    // Per-channel outbound (RFC §10-Q3): every subscribed channel owns its OWN
375    // bounded queue. `forwarders` keeps the task handle (for unsubscribe/teardown
376    // abort) and `queues` keeps the matching receiver in a `StreamMap` the driver
377    // drains with a fair (randomized-start) merge, so a burst on one channel cannot
378    // head-of-line another at the socket. The two maps are keyed by the SAME
379    // channel id and mutated together (see [`subscribe`] / unsubscribe / teardown).
380    let mut forwarders: HashMap<String, tokio::task::JoinHandle<()>> = HashMap::new();
381    let mut queues: StreamMap<String, ReceiverStream<OutFrame>> = StreamMap::new();
382    let (sys_tx, sys_rx) = mpsc::channel::<OutFrame>(SYS_OUTBOUND_BUFFER);
383    queues.insert(SYS_CHANNEL.to_string(), ReceiverStream::new(sys_rx));
384    let mut ping = tokio::time::interval(ping_interval());
385    ping.tick().await; // skip the immediate tick
386
387    // Authorization state (#189). Seeded from the upgrade-time decision so local
388    // / cookie / header clients are authorized immediately; everything else must
389    // present a valid `hello` before any channel is served.
390    let mut authorized = pre_authorized;
391    // The unauthorized-deadline timer. Pinned + biased to fire while `!authorized`
392    // and disarmed once the connection authorizes, so an unauthenticated socket
393    // is closed but an authorized one runs indefinitely.
394    let auth_deadline = tokio::time::sleep(auth_deadline());
395    tokio::pin!(auth_deadline);
396
397    loop {
398        tokio::select! {
399            // While unauthorized, close the socket once the deadline elapses. The
400            // `!authorized` guard disarms this arm the moment the connection
401            // authorizes (a never-completing branch is simply never selected).
402            _ = &mut auth_deadline, if !authorized => {
403                tracing::debug!("ws_v2: closing unauthorized connection after auth deadline");
404                break;
405            }
406            // Drain the per-channel queues to the WS with a fair merge. The
407            // `StreamMap` randomizes its poll start each call, so no single channel's
408            // queue is favored — a bursting channel cannot starve another at the
409            // socket. `(channel, frame)` is yielded; only the frame goes on the
410            // wire. The frame is ALREADY encoded per the connection's `Encoding`
411            // (the forwarder did the encode), so the driver only picks the WS frame
412            // type: `session.text` for a JSON `Text`, `session.binary` for a
413            // MessagePack `Binary`. If the write fails the peer is gone — break and
414            // tear down.
415            Some((_ch, frame)) = queues.next() => {
416                let write = match frame {
417                    OutFrame::Text(s) => session.text(s).await,
418                    OutFrame::Binary(b) => session.binary(b).await,
419                };
420                if write.is_err() {
421                    break;
422                }
423            }
424            // Keepalive. Liveness is write-driven (a dead peer surfaces as a
425            // failed `ping`/`text` write); we do NOT track Pong arrivals or run a
426            // pong timeout — same one-directional keepalive contract as the v1 SSE
427            // stream.
428            //
429            // TWO frames go out per tick (#533):
430            // - a protocol-level ping — the server-side write probe (unchanged);
431            // - an app-level `{ch:"sys", control:{type:"keepalive"}}` DATA frame —
432            //   browsers never expose protocol pings to JS, so without a data
433            //   frame a client on a half-open socket (sleep/wake, NAT idle
434            //   eviction) has NO observable liveness signal and sits "open"
435            //   forever after the server tears down. The sys frame is what the
436            //   lotus watchdog keys on to force a reconnect. Old clients ignore
437            //   the unknown `sys` channel, so this is backward-compatible.
438            _ = ping.tick() => {
439                if session.ping(b"").await.is_err() {
440                    break;
441                }
442                // Only an AUTHORIZED connection gets the data frame: an
443                // unauthenticated socket is served nothing (same posture as
444                // channel data) and is closed by the auth deadline anyway.
445                if authorized {
446                    if let Some(frame) = sys_keepalive_envelope().encode(encoding) {
447                        let write = match frame {
448                            OutFrame::Text(s) => session.text(s).await,
449                            OutFrame::Binary(b) => session.binary(b).await,
450                        };
451                        if write.is_err() {
452                            break;
453                        }
454                    }
455                }
456            }
457            // Client frames. In JSON mode the client sends TEXT frames; in msgpack
458            // mode it sends BINARY frames. A frame carrying the inbound bytes for the
459            // ACTIVE encoding is decoded + dispatched; a frame of the other kind is
460            // ignored (a Binary frame in JSON mode, a Text frame in msgpack mode),
461            // exactly as a malformed frame is — it never tears down the connection.
462            msg = msg_stream.next() => {
463                match msg {
464                    // The inbound frame type that matches the active encoding.
465                    Some(Ok(Message::Text(text))) if encoding == Encoding::Json => {
466                        let keep_open = handle_client_bytes(ClientDispatchContext {
467                            state: &state,
468                            forwarders: &mut forwarders,
469                            queues: &mut queues,
470                            sys_tx: &sys_tx,
471                            batch_ms,
472                            encoding,
473                            authorized: &mut authorized,
474                        }, text.as_bytes())
475                        .await;
476                        if !keep_open {
477                            break;
478                        }
479                    }
480                    Some(Ok(Message::Binary(bytes))) if encoding == Encoding::Msgpack => {
481                        let keep_open = handle_client_bytes(ClientDispatchContext {
482                            state: &state,
483                            forwarders: &mut forwarders,
484                            queues: &mut queues,
485                            sys_tx: &sys_tx,
486                            batch_ms,
487                            encoding,
488                            authorized: &mut authorized,
489                        }, &bytes)
490                        .await;
491                        if !keep_open {
492                            break;
493                        }
494                    }
495                    Some(Ok(Message::Ping(bytes))) => {
496                        if session.pong(&bytes).await.is_err() {
497                            break;
498                        }
499                    }
500                    // A frame of the WRONG kind for the active encoding (a Binary
501                    // frame in JSON mode / a Text frame in msgpack mode), plus Pong /
502                    // Continuation / Nop: ignore, never disconnect.
503                    Some(Ok(Message::Text(_)))
504                    | Some(Ok(Message::Binary(_)))
505                    | Some(Ok(Message::Pong(_)))
506                    | Some(Ok(Message::Continuation(_)))
507                    | Some(Ok(Message::Nop)) => {}
508                    Some(Ok(Message::Close(_))) | None => break,
509                    Some(Err(e)) => {
510                        tracing::debug!("ws_v2: message stream error: {e}");
511                        break;
512                    }
513                }
514            }
515            else => break,
516        }
517    }
518
519    // Clean teardown: abort every forwarder so no orphaned broadcast reader
520    // survives the connection, and drop every per-channel queue (clearing the
521    // `StreamMap` drops each receiver).
522    for (_ch, handle) in forwarders.drain() {
523        handle.abort();
524    }
525    queues.clear();
526    let _ = session.close(None).await;
527}
528
529/// Mutable connection state shared by the decode and dispatch steps.
530struct ClientDispatchContext<'a> {
531    state: &'a web::Data<AppState>,
532    forwarders: &'a mut HashMap<String, tokio::task::JoinHandle<()>>,
533    queues: &'a mut StreamMap<String, ReceiverStream<OutFrame>>,
534    sys_tx: &'a mpsc::Sender<OutFrame>,
535    batch_ms: u64,
536    encoding: Encoding,
537    authorized: &'a mut bool,
538}
539
540/// Decode one inbound frame's bytes per the connection's [`Encoding`] (serde_json
541/// for `Json`, rmp-serde for `Msgpack`) and dispatch the resulting [`ClientFrame`].
542/// A malformed body logs and is ignored — it NEVER tears down the connection, in
543/// EITHER encoding (the msgpack decode error is treated exactly like a malformed
544/// JSON text frame).
545///
546/// This is the thin decode step in front of [`handle_client_frame`]; the dispatch
547/// and auth logic is encoding-agnostic (the SAME `ClientFrame` flows through both
548/// paths), so the JSON behavior is byte-for-byte unchanged.
549///
550/// Returns `false` to signal the driver to CLOSE the connection (a `hello` that
551/// presents an INVALID device credential); `true` to keep it open.
552async fn handle_client_bytes(context: ClientDispatchContext<'_>, bytes: &[u8]) -> bool {
553    let frame: ClientFrame = match decode_client_frame(context.encoding, bytes) {
554        Ok(f) => f,
555        Err(e) => {
556            tracing::debug!("ws_v2: ignoring malformed client frame: {e}");
557            return true;
558        }
559    };
560    handle_client_frame(context, frame).await
561}
562
563/// Dispatch one decoded client frame. A malformed/unknown frame logs and is
564/// ignored — it never tears down the connection.
565///
566/// `authorized` is the per-connection auth state (#189). While it is `false` the
567/// ONLY frame that does anything is a `hello` carrying a VALID device token (it
568/// flips `authorized` to `true`); EVERY other frame — including a
569/// `subscribe`/`unsubscribe`/`stop` — is IGNORED, so a remote unauthenticated
570/// connection gets NO channel data and cancels nothing. The driver's deadline
571/// closes a socket that never authorizes.
572///
573/// Returns `false` to signal the driver to CLOSE the connection (a `hello` that
574/// presents an INVALID device credential); `true` to keep it open.
575async fn handle_client_frame(context: ClientDispatchContext<'_>, frame: ClientFrame) -> bool {
576    let ClientDispatchContext {
577        state,
578        forwarders,
579        queues,
580        sys_tx,
581        batch_ms,
582        encoding,
583        authorized,
584    } = context;
585
586    // Auth gate (#189). Until the connection is authorized, no frame may serve a
587    // channel or cancel a session — the only frame that can change anything is a
588    // `hello` that carries a verifiable device credential.
589    match apply_auth_gate(state, &frame, authorized).await {
590        GateOutcome::Handled => return true,
591        GateOutcome::Close => return false,
592        GateOutcome::Dispatch => {}
593    }
594
595    // Like all server data, heartbeat acknowledgements are emitted only after
596    // authorization. They carry no channel envelope but share the connection's
597    // reserved sys queue and single socket writer.
598    if frame == ClientFrame::Ping {
599        // Drop-on-full is deliberate. The next client heartbeat retries, and
600        // the read loop must never wait behind a slow socket writer.
601        try_enqueue_sys(sys_tx, pong_frame(encoding));
602        return true;
603    }
604
605    // Authorized path: full dispatch.
606    match frame {
607        ClientFrame::Hello { device_id, token } => {
608            // Already authorized. A credentialed hello re-verifies as identity
609            // binding; an invalid one still closes. A token-less hello on an
610            // already-authorized connection is a harmless no-op.
611            match (device_id, token) {
612                (Some(device_id), Some(token)) => {
613                    let config = state.config.read().await.clone();
614                    if crate::handlers::settings::verify_device_token(&config, &device_id, &token) {
615                        // NEVER log the token.
616                        tracing::debug!("ws_v2: hello verified for device {device_id}");
617                    } else {
618                        tracing::warn!(
619                            "ws_v2: hello rejected — invalid device credential for {device_id}; closing"
620                        );
621                        return false;
622                    }
623                }
624                _ => {
625                    tracing::debug!("ws_v2: token-less hello on authorized connection (no-op)");
626                }
627            }
628        }
629        ClientFrame::Subscribe { ch, since } => {
630            subscribe(state, forwarders, queues, batch_ms, encoding, &ch, since).await;
631        }
632        ClientFrame::Unsubscribe { ch } => {
633            if let Some(handle) = forwarders.remove(&ch) {
634                handle.abort();
635                // Drop this channel's queue too: removing it from the StreamMap
636                // drops the receiver so the (now-aborted) forwarder's sender is
637                // dead and no stale frame can leak onto the socket.
638                queues.remove(&ch);
639                tracing::debug!("ws_v2: unsubscribed {ch}");
640            }
641        }
642        ClientFrame::Stop { session_id } => {
643            let cancelled = cancel_session(state, &session_id).await;
644            tracing::debug!("ws_v2: stop {session_id} -> cancelled={cancelled}");
645        }
646        ClientFrame::Ping => unreachable!("authorized ping handled above"),
647        ClientFrame::Unknown => {
648            tracing::debug!("ws_v2: ignoring unknown client frame type");
649        }
650    }
651    true
652}
653
654/// Subscribe to a channel, replacing any existing forwarder for the same `ch`
655/// (a re-subscribe with a new cursor aborts the old one first, so there is never
656/// a duplicate reader on one channel).
657async fn subscribe(
658    state: &web::Data<AppState>,
659    forwarders: &mut HashMap<String, tokio::task::JoinHandle<()>>,
660    queues: &mut StreamMap<String, ReceiverStream<OutFrame>>,
661    batch_ms: u64,
662    encoding: Encoding,
663    ch: &str,
664    since: Option<u64>,
665) {
666    let Some(channel) = Channel::parse(ch) else {
667        tracing::debug!("ws_v2: ignoring subscribe to unknown channel {ch}");
668        return;
669    };
670
671    // Replace any prior forwarder for this exact channel id, dropping its queue
672    // (a re-subscribe with a new cursor must not leave the old queue behind).
673    if let Some(old) = forwarders.remove(ch) {
674        old.abort();
675    }
676    queues.remove(ch);
677
678    // This channel's OWN bounded outbound queue (RFC §10-Q3). The forwarder owns
679    // the sender; the driver drains the receiver via the fair `StreamMap` merge.
680    // Carries already-encoded [`OutFrame`]s (Text/Binary per the connection's
681    // `Encoding`).
682    let (out_tx, out_rx) = mpsc::channel::<OutFrame>(OUTBOUND_BUFFER);
683    let out_tx: OutboundTx = out_tx;
684
685    let handle = match channel {
686        Channel::Feed => {
687            // Subscribe FIRST so events written during journal replay are buffered
688            // in the ring and delivered in the live phase (no gap) — exactly the
689            // v1 SSE handoff discipline.
690            let receiver = state.account_sink.subscribe();
691            let latest_at_start = state.account_sink.latest_seq();
692            let events_dir = state.account_sink.events_dir().to_path_buf();
693            let since = since.unwrap_or(0);
694            spawn_feed_forwarder(
695                out_tx,
696                encoding,
697                receiver,
698                events_dir,
699                since,
700                latest_at_start,
701            )
702        }
703        Channel::Agent(sid) => {
704            // The session must exist; otherwise ignore (parity with the v1
705            // events handler's 404, but here we just skip the subscribe).
706            if state.session_store.get_index_entry(&sid).await.is_none() {
707                tracing::debug!("ws_v2: ignoring subscribe to unknown session {sid}");
708                return;
709            }
710
711            let sender = state.get_session_event_sender(&sid).await;
712            let receiver = sender.subscribe();
713            // Keep the notification relay running for this session (parity with
714            // the v1 events handler).
715            state.ensure_notification_relay(&sid, sender.clone());
716
717            // Snapshot the runner for critical-event + budget replay (mirrors
718            // `events/handler.rs:80-89`).
719            let runner_snapshot = {
720                let runners = state.agent_runners.read().await;
721                runners.get(&sid).cloned()
722            };
723            let budget_event_to_replay = runner_snapshot
724                .as_ref()
725                .and_then(|runner| runner.last_budget_event.clone());
726            let critical_events_to_replay: Vec<_> = runner_snapshot
727                .as_ref()
728                .map(|runner| runner.last_critical_events.clone())
729                .unwrap_or_default();
730
731            // Match the v1 SSE late-subscribe contract. A session that
732            // completed while this socket was half-open must replay cached
733            // state and a synthesized terminal exactly once, rather than open
734            // a live receiver that will never publish another event.
735            let runner_status = runner_snapshot.as_ref().map(|runner| runner.status.clone());
736            if can_attempt_terminal_replay(runner_status.as_ref(), &receiver) {
737                if let Some(terminal_event) =
738                    crate::handlers::agent::events::terminal_event_if_ready(
739                        state,
740                        &sid,
741                        runner_status.clone(),
742                    )
743                    .await
744                {
745                    // Storage and descendant checks above are asynchronous. A
746                    // new runner can be reserved while they run, before it has
747                    // emitted its first broadcast event. Re-read the runner so
748                    // that Pending/Running wins over the stale terminal
749                    // snapshot even while the receiver is still empty.
750                    // We subscribed before the async storage/child checks. If a
751                    // live event arrived meanwhile, preserve its ordering by
752                    // handing the still-buffered receiver to the live forwarder
753                    // instead of sending a synthetic terminal ahead of it.
754                    if current_runner_allows_terminal_replay(state, &sid, &receiver).await {
755                        return finish_subscribe(
756                            forwarders,
757                            queues,
758                            ch,
759                            out_rx,
760                            spawn_agent_terminal_forwarder(
761                                out_tx,
762                                encoding,
763                                ch.to_string(),
764                                budget_event_to_replay,
765                                critical_events_to_replay,
766                                terminal_event,
767                            ),
768                            since,
769                        );
770                    }
771                }
772            }
773
774            spawn_agent_forwarder(
775                state.clone(),
776                sid.clone(),
777                out_tx,
778                encoding,
779                ch.to_string(),
780                receiver,
781                budget_event_to_replay,
782                critical_events_to_replay,
783                batch_ms,
784            )
785        }
786    };
787    finish_subscribe(forwarders, queues, ch, out_rx, handle, since);
788}
789
790fn can_attempt_terminal_replay(
791    runner_status: Option<&crate::app_state::AgentStatus>,
792    receiver: &tokio::sync::broadcast::Receiver<bamboo_agent_core::AgentEvent>,
793) -> bool {
794    matches!(
795        runner_status,
796        None | Some(crate::app_state::AgentStatus::Completed)
797            | Some(crate::app_state::AgentStatus::Cancelled)
798            | Some(crate::app_state::AgentStatus::Error(_))
799    ) && receiver.is_empty()
800}
801
802/// Re-read the runner after asynchronous terminal checks. This is deliberately
803/// a separate helper so the subscribe race can be covered without relying on
804/// scheduler timing: a runner reserved during storage I/O must invalidate the
805/// stale terminal snapshot before the one-shot forwarder is installed.
806async fn current_runner_allows_terminal_replay(
807    state: &web::Data<AppState>,
808    session_id: &str,
809    receiver: &tokio::sync::broadcast::Receiver<bamboo_agent_core::AgentEvent>,
810) -> bool {
811    let current_runner_status = {
812        let runners = state.agent_runners.read().await;
813        runners.get(session_id).map(|runner| runner.status.clone())
814    };
815    can_attempt_terminal_replay(current_runner_status.as_ref(), receiver)
816}
817
818fn finish_subscribe(
819    forwarders: &mut HashMap<String, tokio::task::JoinHandle<()>>,
820    queues: &mut StreamMap<String, ReceiverStream<OutFrame>>,
821    ch: &str,
822    out_rx: mpsc::Receiver<OutFrame>,
823    handle: tokio::task::JoinHandle<()>,
824    since: Option<u64>,
825) {
826    forwarders.insert(ch.to_string(), handle);
827    // Register this channel's queue receiver into the fair-merge drain. Keyed by
828    // the SAME channel id as `forwarders`, so unsubscribe/teardown drops both.
829    queues.insert(ch.to_string(), ReceiverStream::new(out_rx));
830    tracing::debug!("ws_v2: subscribed {ch} (since={since:?})");
831}
832
833#[cfg(test)]
834mod tests {
835    use super::*;
836    use crate::app_state::AppState;
837    use bamboo_config::{AccessControlConfig, DeviceCredential};
838    use tempfile::tempdir;
839
840    fn hello(device_id: Option<&str>, token: Option<&str>) -> ClientFrame {
841        ClientFrame::Hello {
842            device_id: device_id.map(str::to_string),
843            token: token.map(str::to_string),
844        }
845    }
846
847    #[test]
848    fn sys_queue_is_best_effort_and_drops_when_full() {
849        let (tx, mut rx) = mpsc::channel(SYS_OUTBOUND_BUFFER);
850        let first = OutFrame::Text("first".into());
851        try_enqueue_sys(&tx, Some(first.clone()));
852        // Must return immediately and leave the already-queued frame intact.
853        try_enqueue_sys(&tx, Some(OutFrame::Text("newer".into())));
854        assert_eq!(rx.try_recv().unwrap(), first);
855        assert!(rx.try_recv().is_err());
856    }
857
858    // ── Subprotocol negotiation (v2-P3) ───────────────────────────────────────
859
860    fn negotiate_for(header: Option<&str>) -> (Encoding, Option<&'static str>) {
861        let mut req = actix_web::test::TestRequest::default();
862        if let Some(h) = header {
863            req = req.insert_header((actix_web::http::header::SEC_WEBSOCKET_PROTOCOL, h));
864        }
865        negotiate_encoding(&req.to_http_request())
866    }
867
868    #[test]
869    fn negotiate_encoding_branches() {
870        // No header / empty → JSON default, NO echo (unchanged for old clients).
871        assert_eq!(negotiate_for(None), (Encoding::Json, None));
872        assert_eq!(negotiate_for(Some("")), (Encoding::Json, None));
873        // Explicit JSON subprotocol → JSON, echo it.
874        assert_eq!(
875            negotiate_for(Some("bamboo.v2")),
876            (Encoding::Json, Some(SUBPROTOCOL_JSON))
877        );
878        // Msgpack offered → Msgpack, echo it.
879        assert_eq!(
880            negotiate_for(Some("bamboo.v2.msgpack")),
881            (Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
882        );
883        // Multi-offer: msgpack preferred regardless of order, and whitespace is trimmed.
884        assert_eq!(
885            negotiate_for(Some("bamboo.v2.msgpack, bamboo.v2")),
886            (Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
887        );
888        assert_eq!(
889            negotiate_for(Some("  bamboo.v2 ,  bamboo.v2.msgpack ")),
890            (Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
891        );
892        // Unknown-only offer → JSON, and NO bogus echo (echoing a non-offered
893        // subprotocol would violate RFC 6455).
894        assert_eq!(
895            negotiate_for(Some("some.other.proto")),
896            (Encoding::Json, None)
897        );
898    }
899
900    // ── Pure classification (no AppState) ─────────────────────────────────────
901
902    #[test]
903    fn classify_unauthorized_frame_actions() {
904        // A credentialed hello is the only authorize candidate.
905        assert_eq!(
906            UnauthorizedAction::classify(&hello(Some("d"), Some("t"))),
907            UnauthorizedAction::VerifyHello
908        );
909        // A token-less hello (any missing field) is a no-op, never an authorizer.
910        assert_eq!(
911            UnauthorizedAction::classify(&hello(None, None)),
912            UnauthorizedAction::TokenlessHelloNoop
913        );
914        assert_eq!(
915            UnauthorizedAction::classify(&hello(Some("d"), None)),
916            UnauthorizedAction::TokenlessHelloNoop
917        );
918        assert_eq!(
919            UnauthorizedAction::classify(&hello(None, Some("t"))),
920            UnauthorizedAction::TokenlessHelloNoop
921        );
922        // Every channel-touching / control frame is IGNORED while unauthorized.
923        assert_eq!(
924            UnauthorizedAction::classify(&ClientFrame::Subscribe {
925                ch: "feed".into(),
926                since: None
927            }),
928            UnauthorizedAction::Ignore
929        );
930        assert_eq!(
931            UnauthorizedAction::classify(&ClientFrame::Unsubscribe { ch: "feed".into() }),
932            UnauthorizedAction::Ignore
933        );
934        assert_eq!(
935            UnauthorizedAction::classify(&ClientFrame::Stop {
936                session_id: "s".into()
937            }),
938            UnauthorizedAction::Ignore
939        );
940        assert_eq!(
941            UnauthorizedAction::classify(&ClientFrame::Ping),
942            UnauthorizedAction::Ignore
943        );
944        assert_eq!(
945            UnauthorizedAction::classify(&ClientFrame::Unknown),
946            UnauthorizedAction::Ignore
947        );
948    }
949
950    // ── AppState-aware gate (no live Session) ─────────────────────────────────
951
952    async fn app_state_with_device() -> (web::Data<AppState>, DeviceCredential, String) {
953        let dir = tempdir().unwrap();
954        let state = web::Data::new(AppState::new(dir.path().to_path_buf()).await.unwrap());
955        let (cred, token) = crate::handlers::settings::issue_device_token("test-device");
956        {
957            let mut config = state.config.write().await;
958            config.access_control = Some(AccessControlConfig {
959                password_enabled: false,
960                repair_required: false,
961                password_hash: None,
962                password_salt: None,
963                password_credential_ref: None,
964                password_configured: false,
965                updated_at: None,
966                devices: vec![cred.clone()],
967            });
968        }
969        (state, cred, token)
970    }
971
972    #[actix_web::test]
973    async fn subscribe_while_unauthorized_is_ignored_and_stays_unauthorized() {
974        let (state, _cred, _token) = app_state_with_device().await;
975        let mut authorized = false;
976        let frame = ClientFrame::Subscribe {
977            ch: "feed".into(),
978            since: None,
979        };
980        let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
981        // The driver keeps the socket open but does NOT dispatch (no forwarder).
982        assert_eq!(outcome, GateOutcome::Handled);
983        assert!(!authorized, "a subscribe must never authorize a connection");
984    }
985
986    #[actix_web::test]
987    async fn stop_while_unauthorized_is_ignored() {
988        let (state, _cred, _token) = app_state_with_device().await;
989        let mut authorized = false;
990        let frame = ClientFrame::Stop {
991            session_id: "sess".into(),
992        };
993        let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
994        assert_eq!(outcome, GateOutcome::Handled);
995        assert!(!authorized);
996    }
997
998    #[actix_web::test]
999    async fn valid_hello_authorizes() {
1000        let (state, cred, token) = app_state_with_device().await;
1001        let mut authorized = false;
1002        let frame = hello(Some(&cred.device_id), Some(&token));
1003        let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
1004        assert_eq!(outcome, GateOutcome::Handled);
1005        assert!(authorized, "a valid hello must authorize the connection");
1006    }
1007
1008    #[actix_web::test]
1009    async fn invalid_hello_closes_and_does_not_authorize() {
1010        let (state, cred, _token) = app_state_with_device().await;
1011        let mut authorized = false;
1012        let frame = hello(Some(&cred.device_id), Some("bd1_wrongwrongwrong"));
1013        let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
1014        assert_eq!(outcome, GateOutcome::Close);
1015        assert!(!authorized, "an invalid hello must never authorize");
1016    }
1017
1018    #[actix_web::test]
1019    async fn tokenless_hello_does_not_authorize_unauthorized_connection() {
1020        let (state, _cred, _token) = app_state_with_device().await;
1021        let mut authorized = false;
1022        let frame = hello(None, None);
1023        let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
1024        assert_eq!(outcome, GateOutcome::Handled);
1025        assert!(
1026            !authorized,
1027            "a token-less hello must NEVER authorize a non-pre-authorized connection"
1028        );
1029    }
1030
1031    #[actix_web::test]
1032    async fn pre_authorized_connection_dispatches_subscribe() {
1033        let (state, _cred, _token) = app_state_with_device().await;
1034        // Pre-authorized (local / cookie / header equivalent): the gate passes
1035        // every frame straight through to channel dispatch.
1036        let mut authorized = true;
1037        let frame = ClientFrame::Subscribe {
1038            ch: "feed".into(),
1039            since: None,
1040        };
1041        let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
1042        assert_eq!(outcome, GateOutcome::Dispatch);
1043        assert!(authorized);
1044    }
1045
1046    #[actix_web::test]
1047    async fn ping_is_ignored_before_auth_and_enqueued_after_auth() {
1048        let (state, _cred, _token) = app_state_with_device().await;
1049        let mut forwarders = HashMap::new();
1050        let mut queues = StreamMap::new();
1051        let (sys_tx, mut sys_rx) = mpsc::channel(SYS_OUTBOUND_BUFFER);
1052
1053        let mut authorized = false;
1054        assert!(
1055            handle_client_frame(
1056                ClientDispatchContext {
1057                    state: &state,
1058                    forwarders: &mut forwarders,
1059                    queues: &mut queues,
1060                    sys_tx: &sys_tx,
1061                    batch_ms: 0,
1062                    encoding: Encoding::Json,
1063                    authorized: &mut authorized,
1064                },
1065                ClientFrame::Ping,
1066            )
1067            .await
1068        );
1069        assert!(
1070            sys_rx.try_recv().is_err(),
1071            "unauthorized ping must be silent"
1072        );
1073
1074        authorized = true;
1075        assert!(
1076            handle_client_frame(
1077                ClientDispatchContext {
1078                    state: &state,
1079                    forwarders: &mut forwarders,
1080                    queues: &mut queues,
1081                    sys_tx: &sys_tx,
1082                    batch_ms: 0,
1083                    encoding: Encoding::Json,
1084                    authorized: &mut authorized,
1085                },
1086                ClientFrame::Ping,
1087            )
1088            .await
1089        );
1090        assert_eq!(
1091            sys_rx.try_recv().unwrap(),
1092            OutFrame::Text(r#"{"type":"pong"}"#.into())
1093        );
1094    }
1095
1096    #[actix_web::test]
1097    async fn client_cannot_subscribe_or_unsubscribe_reserved_sys_queue() {
1098        let (state, _cred, _token) = app_state_with_device().await;
1099        let mut forwarders = HashMap::new();
1100        let mut queues = StreamMap::new();
1101        let (sys_tx, sys_rx) = mpsc::channel(SYS_OUTBOUND_BUFFER);
1102        queues.insert(SYS_CHANNEL.to_string(), ReceiverStream::new(sys_rx));
1103        let mut authorized = true;
1104
1105        for frame in [
1106            ClientFrame::Subscribe {
1107                ch: SYS_CHANNEL.into(),
1108                since: None,
1109            },
1110            ClientFrame::Unsubscribe {
1111                ch: SYS_CHANNEL.into(),
1112            },
1113        ] {
1114            assert!(
1115                handle_client_frame(
1116                    ClientDispatchContext {
1117                        state: &state,
1118                        forwarders: &mut forwarders,
1119                        queues: &mut queues,
1120                        sys_tx: &sys_tx,
1121                        batch_ms: 0,
1122                        encoding: Encoding::Json,
1123                        authorized: &mut authorized,
1124                    },
1125                    frame,
1126                )
1127                .await
1128            );
1129            assert!(
1130                queues.contains_key(SYS_CHANNEL),
1131                "client channel operations must not remove reserved sys queue"
1132            );
1133        }
1134
1135        try_enqueue_sys(&sys_tx, pong_frame(Encoding::Json));
1136        let (channel, frame) = queues.next().await.expect("sys queue remains drainable");
1137        assert_eq!(channel, SYS_CHANNEL);
1138        assert_eq!(frame, OutFrame::Text(r#"{"type":"pong"}"#.into()));
1139    }
1140
1141    #[test]
1142    fn queued_agent_event_blocks_synthetic_terminal_replay() {
1143        let (tx, rx) = tokio::sync::broadcast::channel(4);
1144        assert!(can_attempt_terminal_replay(None, &rx));
1145        tx.send(bamboo_agent_core::AgentEvent::Token {
1146            content: "first-live-token".into(),
1147        })
1148        .expect("receiver is subscribed");
1149        assert!(
1150            !can_attempt_terminal_replay(None, &rx),
1151            "a queued live frame must win the race with synthetic terminal replay"
1152        );
1153    }
1154
1155    #[test]
1156    fn pending_or_running_runner_blocks_synthetic_terminal_replay() {
1157        let (_tx, rx) = tokio::sync::broadcast::channel(4);
1158        assert!(!can_attempt_terminal_replay(
1159            Some(&crate::app_state::AgentStatus::Pending),
1160            &rx,
1161        ));
1162        assert!(!can_attempt_terminal_replay(
1163            Some(&crate::app_state::AgentStatus::Running),
1164            &rx,
1165        ));
1166    }
1167
1168    #[actix_web::test]
1169    async fn current_runner_reread_blocks_stale_terminal_snapshot() {
1170        let (state, _cred, _token) = app_state_with_device().await;
1171        let (_tx, rx) = tokio::sync::broadcast::channel(4);
1172        let session_id = "runner-reserved-during-terminal-check";
1173
1174        assert!(
1175            current_runner_allows_terminal_replay(&state, session_id, &rx).await,
1176            "the initial no-runner snapshot permits a terminal check"
1177        );
1178
1179        {
1180            let mut runners = state.agent_runners.write().await;
1181            let mut runner = crate::app_state::AgentRunner::new();
1182            runner.status = crate::app_state::AgentStatus::Running;
1183            runners.insert(session_id.to_string(), runner);
1184        }
1185
1186        assert!(
1187            !current_runner_allows_terminal_replay(&state, session_id, &rx).await,
1188            "the post-await re-read must observe the newly Running runner"
1189        );
1190    }
1191
1192    #[actix_web::test]
1193    async fn expired_startup_subscribe_waits_for_locked_live_reconcile() {
1194        let dir = tempdir().expect("temporary app data");
1195        let state = web::Data::new(AppState::new(dir.path().to_path_buf()).await.unwrap());
1196        let session_id = "ws-admission-startup-race";
1197        let channel = format!("agent.{session_id}");
1198        let mut session = bamboo_agent_core::Session::new(session_id, "test-model");
1199        session.add_message(bamboo_agent_core::Message::user("slow startup"));
1200        crate::handlers::agent::events::mark_pending_turn(&mut session);
1201        session.metadata.insert(
1202            "execute.startup_handoff_at".to_string(),
1203            (chrono::Utc::now() - chrono::Duration::seconds(120)).to_rfc3339(),
1204        );
1205        state.save_session(&mut session).await;
1206
1207        // Exercise the real WS admission function, not just the forwarder. An
1208        // execute owner appearing around the async terminal checks must force a
1209        // live subscription until the locked exact-work CAS can run.
1210        let startup_guard =
1211            crate::handlers::agent::events::begin_execute_startup(state.get_ref(), session_id);
1212        let mut forwarders = HashMap::new();
1213        let mut queues = StreamMap::new();
1214        subscribe(
1215            &state,
1216            &mut forwarders,
1217            &mut queues,
1218            0,
1219            Encoding::Json,
1220            &channel,
1221            None,
1222        )
1223        .await;
1224
1225        assert!(
1226            tokio::time::timeout(Duration::from_millis(350), queues.next())
1227                .await
1228                .is_err(),
1229            "an in-flight execute owner must prevent a one-shot terminal"
1230        );
1231        drop(startup_guard);
1232
1233        let (_, terminal_event) = tokio::time::timeout(Duration::from_secs(2), queues.next())
1234            .await
1235            .expect("locked reconcile emits startup failure")
1236            .expect("agent queue remains installed");
1237        let OutFrame::Text(terminal_event) = terminal_event else {
1238            panic!("JSON subscription must emit text frames");
1239        };
1240        let terminal_event: serde_json::Value =
1241            serde_json::from_str(&terminal_event).expect("terminal event JSON");
1242        assert_eq!(terminal_event["event"]["type"], "error");
1243        assert!(terminal_event["event"]["message"]
1244            .as_str()
1245            .is_some_and(|message| message.contains("was not started")));
1246
1247        let (_, terminal_control) = tokio::time::timeout(Duration::from_secs(2), queues.next())
1248            .await
1249            .expect("terminal control follows failure")
1250            .expect("agent queue drains terminal control");
1251        let OutFrame::Text(terminal_control) = terminal_control else {
1252            panic!("JSON subscription must emit text frames");
1253        };
1254        let terminal_control: serde_json::Value =
1255            serde_json::from_str(&terminal_control).expect("terminal control JSON");
1256        assert_eq!(terminal_control["control"]["type"], "terminal");
1257
1258        let stored = state
1259            .storage
1260            .load_session(session_id)
1261            .await
1262            .expect("load session")
1263            .expect("stored session");
1264        assert_eq!(stored.last_run_status().as_deref(), Some("error"));
1265        assert!(crate::handlers::agent::events::startup_work_id(&stored).is_none());
1266    }
1267}