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}