Skip to main content

arete_server/websocket/
server.rs

1use crate::bus::{BusManager, BusMessage};
2use crate::cache::{cmp_seq, EntityCache, SnapshotBatchConfig};
3use crate::compression::maybe_compress;
4use crate::view::{ViewIndex, ViewSpec};
5use crate::websocket::auth::{
6    AuthContext, AuthDecision, AuthDeny, ConnectionAuthRequest, WebSocketAuthPlugin,
7};
8use crate::websocket::client_manager::{ClientManager, RateLimitConfig};
9use crate::websocket::frame::{
10    apply_wire_format, Frame, Mode, SnapshotEntity, SnapshotFrame, SortConfig, SortOrder,
11    SubscribedFrame, UnsubscribedFrame,
12};
13use crate::websocket::subscription::{
14    ClientMessage, RefreshAuthRequest, RefreshAuthResponse, SocketIssueMessage, Subscription,
15    SubscriptionQuery, Unsubscription, PROTOCOL_VERSION,
16};
17use crate::websocket::usage::{WebSocketUsageEmitter, WebSocketUsageEvent};
18use crate::WebSocketDeliveryConfig;
19use anyhow::Result;
20use bytes::Bytes;
21use futures_util::StreamExt;
22use serde::Serialize;
23use serde_json::Value;
24use std::collections::{HashMap, HashSet, VecDeque};
25use std::future::Future;
26use std::net::SocketAddr;
27use std::sync::Arc;
28use std::time::{Duration, Instant};
29use tokio::net::{TcpListener, TcpStream};
30use tokio::sync::{broadcast, watch};
31use tokio_tungstenite::{
32    accept_hdr_async,
33    tungstenite::{
34        handshake::server::{ErrorResponse as HandshakeErrorResponse, Request, Response},
35        http::{header::CONTENT_TYPE, StatusCode},
36        Error as WsError,
37    },
38};
39use tokio_util::sync::CancellationToken;
40use tokio_util::task::TaskTracker;
41use tracing::{debug, error, info, info_span, warn, Instrument};
42use uuid::Uuid;
43
44#[cfg(feature = "otel")]
45use crate::metrics::Metrics;
46
47#[derive(Clone, Default)]
48struct WsMetrics {
49    #[cfg(feature = "otel")]
50    inner: Option<Arc<Metrics>>,
51    #[cfg(test)]
52    probe: Option<Arc<DeliveryProbe>>,
53}
54
55/// Test-only mirror of the delivery instruments recorded through
56/// [`WsMetrics`], so the real-socket load harness can report them without the
57/// `otel` feature. Each counter has the meaning of the metric named beside it.
58#[cfg(test)]
59#[derive(Debug, Default)]
60pub(crate) struct DeliveryProbe {
61    /// `arete.ws.messages.sent`
62    pub(crate) messages_sent: std::sync::atomic::AtomicU64,
63    /// `arete.ws.subscription.lagged`
64    pub(crate) lag_events: std::sync::atomic::AtomicU64,
65    /// `arete.ws.subscription.dropped_updates`
66    pub(crate) dropped_updates: std::sync::atomic::AtomicU64,
67    /// `arete.ws.subscription.resnapshots`
68    pub(crate) resnapshots: std::sync::atomic::AtomicU64,
69    /// `arete.ws.collection.coalesced_updates`
70    pub(crate) coalesced_updates: std::sync::atomic::AtomicU64,
71    /// `arete.ws.collection.coalesced_flushes`
72    pub(crate) coalesced_flushes: std::sync::atomic::AtomicU64,
73    /// `arete.ws.delivery.stopped`, by reason
74    pub(crate) delivery_stopped: std::sync::Mutex<std::collections::BTreeMap<&'static str, u64>>,
75}
76
77#[cfg(test)]
78impl DeliveryProbe {
79    fn add(counter: &std::sync::atomic::AtomicU64, value: u64) {
80        counter.fetch_add(value, std::sync::atomic::Ordering::Relaxed);
81    }
82}
83
84impl WsMetrics {
85    #[cfg(feature = "otel")]
86    fn new(inner: Option<Arc<Metrics>>) -> Self {
87        Self {
88            inner,
89            #[cfg(test)]
90            probe: None,
91        }
92    }
93
94    #[cfg(test)]
95    fn probe(&self, record: impl FnOnce(&DeliveryProbe)) {
96        if let Some(probe) = &self.probe {
97            record(probe);
98        }
99    }
100
101    fn connection_opened(&self, metering_key: Option<&str>) {
102        #[cfg(not(feature = "otel"))]
103        let _ = metering_key;
104        #[cfg(feature = "otel")]
105        if let Some(metrics) = &self.inner {
106            if let Some(metering_key) = metering_key {
107                metrics.record_ws_connection_with_metering(metering_key);
108            } else {
109                metrics.record_ws_connection();
110            }
111        }
112    }
113
114    fn connection_closed(&self, duration_secs: f64, metering_key: Option<&str>) {
115        #[cfg(not(feature = "otel"))]
116        let _ = (duration_secs, metering_key);
117        #[cfg(feature = "otel")]
118        if let Some(metrics) = &self.inner {
119            if let Some(metering_key) = metering_key {
120                metrics.record_ws_disconnection_with_metering(duration_secs, metering_key);
121            } else {
122                metrics.record_ws_disconnection(duration_secs);
123            }
124        }
125    }
126
127    fn message_received(&self, metering_key: Option<&str>) {
128        #[cfg(not(feature = "otel"))]
129        let _ = metering_key;
130        #[cfg(feature = "otel")]
131        if let Some(metrics) = &self.inner {
132            if let Some(metering_key) = metering_key {
133                metrics.record_ws_message_received_with_metering(metering_key);
134            } else {
135                metrics.record_ws_message_received();
136            }
137        }
138    }
139
140    fn message_sent(&self) {
141        #[cfg(test)]
142        self.probe(|probe| DeliveryProbe::add(&probe.messages_sent, 1));
143        #[cfg(feature = "otel")]
144        if let Some(metrics) = &self.inner {
145            metrics.record_ws_message_sent();
146        }
147    }
148
149    fn subscription_created(&self, view: &str, metering_key: Option<&str>) {
150        #[cfg(not(feature = "otel"))]
151        let _ = (view, metering_key);
152        #[cfg(feature = "otel")]
153        if let Some(metrics) = &self.inner {
154            if let Some(metering_key) = metering_key {
155                metrics.record_subscription_created_with_metering(view, metering_key);
156            } else {
157                metrics.record_subscription_created(view);
158            }
159        }
160    }
161
162    fn subscription_removed(&self, view: &str, metering_key: Option<&str>) {
163        #[cfg(not(feature = "otel"))]
164        let _ = (view, metering_key);
165        #[cfg(feature = "otel")]
166        if let Some(metrics) = &self.inner {
167            if let Some(metering_key) = metering_key {
168                metrics.record_subscription_removed_with_metering(view, metering_key);
169            } else {
170                metrics.record_subscription_removed(view);
171            }
172        }
173    }
174
175    fn protocol_error(&self, code: &str) {
176        #[cfg(not(feature = "otel"))]
177        let _ = code;
178        #[cfg(feature = "otel")]
179        if let Some(metrics) = &self.inner {
180            metrics.record_ws_protocol_error(code);
181        }
182    }
183
184    fn subscription_lagged(&self, view: &str, skipped: u64) {
185        #[cfg(test)]
186        self.probe(|probe| {
187            DeliveryProbe::add(&probe.lag_events, 1);
188            DeliveryProbe::add(&probe.dropped_updates, skipped);
189        });
190        #[cfg(not(feature = "otel"))]
191        let _ = (view, skipped);
192        #[cfg(feature = "otel")]
193        if let Some(metrics) = &self.inner {
194            metrics.record_ws_subscription_lagged(view, skipped);
195        }
196    }
197
198    fn subscription_resnapshot(&self, view: &str) {
199        #[cfg(test)]
200        self.probe(|probe| DeliveryProbe::add(&probe.resnapshots, 1));
201        #[cfg(not(feature = "otel"))]
202        let _ = view;
203        #[cfg(feature = "otel")]
204        if let Some(metrics) = &self.inner {
205            metrics.record_ws_subscription_resnapshot(view);
206        }
207    }
208
209    fn collection_coalesced(&self, view: &str, updates: u64) {
210        #[cfg(test)]
211        self.probe(|probe| {
212            DeliveryProbe::add(&probe.coalesced_updates, updates);
213            DeliveryProbe::add(&probe.coalesced_flushes, 1);
214        });
215        #[cfg(not(feature = "otel"))]
216        let _ = (view, updates);
217        #[cfg(feature = "otel")]
218        if let Some(metrics) = &self.inner {
219            metrics.record_ws_collection_coalesced(view, updates);
220        }
221    }
222
223    fn delivery_stopped(&self, view: &str, reason: &'static str) {
224        #[cfg(test)]
225        self.probe(|probe| {
226            *probe
227                .delivery_stopped
228                .lock()
229                .expect("delivery probe lock poisoned")
230                .entry(reason)
231                .or_default() += 1;
232        });
233        #[cfg(not(feature = "otel"))]
234        let _ = (view, reason);
235        #[cfg(feature = "otel")]
236        if let Some(metrics) = &self.inner {
237            metrics.record_ws_delivery_stopped(view, reason);
238        }
239    }
240}
241
242async fn handle_refresh_auth(
243    client_id: Uuid,
244    refresh_req: &RefreshAuthRequest,
245    client_manager: &ClientManager,
246    auth_plugin: &Arc<dyn WebSocketAuthPlugin>,
247) {
248    let refresh_result: Result<AuthContext, String> = if let Some(signed_plugin) = auth_plugin
249        .as_any()
250        .downcast_ref::<crate::websocket::auth::SignedSessionAuthPlugin>()
251    {
252        signed_plugin
253            .verify_refresh_token(&refresh_req.token)
254            .await
255            .map_err(|error| error.reason)
256    } else {
257        Err("In-band auth refresh not supported with current auth plugin".to_string())
258    };
259
260    let response = match refresh_result {
261        Ok(new_context) => {
262            let expires_at = new_context.expires_at;
263            if client_manager.update_client_auth(client_id, new_context) {
264                RefreshAuthResponse {
265                    success: true,
266                    error: None,
267                    expires_at: Some(expires_at),
268                }
269            } else {
270                RefreshAuthResponse {
271                    success: false,
272                    error: Some("client-not-found".to_string()),
273                    expires_at: None,
274                }
275            }
276        }
277        Err(error) => {
278            let code = if error.contains("expired") {
279                "token-expired"
280            } else if error.contains("signature") {
281                "token-invalid-signature"
282            } else if error.contains("issuer") {
283                "token-invalid-issuer"
284            } else if error.contains("audience") {
285                "token-invalid-audience"
286            } else {
287                "token-invalid"
288            };
289            RefreshAuthResponse {
290                success: false,
291                error: Some(code.to_string()),
292                expires_at: None,
293            }
294        }
295    };
296
297    if let Ok(json) = serde_json::to_string(&response) {
298        let _ = client_manager.send_text_to_client(client_id, json).await;
299    }
300}
301
302async fn send_socket_issue(
303    client_id: Uuid,
304    client_manager: &ClientManager,
305    deny: &AuthDeny,
306    fatal: bool,
307    subscription_id: Option<String>,
308) {
309    let message = SocketIssueMessage::from_auth_deny(deny, fatal, subscription_id);
310    if let Ok(json) = serde_json::to_string(&message) {
311        let _ = client_manager.send_text_to_client(client_id, json).await;
312    }
313}
314
315async fn send_protocol_issue(
316    client_id: Uuid,
317    client_manager: &ClientManager,
318    metrics: &WsMetrics,
319    subscription_id: Option<String>,
320    code: &str,
321    message: impl Into<String>,
322) {
323    metrics.protocol_error(code);
324    let issue = SocketIssueMessage::protocol(subscription_id, code, message);
325    if let Ok(json) = serde_json::to_string(&issue) {
326        let _ = client_manager.send_text_to_client(client_id, json).await;
327    }
328}
329
330/// Send an already-built issue, for refusals that carry structured detail.
331async fn send_prepared_issue(
332    client_id: Uuid,
333    client_manager: &ClientManager,
334    metrics: &WsMetrics,
335    issue: SocketIssueMessage,
336) {
337    metrics.protocol_error(&issue.code);
338    if let Ok(json) = serde_json::to_string(&issue) {
339        let _ = client_manager.send_text_to_client(client_id, json).await;
340    }
341}
342
343fn key_class_label(key_class: arete_auth::KeyClass) -> &'static str {
344    match key_class {
345        arete_auth::KeyClass::Secret => "secret",
346        arete_auth::KeyClass::Publishable => "publishable",
347    }
348}
349
350fn emit_usage_event(
351    usage_emitter: &Option<Arc<dyn WebSocketUsageEmitter>>,
352    event: WebSocketUsageEvent,
353) {
354    if let Some(emitter) = usage_emitter.clone() {
355        tokio::spawn(async move {
356            emitter.emit(event).await;
357        });
358    }
359}
360
361fn usage_identity(
362    auth_context: Option<&AuthContext>,
363) -> (
364    Option<String>,
365    Option<String>,
366    Option<String>,
367    Option<String>,
368) {
369    match auth_context {
370        Some(context) => (
371            Some(context.metering_key.clone()),
372            Some(context.subject.clone()),
373            Some(key_class_label(context.key_class).to_string()),
374            context.deployment_id.clone(),
375        ),
376        None => (None, None, None, None),
377    }
378}
379
380fn emit_update_sent_for_client(
381    usage_emitter: &Option<Arc<dyn WebSocketUsageEmitter>>,
382    client_manager: &ClientManager,
383    client_id: Uuid,
384    view_id: &str,
385    bytes: usize,
386) {
387    let auth_context = client_manager.get_auth_context(client_id);
388    let (metering_key, subject, _, deployment_id) = usage_identity(auth_context.as_ref());
389    emit_usage_event(
390        usage_emitter,
391        WebSocketUsageEvent::UpdateSent {
392            client_id: client_id.to_string(),
393            deployment_id,
394            metering_key,
395            subject,
396            view_id: view_id.to_string(),
397            messages: 1,
398            bytes: bytes as u64,
399        },
400    );
401}
402
403#[derive(Clone)]
404struct SubscriptionContext {
405    client_id: Uuid,
406    client_manager: ClientManager,
407    bus_manager: BusManager,
408    entity_cache: EntityCache,
409    view_index: Arc<ViewIndex>,
410    usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
411    journal: Option<Arc<crate::journal::EventJournal>>,
412    metrics: WsMetrics,
413    delivery: WebSocketDeliveryConfig,
414    /// Cancelled when the server stops; every session ends through its normal
415    /// cleanup path rather than being dropped mid-flight.
416    shutdown: CancellationToken,
417}
418
419pub struct WebSocketServer {
420    bind_addr: SocketAddr,
421    client_manager: ClientManager,
422    bus_manager: BusManager,
423    entity_cache: EntityCache,
424    view_index: Arc<ViewIndex>,
425    max_clients: usize,
426    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
427    usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
428    rate_limit_config: Option<RateLimitConfig>,
429    journal: Option<Arc<crate::journal::EventJournal>>,
430    delivery: WebSocketDeliveryConfig,
431    #[cfg(feature = "otel")]
432    metrics: Option<Arc<Metrics>>,
433}
434
435impl WebSocketServer {
436    #[cfg(feature = "otel")]
437    pub fn new(
438        bind_addr: SocketAddr,
439        bus_manager: BusManager,
440        entity_cache: EntityCache,
441        view_index: Arc<ViewIndex>,
442        metrics: Option<Arc<Metrics>>,
443    ) -> Self {
444        Self {
445            bind_addr,
446            client_manager: ClientManager::new(),
447            bus_manager,
448            entity_cache,
449            view_index,
450            max_clients: 10_000,
451            auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
452            usage_emitter: None,
453            rate_limit_config: None,
454            journal: None,
455            delivery: WebSocketDeliveryConfig::default(),
456            metrics,
457        }
458    }
459
460    #[cfg(not(feature = "otel"))]
461    pub fn new(
462        bind_addr: SocketAddr,
463        bus_manager: BusManager,
464        entity_cache: EntityCache,
465        view_index: Arc<ViewIndex>,
466    ) -> Self {
467        Self {
468            bind_addr,
469            client_manager: ClientManager::new(),
470            bus_manager,
471            entity_cache,
472            view_index,
473            max_clients: 10_000,
474            auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
475            usage_emitter: None,
476            rate_limit_config: None,
477            journal: None,
478            delivery: WebSocketDeliveryConfig::default(),
479        }
480    }
481
482    pub fn with_max_clients(mut self, max_clients: usize) -> Self {
483        self.max_clients = max_clients;
484        self
485    }
486
487    pub fn with_auth_plugin(mut self, auth_plugin: Arc<dyn WebSocketAuthPlugin>) -> Self {
488        self.auth_plugin = auth_plugin;
489        self
490    }
491
492    pub fn with_usage_emitter(mut self, usage_emitter: Arc<dyn WebSocketUsageEmitter>) -> Self {
493        self.usage_emitter = Some(usage_emitter);
494        self
495    }
496
497    /// Serve replayable append subscriptions from the retained event journal.
498    pub fn with_journal(mut self, journal: Arc<crate::journal::EventJournal>) -> Self {
499        self.journal = Some(journal);
500        self
501    }
502
503    pub fn with_rate_limit_config(mut self, config: RateLimitConfig) -> Self {
504        self.rate_limit_config = Some(config);
505        self
506    }
507
508    pub fn with_delivery_config(mut self, config: WebSocketDeliveryConfig) -> Self {
509        self.delivery = config;
510        self
511    }
512
513    /// Bind the configured address and serve connections until the task is
514    /// dropped. Equivalent to [`into_acceptor`](Self::into_acceptor) followed by
515    /// [`ConnectionAcceptor::serve_listener`].
516    pub async fn start(self) -> Result<()> {
517        info!(
518            "Starting WebSocket server on {} (max_clients: {})",
519            self.bind_addr, self.max_clients
520        );
521        let listener = TcpListener::bind(&self.bind_addr).await?;
522        let (acceptor, _cleanup) = self.into_acceptor();
523        acceptor.serve_listener(listener).await
524    }
525
526    /// Split this server into the part that serves connections and the
527    /// client-manager cleanup task, leaving the caller to own the listener.
528    ///
529    /// The cleanup handle is returned rather than detached so a caller that
530    /// stops serving can stop it too.
531    pub(crate) fn into_acceptor(self) -> (ConnectionAcceptor, tokio::task::JoinHandle<()>) {
532        let client_manager = self
533            .rate_limit_config
534            .map(ClientManager::with_config)
535            .unwrap_or(self.client_manager);
536        let cleanup = client_manager.start_cleanup_task();
537
538        #[cfg(feature = "otel")]
539        let metrics = WsMetrics::new(self.metrics.clone());
540        #[cfg(not(feature = "otel"))]
541        let metrics = WsMetrics::default();
542
543        let acceptor = ConnectionAcceptor {
544            client_manager,
545            bus_manager: self.bus_manager,
546            entity_cache: self.entity_cache,
547            view_index: self.view_index,
548            max_clients: self.max_clients,
549            auth_plugin: self.auth_plugin,
550            usage_emitter: self.usage_emitter,
551            journal: self.journal,
552            delivery: self.delivery,
553            metrics,
554            shutdown: CancellationToken::new(),
555            sessions: TaskTracker::new(),
556        };
557        (acceptor, cleanup)
558    }
559}
560
561/// Serves already-accepted TCP connections against one server's buses, cache
562/// and views.
563///
564/// This is what [`WebSocketServer::start`] runs behind its listener, separated
565/// so that a caller that owns the listener (an application that terminates
566/// TLS itself, a test with an ephemeral port) can hand streams in without the
567/// server binding anything.
568#[derive(Clone)]
569pub(crate) struct ConnectionAcceptor {
570    client_manager: ClientManager,
571    bus_manager: BusManager,
572    entity_cache: EntityCache,
573    view_index: Arc<ViewIndex>,
574    max_clients: usize,
575    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
576    usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
577    journal: Option<Arc<crate::journal::EventJournal>>,
578    delivery: WebSocketDeliveryConfig,
579    metrics: WsMetrics,
580    shutdown: CancellationToken,
581    /// Sessions spawned by [`serve_listener`](Self::serve_listener), so a
582    /// stop can wait for them. Sessions a caller serves on its own tasks are
583    /// the caller's to wait for.
584    sessions: TaskTracker,
585}
586
587impl ConnectionAcceptor {
588    /// Mirror this acceptor's delivery instruments into `probe`. Must be set
589    /// before serving: each session copies the metrics handle when it starts.
590    #[cfg(test)]
591    pub(crate) fn with_delivery_probe(mut self, probe: Arc<DeliveryProbe>) -> Self {
592        self.metrics.probe = Some(probe);
593        self
594    }
595
596    /// Number of clients currently connected to this server.
597    pub(crate) fn client_count(&self) -> usize {
598        self.client_manager.client_count()
599    }
600
601    /// End every session this acceptor is serving and stop accepting.
602    ///
603    /// Sessions notice on their next poll and leave through the same cleanup
604    /// as a client disconnect, so the client manager, buses and usage events
605    /// see an ordinary close.
606    pub(crate) fn shutdown(&self) {
607        self.shutdown.cancel();
608        self.sessions.close();
609    }
610
611    /// Resolves once every listener-spawned session has finished cleaning
612    /// up. Call after [`shutdown`](Self::shutdown).
613    pub(crate) async fn wait_for_sessions(&self) {
614        self.sessions.wait().await;
615    }
616
617    /// Serve one accepted connection: WebSocket handshake, authentication,
618    /// then the subscription session until the peer disconnects.
619    ///
620    /// Returns `Ok(())` without serving when the server is at its client
621    /// limit, exactly as the listener loop does.
622    pub(crate) async fn serve(&self, stream: TcpStream, remote_addr: SocketAddr) -> Result<()> {
623        if self.client_manager.client_count() >= self.max_clients {
624            warn!(
625                "Rejecting connection from {}: max clients reached",
626                remote_addr
627            );
628            return Ok(());
629        }
630
631        let context = SubscriptionContext {
632            client_id: Uuid::nil(),
633            client_manager: self.client_manager.clone(),
634            bus_manager: self.bus_manager.clone(),
635            entity_cache: self.entity_cache.clone(),
636            view_index: self.view_index.clone(),
637            usage_emitter: self.usage_emitter.clone(),
638            journal: self.journal.clone(),
639            delivery: self.delivery.clone(),
640            metrics: self.metrics.clone(),
641            shutdown: self.shutdown.clone(),
642        };
643        handle_connection(stream, context, remote_addr, self.auth_plugin.clone()).await
644    }
645
646    /// Accept from `listener` until [`shutdown`](Self::shutdown), serving each
647    /// connection on its own task.
648    pub(crate) async fn serve_listener(self, listener: TcpListener) -> Result<()> {
649        loop {
650            let accepted = tokio::select! {
651                _ = self.shutdown.cancelled() => return Ok(()),
652                accepted = listener.accept() => accepted,
653            };
654            match accepted {
655                Ok((stream, addr)) => {
656                    let acceptor = self.clone();
657                    self.sessions.spawn(
658                        async move {
659                            if let Err(error) = acceptor.serve(stream, addr).await {
660                                error!("WebSocket connection error: {}", error);
661                            }
662                        }
663                        .instrument(info_span!("ws.connection", %addr)),
664                    );
665                }
666                Err(error) => error!("Failed to accept connection: {}", error),
667            }
668        }
669    }
670}
671
672#[derive(Debug, Clone)]
673struct HandshakeReject {
674    status: StatusCode,
675    body: crate::websocket::auth::ErrorResponse,
676    error_code: String,
677    retry_after_secs: Option<u64>,
678}
679
680impl HandshakeReject {
681    fn from_deny(deny: &AuthDeny) -> Self {
682        let retry_after_secs = match deny.retry_policy {
683            crate::websocket::auth::RetryPolicy::RetryAfter(duration) => Some(duration.as_secs()),
684            _ => None,
685        };
686        Self {
687            status: StatusCode::from_u16(deny.http_status).unwrap_or(StatusCode::UNAUTHORIZED),
688            body: deny.to_error_response(),
689            error_code: deny.code.to_string(),
690            retry_after_secs,
691        }
692    }
693}
694
695fn build_handshake_error_response(
696    response: &Response,
697    reject: &HandshakeReject,
698) -> HandshakeErrorResponse {
699    let mut builder = Response::builder()
700        .status(reject.status)
701        .version(response.version())
702        .header(CONTENT_TYPE, "application/json; charset=utf-8")
703        .header("X-Error-Code", &reject.error_code)
704        .header("Cache-Control", "no-store");
705    if let Some(retry_after_secs) = reject.retry_after_secs {
706        builder = builder.header("Retry-After", retry_after_secs.to_string());
707    }
708    let body = serde_json::to_string(&reject.body).unwrap_or_else(|_| {
709        format!(
710            r#"{{"error":"{}","message":"{}","code":"{}","retryable":false}}"#,
711            reject.body.error, reject.body.message, reject.body.code
712        )
713    });
714    builder
715        .body(Some(body))
716        .expect("handshake rejection response should build")
717}
718
719#[allow(clippy::result_large_err)]
720async fn accept_authorized_connection(
721    stream: TcpStream,
722    remote_addr: SocketAddr,
723    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
724    client_manager: ClientManager,
725) -> Result<Option<(tokio_tungstenite::WebSocketStream<TcpStream>, AuthContext)>> {
726    use std::sync::Mutex;
727
728    let capture: Arc<Mutex<Option<Result<AuthContext, HandshakeReject>>>> =
729        Arc::new(Mutex::new(None));
730    let capture_ref = capture.clone();
731    let auth_plugin_ref = auth_plugin.clone();
732    let manager_ref = client_manager.clone();
733
734    let handshake_result = accept_hdr_async(stream, move |request: &Request, response| {
735        let request = ConnectionAuthRequest::from_http_request(remote_addr, request);
736        let result = tokio::task::block_in_place(|| {
737            tokio::runtime::Handle::current().block_on(async {
738                match auth_plugin_ref.authorize(&request).await {
739                    AuthDecision::Allow(context) => manager_ref
740                        .check_connection_allowed(remote_addr, &Some(context.clone()))
741                        .await
742                        .map(|()| context)
743                        .map_err(|deny| HandshakeReject::from_deny(&deny)),
744                    AuthDecision::Deny(deny) => Err(HandshakeReject::from_deny(&deny)),
745                }
746            })
747        });
748        *capture_ref.lock().expect("capture lock poisoned") = Some(result.clone());
749        match result {
750            Ok(_) => Ok(response),
751            Err(reject) => Err(build_handshake_error_response(&response, &reject)),
752        }
753    })
754    .await;
755
756    let auth_result = capture.lock().expect("capture lock poisoned").take();
757    match handshake_result {
758        Ok(stream) => match auth_result {
759            Some(Ok(context)) => Ok(Some((stream, context))),
760            Some(Err(reject)) => Err(anyhow::anyhow!(
761                "handshake unexpectedly succeeded after rejection: {}",
762                reject.body.message
763            )),
764            None => Err(anyhow::anyhow!("no auth result captured during handshake")),
765        },
766        Err(WsError::Http(_)) => Ok(None),
767        Err(error) => Err(error.into()),
768    }
769}
770
771async fn handle_connection(
772    stream: TcpStream,
773    mut context: SubscriptionContext,
774    remote_addr: SocketAddr,
775    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
776) -> Result<()> {
777    // The handshake is raced against shutdown too: a peer that stalls it must
778    // not keep a task alive after the server has stopped.
779    let accepted = tokio::select! {
780        _ = context.shutdown.cancelled() => return Ok(()),
781        accepted = accept_authorized_connection(
782            stream,
783            remote_addr,
784            auth_plugin.clone(),
785            context.client_manager.clone(),
786        ) => accepted?,
787    };
788    let Some((ws_stream, auth_context)) = accepted else {
789        return Ok(());
790    };
791
792    let client_id = Uuid::new_v4();
793    context.client_id = client_id;
794    let connection_start = Instant::now();
795    let (metering_key, subject, key_class, deployment_id) = usage_identity(Some(&auth_context));
796    context.metrics.connection_opened(metering_key.as_deref());
797
798    let (ws_sender, mut ws_receiver) = ws_stream.split();
799    context
800        .client_manager
801        .add_client(client_id, ws_sender, Some(auth_context), remote_addr);
802    emit_usage_event(
803        &context.usage_emitter,
804        WebSocketUsageEvent::ConnectionEstablished {
805            client_id: client_id.to_string(),
806            remote_addr: remote_addr.to_string(),
807            deployment_id: deployment_id.clone(),
808            metering_key: metering_key.clone(),
809            subject: subject.clone(),
810            key_class,
811        },
812    );
813
814    let mut active_subscriptions: HashMap<String, String> = HashMap::new();
815    loop {
816        let message = tokio::select! {
817            _ = context.shutdown.cancelled() => break,
818            next = ws_receiver.next() => match next {
819                Some(message) => message,
820                None => break,
821            },
822        };
823        let message = match message {
824            Ok(message) => message,
825            Err(error) => {
826                warn!("WebSocket error for client {}: {}", client_id, error);
827                break;
828            }
829        };
830        if message.is_close() {
831            break;
832        }
833        context.client_manager.update_client_last_seen(client_id);
834        if !message.is_text() {
835            continue;
836        }
837        if let Err(deny) = context
838            .client_manager
839            .check_inbound_message_allowed(client_id)
840        {
841            send_socket_issue(client_id, &context.client_manager, &deny, true, None).await;
842            break;
843        }
844        context.metrics.message_received(metering_key.as_deref());
845
846        let text = match message.to_text() {
847            Ok(text) => text,
848            Err(_) => continue,
849        };
850        let client_message = match serde_json::from_str::<ClientMessage>(text) {
851            Ok(message) => message,
852            Err(parse_error) => {
853                let subscription_id = extract_subscription_id(text);
854                send_protocol_issue(
855                    client_id,
856                    &context.client_manager,
857                    &context.metrics,
858                    subscription_id,
859                    "malformed-message",
860                    format!("invalid protocol v2 message: {parse_error}"),
861                )
862                .await;
863                continue;
864            }
865        };
866
867        match client_message {
868            ClientMessage::Subscribe(subscription) => {
869                let subscription_id = subscription.subscription_id.clone();
870                if let Err(message) = subscription.validate() {
871                    send_protocol_issue(
872                        client_id,
873                        &context.client_manager,
874                        &context.metrics,
875                        Some(subscription_id),
876                        "invalid-subscription",
877                        message,
878                    )
879                    .await;
880                    continue;
881                }
882                if let Err(deny) = context
883                    .client_manager
884                    .check_subscription_allowed(client_id)
885                    .await
886                {
887                    send_socket_issue(
888                        client_id,
889                        &context.client_manager,
890                        &deny,
891                        false,
892                        Some(subscription_id),
893                    )
894                    .await;
895                    continue;
896                }
897
898                let cancel_token = CancellationToken::new();
899                if !context
900                    .client_manager
901                    .add_client_subscription(
902                        client_id,
903                        subscription_id.clone(),
904                        cancel_token.clone(),
905                    )
906                    .await
907                {
908                    send_protocol_issue(
909                        client_id,
910                        &context.client_manager,
911                        &context.metrics,
912                        Some(subscription_id),
913                        "duplicate-subscription-id",
914                        "subscriptionId is already active on this connection",
915                    )
916                    .await;
917                    continue;
918                }
919
920                let view = subscription.query.view.clone();
921                if let Err(error) = attach_client_to_bus(&context, subscription, cancel_token).await
922                {
923                    context
924                        .client_manager
925                        .remove_client_subscription(client_id, &subscription_id)
926                        .await;
927                    // A refusal that already knows what to tell the client
928                    // (an expired cursor, a changed epoch) keeps its own
929                    // frame; anything else is a generic rejection.
930                    match error.downcast::<RejectedSubscription>() {
931                        Ok(rejected) => {
932                            send_prepared_issue(
933                                client_id,
934                                &context.client_manager,
935                                &context.metrics,
936                                rejected.0,
937                            )
938                            .await;
939                        }
940                        Err(error) => {
941                            send_protocol_issue(
942                                client_id,
943                                &context.client_manager,
944                                &context.metrics,
945                                Some(subscription_id),
946                                "subscription-rejected",
947                                error.to_string(),
948                            )
949                            .await;
950                        }
951                    }
952                    continue;
953                }
954
955                active_subscriptions.insert(subscription_id, view.clone());
956                context
957                    .metrics
958                    .subscription_created(&view, metering_key.as_deref());
959                emit_usage_event(
960                    &context.usage_emitter,
961                    WebSocketUsageEvent::SubscriptionCreated {
962                        client_id: client_id.to_string(),
963                        deployment_id: deployment_id.clone(),
964                        metering_key: metering_key.clone(),
965                        subject: subject.clone(),
966                        view_id: view,
967                    },
968                );
969            }
970            ClientMessage::Unsubscribe(unsubscription) => {
971                handle_unsubscribe(
972                    &context,
973                    unsubscription,
974                    &mut active_subscriptions,
975                    metering_key.as_deref(),
976                    &deployment_id,
977                    &metering_key,
978                    &subject,
979                )
980                .await;
981            }
982            ClientMessage::Ping => debug!("Received ping from client {}", client_id),
983            ClientMessage::RefreshAuth(request) => {
984                handle_refresh_auth(client_id, &request, &context.client_manager, &auth_plugin)
985                    .await;
986            }
987        }
988    }
989
990    context
991        .client_manager
992        .cancel_all_client_subscriptions(client_id)
993        .await;
994    context.client_manager.remove_client(client_id);
995    if let Some(rate_limiter) = context.client_manager.rate_limiter().cloned() {
996        rate_limiter.remove_client_buckets(client_id).await;
997    }
998    for view in active_subscriptions.values() {
999        context
1000            .metrics
1001            .subscription_removed(view, metering_key.as_deref());
1002        emit_usage_event(
1003            &context.usage_emitter,
1004            WebSocketUsageEvent::SubscriptionRemoved {
1005                client_id: client_id.to_string(),
1006                deployment_id: deployment_id.clone(),
1007                metering_key: metering_key.clone(),
1008                subject: subject.clone(),
1009                view_id: view.clone(),
1010            },
1011        );
1012    }
1013    let duration = connection_start.elapsed().as_secs_f64();
1014    context
1015        .metrics
1016        .connection_closed(duration, metering_key.as_deref());
1017    emit_usage_event(
1018        &context.usage_emitter,
1019        WebSocketUsageEvent::ConnectionClosed {
1020            client_id: client_id.to_string(),
1021            deployment_id,
1022            metering_key,
1023            subject,
1024            duration_secs: Some(duration),
1025            subscription_count: u32::try_from(active_subscriptions.len()).unwrap_or(u32::MAX),
1026        },
1027    );
1028    Ok(())
1029}
1030
1031#[allow(clippy::too_many_arguments)]
1032async fn handle_unsubscribe(
1033    context: &SubscriptionContext,
1034    unsubscription: Unsubscription,
1035    active_subscriptions: &mut HashMap<String, String>,
1036    metrics_metering_key: Option<&str>,
1037    deployment_id: &Option<String>,
1038    usage_metering_key: &Option<String>,
1039    subject: &Option<String>,
1040) {
1041    let subscription_id = unsubscription.subscription_id.clone();
1042    if let Err(message) = unsubscription.validate() {
1043        send_protocol_issue(
1044            context.client_id,
1045            &context.client_manager,
1046            &context.metrics,
1047            Some(subscription_id),
1048            "invalid-unsubscription",
1049            message,
1050        )
1051        .await;
1052        return;
1053    }
1054
1055    if !context
1056        .client_manager
1057        .remove_client_subscription(context.client_id, &subscription_id)
1058        .await
1059    {
1060        send_protocol_issue(
1061            context.client_id,
1062            &context.client_manager,
1063            &context.metrics,
1064            Some(subscription_id),
1065            "unknown-subscription-id",
1066            "subscriptionId is not active on this connection",
1067        )
1068        .await;
1069        return;
1070    }
1071
1072    let Some(view) = active_subscriptions.remove(&subscription_id) else {
1073        return;
1074    };
1075    let _ = send_control_frame(context, &UnsubscribedFrame::new(subscription_id), &view);
1076    context
1077        .metrics
1078        .subscription_removed(&view, metrics_metering_key);
1079    emit_usage_event(
1080        &context.usage_emitter,
1081        WebSocketUsageEvent::SubscriptionRemoved {
1082            client_id: context.client_id.to_string(),
1083            deployment_id: deployment_id.clone(),
1084            metering_key: usage_metering_key.clone(),
1085            subject: subject.clone(),
1086            view_id: view,
1087        },
1088    );
1089}
1090
1091fn extract_subscription_id(text: &str) -> Option<String> {
1092    serde_json::from_str::<Value>(text)
1093        .ok()?
1094        .get("subscriptionId")?
1095        .as_str()
1096        .map(str::to_string)
1097}
1098
1099struct SnapshotMetadata<'a> {
1100    subscription_id: &'a str,
1101    snapshot_id: &'a str,
1102    authoritative: bool,
1103    mode: Mode,
1104    view_id: &'a str,
1105    key: Option<&'a str>,
1106}
1107
1108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1109enum SnapshotPurpose {
1110    Initial,
1111    Recovery,
1112}
1113
1114impl SnapshotPurpose {
1115    fn authoritative(self, subscription: &Subscription) -> bool {
1116        match self {
1117            Self::Initial => subscription.query.after.is_none(),
1118            // Recovery replaces the exact query membership even when `after`
1119            // made the initial snapshot incremental. A merge-only snapshot
1120            // cannot remove members whose delete frames were skipped.
1121            Self::Recovery => true,
1122        }
1123    }
1124}
1125
1126fn create_snapshot_batches(
1127    entities: &[SnapshotEntity],
1128    metadata: SnapshotMetadata<'_>,
1129    batch_config: &SnapshotBatchConfig,
1130) -> Vec<SnapshotFrame> {
1131    if entities.is_empty() {
1132        return vec![SnapshotFrame {
1133            protocol_version: PROTOCOL_VERSION,
1134            subscription_id: metadata.subscription_id.to_string(),
1135            snapshot_id: metadata.snapshot_id.to_string(),
1136            authoritative: metadata.authoritative,
1137            mode: metadata.mode,
1138            export: metadata.view_id.to_string(),
1139            op: "snapshot",
1140            key: metadata.key.map(str::to_string),
1141            data: vec![],
1142            complete: true,
1143        }];
1144    }
1145
1146    let mut batches = Vec::new();
1147    let mut offset = 0;
1148    while offset < entities.len() {
1149        let configured_size = if offset == 0 {
1150            batch_config.initial_batch_size
1151        } else {
1152            batch_config.subsequent_batch_size
1153        };
1154        let end = (offset + configured_size.max(1)).min(entities.len());
1155        batches.push(SnapshotFrame {
1156            protocol_version: PROTOCOL_VERSION,
1157            subscription_id: metadata.subscription_id.to_string(),
1158            snapshot_id: metadata.snapshot_id.to_string(),
1159            authoritative: metadata.authoritative,
1160            mode: metadata.mode,
1161            export: metadata.view_id.to_string(),
1162            op: "snapshot",
1163            key: metadata.key.map(str::to_string),
1164            data: entities[offset..end].to_vec(),
1165            complete: end == entities.len(),
1166        });
1167        offset = end;
1168    }
1169    batches
1170}
1171
1172async fn send_snapshot_batches(
1173    context: &SubscriptionContext,
1174    subscription: &Subscription,
1175    entities: &[SnapshotEntity],
1176    mode: Mode,
1177    purpose: SnapshotPurpose,
1178    batch_config: &SnapshotBatchConfig,
1179) -> Result<()> {
1180    let snapshot_id = Uuid::new_v4().to_string();
1181    let authoritative = purpose.authoritative(subscription);
1182    let frames = create_snapshot_batches(
1183        entities,
1184        SnapshotMetadata {
1185            subscription_id: &subscription.subscription_id,
1186            snapshot_id: &snapshot_id,
1187            authoritative,
1188            mode,
1189            view_id: &subscription.query.view,
1190            key: subscription.query.key.as_deref(),
1191        },
1192        batch_config,
1193    );
1194
1195    for frame in frames {
1196        let rows = frame.data.len() as u32;
1197        let json = serde_json::to_vec(&frame)?;
1198        let payload = maybe_compress(&json);
1199        let bytes = payload.as_bytes().len() as u64;
1200        context
1201            .client_manager
1202            .send_compressed_async(context.client_id, payload)
1203            .await
1204            .map_err(|error| anyhow::anyhow!("failed to send snapshot: {error}"))?;
1205        context.metrics.message_sent();
1206
1207        let auth_context = context.client_manager.get_auth_context(context.client_id);
1208        let (metering_key, subject, _, deployment_id) = usage_identity(auth_context.as_ref());
1209        emit_usage_event(
1210            &context.usage_emitter,
1211            WebSocketUsageEvent::SnapshotSent {
1212                client_id: context.client_id.to_string(),
1213                deployment_id,
1214                metering_key,
1215                subject,
1216                view_id: subscription.query.view.clone(),
1217                rows,
1218                messages: 1,
1219                bytes,
1220            },
1221        );
1222    }
1223    Ok(())
1224}
1225
1226fn extract_sort_config(view_spec: &ViewSpec) -> Option<SortConfig> {
1227    if let Some(sort) = view_spec
1228        .pipeline
1229        .as_ref()
1230        .and_then(|pipeline| pipeline.sort.as_ref())
1231    {
1232        return Some(SortConfig {
1233            field: sort.field_path.clone(),
1234            order: match sort.order {
1235                crate::materialized_view::SortOrder::Asc => SortOrder::Asc,
1236                crate::materialized_view::SortOrder::Desc => SortOrder::Desc,
1237            },
1238        });
1239    }
1240    (view_spec.mode == Mode::List).then(|| SortConfig {
1241        field: vec!["_seq".to_string()],
1242        order: SortOrder::Desc,
1243    })
1244}
1245
1246fn send_control_frame<T: Serialize>(
1247    context: &SubscriptionContext,
1248    frame: &T,
1249    view_id: &str,
1250) -> Result<()> {
1251    let json = serde_json::to_vec(frame)?;
1252    let bytes = json.len();
1253    context
1254        .client_manager
1255        .send_to_client(context.client_id, Arc::new(Bytes::from(json)))
1256        .map_err(|error| anyhow::anyhow!("failed to send control frame: {error}"))?;
1257    context.metrics.message_sent();
1258    emit_update_sent_for_client(
1259        &context.usage_emitter,
1260        &context.client_manager,
1261        context.client_id,
1262        view_id,
1263        bytes,
1264    );
1265    Ok(())
1266}
1267
1268fn send_subscribed_frame(
1269    context: &SubscriptionContext,
1270    subscription: &Subscription,
1271    view_spec: &ViewSpec,
1272) -> Result<()> {
1273    let frame = SubscribedFrame::new(
1274        subscription.subscription_id.clone(),
1275        subscription.query.clone(),
1276        view_spec.mode,
1277        extract_sort_config(view_spec),
1278    );
1279    send_control_frame(context, &frame, &subscription.query.view)
1280}
1281
1282fn enforce_snapshot_limit(context: &SubscriptionContext, rows: usize) -> Result<()> {
1283    context
1284        .client_manager
1285        .check_snapshot_allowed(context.client_id, u32::try_from(rows).unwrap_or(u32::MAX))
1286        .map_err(|deny| anyhow::anyhow!(deny.reason))
1287}
1288
1289async fn subscribe_state_then_snapshot<F, Fut, T>(
1290    bus_manager: &BusManager,
1291    view_id: &str,
1292    key: &str,
1293    snapshot: F,
1294) -> (watch::Receiver<Arc<Bytes>>, T)
1295where
1296    F: FnOnce() -> Fut,
1297    Fut: Future<Output = T>,
1298{
1299    let mut receiver = bus_manager.get_or_create_state_bus(view_id, key).await;
1300    receiver.borrow_and_update();
1301    let snapshot = snapshot().await;
1302    (receiver, snapshot)
1303}
1304
1305async fn subscribe_list_then_snapshot<F, Fut, T>(
1306    bus_manager: &BusManager,
1307    view_id: &str,
1308    snapshot: F,
1309) -> (broadcast::Receiver<Arc<BusMessage>>, T)
1310where
1311    F: FnOnce() -> Fut,
1312    Fut: Future<Output = T>,
1313{
1314    let receiver = bus_manager.get_or_create_list_bus(view_id).await;
1315    let snapshot = snapshot().await;
1316    (receiver, snapshot)
1317}
1318
1319async fn attach_client_to_bus(
1320    context: &SubscriptionContext,
1321    mut subscription: Subscription,
1322    cancel_token: CancellationToken,
1323) -> Result<()> {
1324    let view_spec = context
1325        .view_index
1326        .get_view(&subscription.query.view)
1327        .cloned()
1328        .ok_or_else(|| anyhow::anyhow!("unknown view: {}", subscription.query.view))?;
1329
1330    if view_spec.mode == Mode::State && !view_spec.is_derived() && subscription.query.key.is_none()
1331    {
1332        return Err(anyhow::anyhow!("state subscriptions require query.key"));
1333    }
1334    if view_spec.is_derived() && subscription.query.take.is_none() {
1335        subscription.query.take = view_spec
1336            .pipeline
1337            .as_ref()
1338            .and_then(|pipeline| pipeline.limit);
1339    }
1340
1341    // A retained tape takes precedence for append views: it is the only
1342    // delivery that can honour a cursor. Without one, fall through to the
1343    // previous latest-state behaviour.
1344    let journal = context
1345        .journal
1346        .clone()
1347        .filter(|journal| journal.is_enabled() && view_spec.mode == Mode::Append);
1348    if let Some(journal) = journal {
1349        return attach_journal_subscription(
1350            context,
1351            subscription,
1352            view_spec,
1353            journal,
1354            cancel_token,
1355        )
1356        .await;
1357    }
1358
1359    if view_spec.mode == Mode::State && !view_spec.is_derived() {
1360        attach_state_subscription(context, subscription, view_spec, cancel_token).await
1361    } else {
1362        attach_collection_subscription(context, subscription, view_spec, cancel_token).await
1363    }
1364}
1365
1366async fn attach_state_subscription(
1367    context: &SubscriptionContext,
1368    subscription: Subscription,
1369    view_spec: ViewSpec,
1370    cancel_token: CancellationToken,
1371) -> Result<()> {
1372    let view_id = subscription.query.view.clone();
1373    let key = subscription.query.key.clone().unwrap_or_default();
1374    let query = subscription.query.clone();
1375    let cache = context.entity_cache.clone();
1376    let view_spec_for_snapshot = view_spec.clone();
1377    let (mut receiver, initial) =
1378        subscribe_state_then_snapshot(&context.bus_manager, &view_id, &key, move || async move {
1379            load_query_entities(&cache, None, &view_spec_for_snapshot, &query, false).await
1380        })
1381        .await;
1382
1383    let mut snapshot_entities = initial.clone();
1384    if let Some(limit) = subscription.query.snapshot_limit {
1385        snapshot_entities.truncate(limit);
1386    }
1387    enforce_snapshot_limit(context, snapshot_entities.len())?;
1388    send_subscribed_frame(context, &subscription, &view_spec)?;
1389    if subscription.snapshot.enabled {
1390        send_snapshot_batches(
1391            context,
1392            &subscription,
1393            &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1394            view_spec.mode,
1395            SnapshotPurpose::Initial,
1396            &context.entity_cache.snapshot_config(),
1397        )
1398        .await?;
1399    }
1400
1401    let task_context = context.clone();
1402    let subscription_id = subscription.subscription_id.clone();
1403    let query = subscription.query.clone();
1404    let view_spec_task = view_spec.clone();
1405    let span_view = view_id.clone();
1406    let span_key = key.clone();
1407    tokio::spawn(
1408        async move {
1409            let mut member = !initial.is_empty();
1410            loop {
1411                tokio::select! {
1412                    _ = cancel_token.cancelled() => break,
1413                    changed = receiver.changed() => {
1414                        if changed.is_err() {
1415                            break;
1416                        }
1417                        let payload = receiver.borrow().clone();
1418                        let metadata = source_frame_metadata(&payload);
1419                        if metadata.op == "delete" {
1420                            task_context.entity_cache.remove(&query.view, &key).await;
1421                            if member && send_membership_frame(
1422                                &task_context,
1423                                &subscription_id,
1424                                &view_spec_task,
1425                                "delete",
1426                                &key,
1427                                Value::Null,
1428                                metadata.seq,
1429                            ).is_err() {
1430                                break;
1431                            }
1432                            member = false;
1433                            continue;
1434                        }
1435
1436                        let selected = load_query_entities(
1437                            &task_context.entity_cache,
1438                            None,
1439                            &view_spec_task,
1440                            &query,
1441                            false,
1442                        ).await;
1443                        let is_member = !selected.is_empty();
1444                        let result = match (member, is_member) {
1445                            (true, true) => send_scoped_source_payload(
1446                                &task_context,
1447                                &subscription_id,
1448                                &query.view,
1449                                payload,
1450                            ),
1451                            (false, true) => {
1452                                let (entity_key, data) = selected.into_iter().next().unwrap();
1453                                send_membership_frame(
1454                                    &task_context,
1455                                    &subscription_id,
1456                                    &view_spec_task,
1457                                    "upsert",
1458                                    &entity_key,
1459                                    data,
1460                                    metadata.seq,
1461                                )
1462                            }
1463                            (true, false) => send_membership_frame(
1464                                &task_context,
1465                                &subscription_id,
1466                                &view_spec_task,
1467                                "remove",
1468                                &key,
1469                                Value::Null,
1470                                metadata.seq,
1471                            ),
1472                            (false, false) => Ok(()),
1473                        };
1474                        if result.is_err() {
1475                            break;
1476                        }
1477                        member = is_member;
1478                    }
1479                }
1480            }
1481        }
1482        .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
1483    );
1484    Ok(())
1485}
1486
1487async fn attach_collection_subscription(
1488    context: &SubscriptionContext,
1489    subscription: Subscription,
1490    view_spec: ViewSpec,
1491    cancel_token: CancellationToken,
1492) -> Result<()> {
1493    let view_id = subscription.query.view.clone();
1494    let source_view_id = view_spec
1495        .source_view
1496        .clone()
1497        .unwrap_or_else(|| view_id.clone());
1498    let (mut receiver, initial_membership) = subscribe_collection_then_snapshot(
1499        context,
1500        &source_view_id,
1501        &view_spec,
1502        &subscription.query,
1503    )
1504    .await;
1505
1506    let mut snapshot_entities = initial_membership.clone();
1507    if let Some(limit) = subscription.query.snapshot_limit {
1508        snapshot_entities.truncate(limit);
1509    }
1510    enforce_snapshot_limit(context, snapshot_entities.len())?;
1511    send_subscribed_frame(context, &subscription, &view_spec)?;
1512    if subscription.snapshot.enabled {
1513        send_snapshot_batches(
1514            context,
1515            &subscription,
1516            &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1517            view_spec.mode,
1518            SnapshotPurpose::Initial,
1519            &context.entity_cache.snapshot_config(),
1520        )
1521        .await?;
1522    }
1523
1524    let task_context = context.clone();
1525    let task_subscription = subscription.clone();
1526    let subscription_id = subscription.subscription_id.clone();
1527    let query = subscription.query.clone();
1528    let view_spec_task = view_spec.clone();
1529    let span_view = view_id.clone();
1530    tokio::spawn(
1531        async move {
1532            let mut current = initial_membership;
1533            // Append views are event tapes: collapsing two records would lose
1534            // observable history. Coalescing is only valid for latest-state
1535            // list membership, where a full final entity preserves meaning.
1536            let coalesce_ms = (view_spec_task.mode == Mode::List)
1537                .then(|| {
1538                    view_spec_task
1539                        .delivery
1540                        .coalesce_ms
1541                        .or(task_context.delivery.collection_coalesce_ms)
1542                        .filter(|milliseconds| *milliseconds > 0)
1543                })
1544                .flatten();
1545            let mut flush_interval = coalesce_ms.map(|milliseconds| {
1546                let period = Duration::from_millis(milliseconds);
1547                let mut interval = tokio::time::interval_at(
1548                    tokio::time::Instant::now() + period,
1549                    period,
1550                );
1551                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
1552                interval
1553            });
1554            let mut pending = HashMap::<String, Arc<BusMessage>>::new();
1555            let mut pending_updates = 0_u64;
1556
1557            loop {
1558                tokio::select! {
1559                    _ = cancel_token.cancelled() => break,
1560                    _ = async {
1561                        flush_interval
1562                            .as_mut()
1563                            .expect("coalescing interval is guarded")
1564                            .tick()
1565                            .await;
1566                    }, if flush_interval.is_some() => {
1567                        if pending.is_empty() {
1568                            continue;
1569                        }
1570                        let sorted_caches = view_spec_task
1571                            .is_derived()
1572                            .then(|| task_context.view_index.sorted_caches());
1573                        let next = load_query_entities(
1574                            &task_context.entity_cache,
1575                            sorted_caches,
1576                            &view_spec_task,
1577                            &query,
1578                            false,
1579                        ).await;
1580                        if emit_coalesced_collection_delta(
1581                            &task_context,
1582                            &subscription_id,
1583                            &view_spec_task,
1584                            &current,
1585                            &next,
1586                            &pending,
1587                        ).is_err() {
1588                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
1589                            break;
1590                        }
1591                        task_context
1592                            .metrics
1593                            .collection_coalesced(&view_id, pending_updates);
1594                        current = next;
1595                        pending.clear();
1596                        pending_updates = 0;
1597                    }
1598                    received = receiver.recv() => {
1599                        let envelope = match received {
1600                            Ok(envelope) => envelope,
1601                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
1602                                task_context.metrics.subscription_lagged(&view_id, skipped);
1603                                warn!(
1604                                    "Subscription {} lagged by {} updates",
1605                                    subscription_id, skipped
1606                                );
1607                                if view_spec_task.mode == Mode::Append {
1608                                    let _ = send_control_frame(
1609                                        &task_context,
1610                                        &SocketIssueMessage::append_subscription_lagged(
1611                                            subscription_id.clone(),
1612                                            skipped,
1613                                        ),
1614                                        &view_id,
1615                                    );
1616                                    task_context.metrics.delivery_stopped(&view_id, "append-lagged-without-replay");
1617                                    break;
1618                                }
1619                                if !task_subscription.snapshot.enabled {
1620                                    let _ = send_control_frame(
1621                                        &task_context,
1622                                        &SocketIssueMessage::subscription_lagged(
1623                                            subscription_id.clone(),
1624                                            skipped,
1625                                        ),
1626                                        &view_id,
1627                                    );
1628                                    task_context.metrics.delivery_stopped(&view_id, "lagged-without-snapshot");
1629                                    break;
1630                                }
1631                                info!(
1632                                    "Subscription {} is recovering from an authoritative snapshot",
1633                                    subscription_id
1634                                );
1635                                match recover_collection_subscription(
1636                                    &task_context,
1637                                    &task_subscription,
1638                                    &view_spec_task,
1639                                    &source_view_id,
1640                                ).await {
1641                                    Ok((next_receiver, recovered)) => {
1642                                        receiver = next_receiver;
1643                                        current = recovered;
1644                                        pending.clear();
1645                                        pending_updates = 0;
1646                                        task_context.metrics.subscription_resnapshot(&view_id);
1647                                        continue;
1648                                    }
1649                                    Err(error) => {
1650                                        warn!(
1651                                            "Subscription {} failed to recover from lag: {error:#}",
1652                                            subscription_id
1653                                        );
1654                                        task_context.metrics.delivery_stopped(&view_id, "resnapshot-failed");
1655                                        break;
1656                                    }
1657                                }
1658                            }
1659                            Err(broadcast::error::RecvError::Closed) => break,
1660                        };
1661
1662                        apply_collection_source_event(
1663                            &task_context,
1664                            &source_view_id,
1665                            &view_spec_task,
1666                            &query,
1667                            &envelope,
1668                        ).await;
1669
1670                        if flush_interval.is_some() {
1671                            pending_updates = pending_updates.saturating_add(1);
1672                            pending.insert(envelope.key.clone(), envelope);
1673                            continue;
1674                        }
1675
1676                        let metadata = source_frame_metadata(&envelope.payload);
1677                        let sorted_caches = view_spec_task
1678                            .is_derived()
1679                            .then(|| task_context.view_index.sorted_caches());
1680                        let next = load_query_entities(
1681                            &task_context.entity_cache,
1682                            sorted_caches,
1683                            &view_spec_task,
1684                            &query,
1685                            false,
1686                        ).await;
1687                        if emit_collection_delta(
1688                            &task_context,
1689                            &subscription_id,
1690                            &view_spec_task,
1691                            &current,
1692                            &next,
1693                            &envelope,
1694                            &metadata,
1695                        ).is_err() {
1696                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
1697                            break;
1698                        }
1699                        current = next;
1700                    }
1701                }
1702            }
1703        }
1704        .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
1705    );
1706    Ok(())
1707}
1708
1709async fn subscribe_collection_then_snapshot(
1710    context: &SubscriptionContext,
1711    source_view_id: &str,
1712    view_spec: &ViewSpec,
1713    query: &SubscriptionQuery,
1714) -> (broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>) {
1715    let cache = context.entity_cache.clone();
1716    let sorted_caches = view_spec
1717        .is_derived()
1718        .then(|| context.view_index.sorted_caches());
1719    let view_spec = view_spec.clone();
1720    let query = query.clone();
1721    subscribe_list_then_snapshot(&context.bus_manager, source_view_id, move || async move {
1722        load_query_entities(&cache, sorted_caches, &view_spec, &query, false).await
1723    })
1724    .await
1725}
1726
1727async fn recover_collection_subscription(
1728    context: &SubscriptionContext,
1729    subscription: &Subscription,
1730    view_spec: &ViewSpec,
1731    source_view_id: &str,
1732) -> Result<(broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>)> {
1733    let (receiver, membership) =
1734        subscribe_collection_then_snapshot(context, source_view_id, view_spec, &subscription.query)
1735            .await;
1736    // `snapshotLimit` caps only the initial transfer. Recovery has to replace
1737    // the complete live membership; truncating it would make the replacement
1738    // authoritative while immediately omitting members the server still
1739    // considers current.
1740    enforce_snapshot_limit(context, membership.len())?;
1741    send_snapshot_batches(
1742        context,
1743        subscription,
1744        &to_wire_snapshot_entities(membership.clone(), view_spec),
1745        view_spec.mode,
1746        SnapshotPurpose::Recovery,
1747        &context.entity_cache.snapshot_config(),
1748    )
1749    .await?;
1750    Ok((receiver, membership))
1751}
1752
1753async fn apply_collection_source_event(
1754    context: &SubscriptionContext,
1755    source_view_id: &str,
1756    view_spec: &ViewSpec,
1757    query: &SubscriptionQuery,
1758    envelope: &BusMessage,
1759) {
1760    let metadata = source_frame_metadata(&envelope.payload);
1761    if metadata.op != "delete" {
1762        return;
1763    }
1764    // A slow subscription can observe an old delete after the projector has
1765    // already recreated the key. Never let subscriber-local lag erase newer
1766    // shared cache state.
1767    let current = context
1768        .entity_cache
1769        .get(source_view_id, &envelope.key)
1770        .await;
1771    if source_delete_is_stale(current.as_ref(), metadata.seq.as_deref()) {
1772        return;
1773    }
1774    context
1775        .entity_cache
1776        .remove(source_view_id, &envelope.key)
1777        .await;
1778    if view_spec.is_derived() {
1779        let caches = context.view_index.sorted_caches();
1780        let mut guard = caches.write().await;
1781        if let Some(cache) = guard.get_mut(&query.view) {
1782            cache.remove(&envelope.key);
1783        }
1784    }
1785}
1786
1787fn source_delete_is_stale(current: Option<&Value>, delete_seq: Option<&str>) -> bool {
1788    match (current, delete_seq) {
1789        (Some(current), Some(delete_seq)) => current
1790            .get("_seq")
1791            .and_then(Value::as_str)
1792            .is_some_and(|current_seq| {
1793                cmp_seq(current_seq, delete_seq) == std::cmp::Ordering::Greater
1794            }),
1795        _ => false,
1796    }
1797}
1798
1799/// A subscription refused for a reason the client needs spelled out.
1800///
1801/// Attach paths that return this get the registration released by the
1802/// connection loop, the same as any other failure, while the client still
1803/// receives the specific error rather than a generic `subscription-rejected`.
1804/// Sending the frame and returning `Ok` instead would leave a registered
1805/// subscription with nothing attached: it would hold a slot against the
1806/// client's limit and make the advertised "resubscribe" remediation fail with
1807/// `duplicate-subscription-id`.
1808#[derive(Debug)]
1809pub(crate) struct RejectedSubscription(pub SocketIssueMessage);
1810
1811impl RejectedSubscription {
1812    fn into_error(issue: SocketIssueMessage) -> anyhow::Error {
1813        anyhow::Error::new(Self(issue))
1814    }
1815}
1816
1817impl std::fmt::Display for RejectedSubscription {
1818    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1819        formatter.write_str(&self.0.message)
1820    }
1821}
1822
1823impl std::error::Error for RejectedSubscription {}
1824
1825/// Deliver an append view as an event tape: replay the retained records after
1826/// the cursor, then forward live frames.
1827///
1828/// This deliberately does not recompute membership from the entity cache the
1829/// way [`attach_collection_subscription`] does. The cache folds each patch
1830/// into the resident entity, so a membership diff cannot express "these three
1831/// events happened"; the retained records can.
1832async fn attach_journal_subscription(
1833    context: &SubscriptionContext,
1834    subscription: Subscription,
1835    view_spec: ViewSpec,
1836    journal: Arc<crate::journal::EventJournal>,
1837    cancel_token: CancellationToken,
1838) -> Result<()> {
1839    let view_id = view_spec.id.clone();
1840    let subscription_id = subscription.subscription_id.clone();
1841
1842    // A tape has no membership window, so `take`/`skip` cannot mean what they
1843    // mean on a list view. Refuse them rather than accept and ignore them.
1844    if subscription.query.take.is_some()
1845        || subscription.query.skip.is_some()
1846        || subscription.query.snapshot_limit.is_some()
1847    {
1848        return Err(RejectedSubscription::into_error(SocketIssueMessage::protocol(
1849            Some(subscription_id),
1850            "invalid-subscription",
1851            format!(
1852                "take, skip and snapshotLimit are window options and do not apply to the replayable view {view_id}"
1853            ),
1854        )));
1855    }
1856
1857    // Reject a malformed cursor rather than silently replaying from the start,
1858    // which would look like success and duplicate everything.
1859    let cursor = match subscription.query.after.as_deref() {
1860        Some(raw) => match crate::journal::Cursor::parse(raw) {
1861            Some(cursor) => Some(cursor),
1862            None => {
1863                return Err(RejectedSubscription::into_error(
1864                    SocketIssueMessage::protocol(
1865                        Some(subscription_id),
1866                        "invalid-cursor",
1867                        format!(
1868                            "`after` must be an {{epoch}}:{{offset}} replay cursor for view {view_id}, got {raw:?}"
1869                        ),
1870                    ),
1871                ));
1872            }
1873        },
1874        None => None,
1875    };
1876
1877    // Subscribe before reading the journal so anything published during the
1878    // replay is still delivered; the offset filter below drops the overlap.
1879    let mut receiver = context.bus_manager.get_or_create_list_bus(&view_id).await;
1880
1881    let replayed = match journal.replay_after(&view_id, cursor.as_ref()).await {
1882        Ok(records) => records,
1883        Err(error) => {
1884            return Err(RejectedSubscription::into_error(
1885                SocketIssueMessage::replay_refused(Some(subscription_id), &error),
1886            ));
1887        }
1888    };
1889
1890    let frame = SubscribedFrame::new(
1891        subscription.subscription_id.clone(),
1892        subscription.query.clone(),
1893        view_spec.mode,
1894        extract_sort_config(&view_spec),
1895    )
1896    .with_replay_window(journal.window(&view_id).await);
1897    send_control_frame(context, &frame, &view_id)?;
1898
1899    // Everything past the acknowledgement runs on its own task. The replay can
1900    // be long and applies real backpressure, and `attach_client_to_bus` is
1901    // awaited directly on the connection's inbound loop — doing it there would
1902    // block unsubscribe, auth refresh and pong for the whole replay.
1903    let task_context = context.clone();
1904    let task_subscription_id = subscription.subscription_id.clone();
1905    let task_query = subscription.query.clone();
1906    let task_epoch = journal.epoch().await;
1907    let span_view = view_id.clone();
1908    tokio::spawn(
1909        async move {
1910            let mut last_sent = cursor.map(|cursor| cursor.offset);
1911            // Frames that published while the replay was still running. The
1912            // bus is a bounded broadcast, so it has to be drained as we go or
1913            // a busy view laps us before the replay finishes.
1914            let mut pending: VecDeque<Arc<BusMessage>> = VecDeque::new();
1915            let mut lagged: Option<u64> = None;
1916
1917            for record in replayed {
1918                if cancel_token.is_cancelled() {
1919                    return;
1920                }
1921                drain_available(&mut receiver, &mut pending, &mut lagged);
1922                if journal_record_matches(&task_query, &record)
1923                    && send_scoped_source_payload_async(
1924                        &task_context,
1925                        &task_subscription_id,
1926                        &span_view,
1927                        record.payload,
1928                    )
1929                    .await
1930                    .is_err()
1931                {
1932                    return;
1933                }
1934                // Advance past filtered records too: they were considered.
1935                last_sent = Some(record.offset);
1936            }
1937
1938            // Flush what arrived during the replay before going live, so the
1939            // handover keeps offset order.
1940            while let Some(envelope) = pending.pop_front() {
1941                if !forward_live_frame(
1942                    &task_context,
1943                    &task_subscription_id,
1944                    &span_view,
1945                    &task_query,
1946                    &envelope,
1947                    &mut last_sent,
1948                )
1949                .await
1950                {
1951                    return;
1952                }
1953            }
1954
1955            if let Some(skipped) = lagged {
1956                report_replay_gap(
1957                    &task_context,
1958                    &task_subscription_id,
1959                    &span_view,
1960                    &task_epoch,
1961                    skipped,
1962                    last_sent,
1963                );
1964                return;
1965            }
1966
1967            loop {
1968                tokio::select! {
1969                    _ = cancel_token.cancelled() => break,
1970                    received = receiver.recv() => {
1971                        let envelope = match received {
1972                            Ok(envelope) => envelope,
1973                            // A lagged tape is a gap. Report it with the last
1974                            // offset delivered *before* the gap and stop
1975                            // delivering on this subscription.
1976                            //
1977                            // Continuing would hand the consumer frames from
1978                            // after the gap, advancing its checkpoint past
1979                            // the skipped records so they could never be
1980                            // replayed. Stopping is not a silent stall: the
1981                            // consumer has an explicit error and a cursor
1982                            // that recovers exactly what it missed.
1983                            //
1984                            // The registration is deliberately left alone.
1985                            // Its lifecycle belongs to the connection loop,
1986                            // which holds the only handle to
1987                            // `active_subscriptions`; releasing half of it
1988                            // here would desynchronise unsubscribe, the
1989                            // duplicate-ID gate and close-time usage.
1990                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
1991                                report_replay_gap(
1992                                    &task_context,
1993                                    &task_subscription_id,
1994                                    &span_view,
1995                                    &task_epoch,
1996                                    skipped,
1997                                    last_sent,
1998                                );
1999                                break;
2000                            }
2001                            Err(broadcast::error::RecvError::Closed) => break,
2002                        };
2003
2004                        if !forward_live_frame(
2005                            &task_context,
2006                            &task_subscription_id,
2007                            &span_view,
2008                            &task_query,
2009                            &envelope,
2010                            &mut last_sent,
2011                        )
2012                        .await
2013                        {
2014                            break;
2015                        }
2016                    }
2017                }
2018            }
2019        }
2020        .instrument(info_span!(
2021            "ws.subscribe.replay",
2022            client_id = %context.client_id,
2023            view = %view_id
2024        )),
2025    );
2026    Ok(())
2027}
2028
2029/// Take whatever the bus already has without waiting, so a long replay cannot
2030/// be lapped by a busy view.
2031fn drain_available(
2032    receiver: &mut broadcast::Receiver<Arc<BusMessage>>,
2033    pending: &mut VecDeque<Arc<BusMessage>>,
2034    lagged: &mut Option<u64>,
2035) {
2036    // Once a gap is known, everything still on the bus is on the far side of
2037    // it. Buffering it would put post-gap frames in front of the lag report,
2038    // advancing `last_sent` past the hole and making `recoverFrom` point
2039    // after the very records it is supposed to recover.
2040    if lagged.is_some() {
2041        return;
2042    }
2043    // Bounded so a view publishing faster than the client drains cannot turn
2044    // the buffer into an unbounded queue. Filling it is not itself a gap:
2045    // nothing has been skipped at that instant, the buffer simply stopped
2046    // accepting. Stop buffering and let the bus report the loss, with the
2047    // count it actually measures, when delivery reaches it.
2048    const MAX_PENDING: usize = 8_192;
2049    loop {
2050        if pending.len() >= MAX_PENDING {
2051            return;
2052        }
2053        match receiver.try_recv() {
2054            Ok(envelope) => pending.push_back(envelope),
2055            Err(broadcast::error::TryRecvError::Empty)
2056            | Err(broadcast::error::TryRecvError::Closed) => return,
2057            Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
2058                *lagged = Some(skipped);
2059                return;
2060            }
2061        }
2062    }
2063}
2064
2065/// Whether a live frame was already delivered by the replay that preceded it.
2066///
2067/// The bus is subscribed before the tape is read, so a record published in
2068/// between appears on both paths; without this the consumer sees it twice.
2069/// Advances the high-water mark as a side effect.
2070fn already_delivered(offset: Option<u64>, last_sent: &mut Option<u64>) -> bool {
2071    let Some(offset) = offset else {
2072        // A frame with no offset predates the tape, so it cannot have been
2073        // replayed and must not move the mark.
2074        return false;
2075    };
2076    if last_sent.is_some_and(|last| offset <= last) {
2077        return true;
2078    }
2079    *last_sent = Some(offset);
2080    false
2081}
2082
2083/// Deliver one live frame, skipping anything the replay already sent.
2084///
2085/// Returns false when the subscription should end.
2086async fn forward_live_frame(
2087    context: &SubscriptionContext,
2088    subscription_id: &str,
2089    view_id: &str,
2090    query: &SubscriptionQuery,
2091    envelope: &Arc<BusMessage>,
2092    last_sent: &mut Option<u64>,
2093) -> bool {
2094    let metadata = source_frame_metadata(&envelope.payload);
2095    if already_delivered(metadata.offset, last_sent) {
2096        return true;
2097    }
2098    if !live_frame_matches(query, &envelope.key, &envelope.payload) {
2099        return true;
2100    }
2101    send_scoped_source_payload(context, subscription_id, view_id, envelope.payload.clone()).is_ok()
2102}
2103
2104fn report_replay_gap(
2105    context: &SubscriptionContext,
2106    subscription_id: &str,
2107    view_id: &str,
2108    epoch: &crate::journal::JournalEpoch,
2109    skipped: u64,
2110    last_sent: Option<u64>,
2111) {
2112    warn!(
2113        "Replay subscription {} lagged past {} records; stopping with a recovery cursor",
2114        subscription_id, skipped
2115    );
2116    let recover_from = last_sent.map(|offset| crate::journal::Cursor {
2117        epoch: epoch.clone(),
2118        offset,
2119    });
2120    let _ = send_control_frame(
2121        context,
2122        &SocketIssueMessage::replay_lagged(
2123            Some(subscription_id.to_string()),
2124            skipped,
2125            recover_from,
2126        ),
2127        view_id,
2128    );
2129}
2130
2131/// Apply the subscription's `key`, `partition` and `filters` to a retained
2132/// record. A replay must honour the same predicates a live subscription does.
2133fn journal_record_matches(
2134    query: &SubscriptionQuery,
2135    record: &crate::journal::JournalRecord,
2136) -> bool {
2137    live_frame_matches(query, &record.key, &record.payload)
2138}
2139
2140fn live_frame_matches(query: &SubscriptionQuery, key: &str, payload: &[u8]) -> bool {
2141    if !query.matches_key(key) {
2142        return false;
2143    }
2144    if query.partition.is_none() && query.filters.is_empty() {
2145        return true;
2146    }
2147    let Ok(frame) = serde_json::from_slice::<Value>(payload) else {
2148        return false;
2149    };
2150    let Some(data) = frame.get("data") else {
2151        return false;
2152    };
2153    if let Some(partition) = &query.partition {
2154        if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2155            return false;
2156        }
2157    }
2158    query
2159        .filters
2160        .iter()
2161        .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2162}
2163
2164/// Awaiting variant of [`send_scoped_source_payload`], for replays that can
2165/// exceed the client's send queue.
2166async fn send_scoped_source_payload_async(
2167    context: &SubscriptionContext,
2168    subscription_id: &str,
2169    view_id: &str,
2170    payload: Arc<Bytes>,
2171) -> Result<()> {
2172    let mut value: Value = serde_json::from_slice(&payload)?;
2173    let object = value
2174        .as_object_mut()
2175        .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2176    object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2177    object.insert(
2178        "subscriptionId".to_string(),
2179        Value::String(subscription_id.to_string()),
2180    );
2181    let json = serde_json::to_vec(&value)?;
2182    let compressed = maybe_compress(&json);
2183    let bytes = compressed.as_bytes().len();
2184    context
2185        .client_manager
2186        .send_compressed_async(context.client_id, compressed)
2187        .await
2188        .map_err(|error| anyhow::anyhow!("failed to send replayed frame: {error}"))?;
2189    context.metrics.message_sent();
2190    emit_update_sent_for_client(
2191        &context.usage_emitter,
2192        &context.client_manager,
2193        context.client_id,
2194        view_id,
2195        bytes,
2196    );
2197    Ok(())
2198}
2199
2200#[derive(Default)]
2201struct SourceFrameMetadata {
2202    op: String,
2203    seq: Option<String>,
2204    offset: Option<u64>,
2205}
2206
2207fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
2208    serde_json::from_slice::<Value>(payload)
2209        .ok()
2210        .map(|value| SourceFrameMetadata {
2211            op: value
2212                .get("op")
2213                .and_then(Value::as_str)
2214                .unwrap_or_default()
2215                .to_string(),
2216            seq: value.get("seq").and_then(Value::as_str).map(str::to_string),
2217            offset: value.get("offset").and_then(Value::as_u64),
2218        })
2219        .unwrap_or_default()
2220}
2221
2222fn send_scoped_source_payload(
2223    context: &SubscriptionContext,
2224    subscription_id: &str,
2225    view_id: &str,
2226    payload: Arc<Bytes>,
2227) -> Result<()> {
2228    let mut value: Value = serde_json::from_slice(&payload)?;
2229    let object = value
2230        .as_object_mut()
2231        .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2232    object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2233    object.insert(
2234        "subscriptionId".to_string(),
2235        Value::String(subscription_id.to_string()),
2236    );
2237    let encoded = Arc::new(Bytes::from(serde_json::to_vec(&value)?));
2238    let bytes = encoded.len();
2239    context
2240        .client_manager
2241        .send_to_client(context.client_id, encoded)
2242        .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
2243    context.metrics.message_sent();
2244    emit_update_sent_for_client(
2245        &context.usage_emitter,
2246        &context.client_manager,
2247        context.client_id,
2248        view_id,
2249        bytes,
2250    );
2251    Ok(())
2252}
2253
2254fn send_membership_frame(
2255    context: &SubscriptionContext,
2256    subscription_id: &str,
2257    view_spec: &ViewSpec,
2258    op: &str,
2259    key: &str,
2260    mut data: Value,
2261    seq: Option<String>,
2262) -> Result<()> {
2263    apply_wire_format(&mut data, &view_spec.wire_format);
2264    let frame = Frame::scoped(
2265        subscription_id,
2266        view_spec.mode,
2267        &view_spec.id,
2268        op,
2269        key,
2270        data,
2271        seq,
2272    );
2273    let encoded = Arc::new(Bytes::from(serde_json::to_vec(&frame)?));
2274    let bytes = encoded.len();
2275    context
2276        .client_manager
2277        .send_to_client(context.client_id, encoded)
2278        .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
2279    context.metrics.message_sent();
2280    emit_update_sent_for_client(
2281        &context.usage_emitter,
2282        &context.client_manager,
2283        context.client_id,
2284        &view_spec.id,
2285        bytes,
2286    );
2287    Ok(())
2288}
2289
2290fn emit_collection_delta(
2291    context: &SubscriptionContext,
2292    subscription_id: &str,
2293    view_spec: &ViewSpec,
2294    current: &[(String, Value)],
2295    next: &[(String, Value)],
2296    envelope: &BusMessage,
2297    metadata: &SourceFrameMetadata,
2298) -> Result<()> {
2299    let current_keys: Vec<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
2300    let next_keys: Vec<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2301    let next_set: HashSet<&str> = next_keys.iter().copied().collect();
2302
2303    for key in current_keys
2304        .iter()
2305        .copied()
2306        .filter(|key| !next_set.contains(key))
2307    {
2308        let op = if metadata.op == "delete" && key == envelope.key {
2309            "delete"
2310        } else {
2311            "remove"
2312        };
2313        send_membership_frame(
2314            context,
2315            subscription_id,
2316            view_spec,
2317            op,
2318            key,
2319            Value::Null,
2320            metadata.seq.clone(),
2321        )?;
2322    }
2323
2324    for (key, data) in next.iter() {
2325        let was_member = current_keys.iter().any(|candidate| *candidate == key);
2326        match member_action(
2327            was_member,
2328            key == &envelope.key,
2329            view_spec.is_derived(),
2330            &metadata.op,
2331        ) {
2332            MemberAction::Skip => {}
2333            MemberAction::ForwardPatch => send_scoped_source_payload(
2334                context,
2335                subscription_id,
2336                &view_spec.id,
2337                envelope.payload.clone(),
2338            )?,
2339            MemberAction::Upsert => {
2340                let seq = metadata
2341                    .seq
2342                    .clone()
2343                    .or_else(|| data.get("_seq").and_then(Value::as_str).map(str::to_string));
2344                send_membership_frame(
2345                    context,
2346                    subscription_id,
2347                    view_spec,
2348                    "upsert",
2349                    key,
2350                    data.clone(),
2351                    seq,
2352                )?;
2353            }
2354        }
2355    }
2356    Ok(())
2357}
2358
2359/// Emit the net effect of every source mutation seen during one coalescing
2360/// interval. Full entities are used for changed members because intermediate
2361/// sparse patches were intentionally discarded and can no longer be merged
2362/// safely by a client.
2363fn emit_coalesced_collection_delta(
2364    context: &SubscriptionContext,
2365    subscription_id: &str,
2366    view_spec: &ViewSpec,
2367    current: &[(String, Value)],
2368    next: &[(String, Value)],
2369    pending: &HashMap<String, Arc<BusMessage>>,
2370) -> Result<()> {
2371    for change in plan_coalesced_collection_delta(current, next, pending) {
2372        send_membership_frame(
2373            context,
2374            subscription_id,
2375            view_spec,
2376            change.op,
2377            &change.key,
2378            change.data,
2379            change.seq,
2380        )?;
2381    }
2382    Ok(())
2383}
2384
2385#[derive(Debug, PartialEq)]
2386struct CollectionChange {
2387    op: &'static str,
2388    key: String,
2389    data: Value,
2390    seq: Option<String>,
2391}
2392
2393fn plan_coalesced_collection_delta(
2394    current: &[(String, Value)],
2395    next: &[(String, Value)],
2396    pending: &HashMap<String, Arc<BusMessage>>,
2397) -> Vec<CollectionChange> {
2398    let current_by_key: HashMap<&str, &Value> = current
2399        .iter()
2400        .map(|(key, data)| (key.as_str(), data))
2401        .collect();
2402    let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2403    let latest_seq = pending
2404        .values()
2405        .filter_map(|envelope| source_frame_metadata(&envelope.payload).seq)
2406        .max_by(|left, right| cmp_seq(left, right));
2407    let mut changes = Vec::new();
2408
2409    for (key, _) in current
2410        .iter()
2411        .filter(|(key, _)| !next_keys.contains(key.as_str()))
2412    {
2413        let metadata = pending
2414            .get(key)
2415            .map(|envelope| source_frame_metadata(&envelope.payload));
2416        let op = if metadata
2417            .as_ref()
2418            .is_some_and(|metadata| metadata.op == "delete")
2419        {
2420            "delete"
2421        } else {
2422            "remove"
2423        };
2424        let seq = metadata
2425            .and_then(|metadata| metadata.seq)
2426            .or_else(|| latest_seq.clone());
2427        changes.push(CollectionChange {
2428            op,
2429            key: key.clone(),
2430            data: Value::Null,
2431            seq,
2432        });
2433    }
2434
2435    for (key, data) in next {
2436        let changed = current_by_key
2437            .get(key.as_str())
2438            .is_none_or(|previous| *previous != data);
2439        if !changed && !pending.contains_key(key) {
2440            continue;
2441        }
2442        let seq = data
2443            .get("_seq")
2444            .and_then(Value::as_str)
2445            .map(str::to_string)
2446            .or_else(|| {
2447                pending
2448                    .get(key)
2449                    .and_then(|envelope| source_frame_metadata(&envelope.payload).seq)
2450            })
2451            .or_else(|| latest_seq.clone());
2452        changes.push(CollectionChange {
2453            op: "upsert",
2454            key: key.clone(),
2455            data: data.clone(),
2456            seq,
2457        });
2458    }
2459    changes
2460}
2461
2462/// What one in-window key owes a subscriber after a source mutation.
2463#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2464enum MemberAction {
2465    /// Send nothing: the subscriber already holds this entity and its data did
2466    /// not change.
2467    Skip,
2468    /// Forward the source patch verbatim.
2469    ForwardPatch,
2470    /// Send the whole entity.
2471    Upsert,
2472}
2473
2474/// Decide what to send for one key in the query window.
2475///
2476/// Only the mutated key's data changed, so a key that was already a member has
2477/// at most moved index — and index is not the server's to communicate.
2478/// `subscribed` announces the window's sort (see [`extract_sort_config`], which
2479/// is `_seq` descending for a plain list) and every SDK re-sorts locally from
2480/// it, so resending an unchanged entity to convey its new position is pure
2481/// waste. This matters because on a `_seq`-ordered list the mutated entity
2482/// jumps to the front on *every* mutation: keying the decision on position
2483/// change meant rebroadcasting the whole window each time, and the mutated
2484/// entity itself never kept its position long enough to have its patch
2485/// forwarded.
2486fn member_action(
2487    was_member: bool,
2488    is_mutated_key: bool,
2489    is_derived: bool,
2490    op: &str,
2491) -> MemberAction {
2492    if !is_mutated_key {
2493        // A key entering the window has no local state to merge into, so it
2494        // needs the whole entity; one already held is unchanged.
2495        return if was_member {
2496            MemberAction::Skip
2497        } else {
2498            MemberAction::Upsert
2499        };
2500    }
2501    // The mutated entity rides its own patch through untouched, but only when
2502    // the subscriber already holds a copy. Derived views still send whole
2503    // entities: the patch on the bus is scoped to the source view, not this
2504    // one (see A4-150).
2505    if was_member && !is_derived && op != "delete" {
2506        MemberAction::ForwardPatch
2507    } else {
2508        MemberAction::Upsert
2509    }
2510}
2511
2512fn to_wire_snapshot_entities(
2513    entities: Vec<(String, Value)>,
2514    view_spec: &ViewSpec,
2515) -> Vec<SnapshotEntity> {
2516    entities
2517        .into_iter()
2518        .map(|(key, mut data)| {
2519            apply_wire_format(&mut data, &view_spec.wire_format);
2520            SnapshotEntity { key, data }
2521        })
2522        .collect()
2523}
2524
2525async fn load_query_entities(
2526    entity_cache: &EntityCache,
2527    sorted_caches: Option<
2528        Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
2529    >,
2530    view_spec: &ViewSpec,
2531    query: &SubscriptionQuery,
2532    apply_snapshot_limit: bool,
2533) -> Vec<(String, Value)> {
2534    let (entities, preordered) = if let Some(sorted_caches) = sorted_caches {
2535        let mut caches = sorted_caches.write().await;
2536        let entities = caches
2537            .get_mut(&view_spec.id)
2538            .map(|cache| cache.get_all_ordered())
2539            .unwrap_or_default();
2540        (entities, true)
2541    } else if view_spec.mode == Mode::State {
2542        let entity = match query.key.as_deref() {
2543            Some(key) => entity_cache
2544                .get(&view_spec.id, key)
2545                .await
2546                .map(|data| vec![(key.to_string(), data)])
2547                .unwrap_or_default(),
2548            None => vec![],
2549        };
2550        (entity, true)
2551    } else {
2552        (entity_cache.get_all(&view_spec.id).await, false)
2553    };
2554    select_query_entities(entities, query, preordered, apply_snapshot_limit)
2555}
2556
2557fn select_query_entities(
2558    mut entities: Vec<(String, Value)>,
2559    query: &SubscriptionQuery,
2560    preordered: bool,
2561    apply_snapshot_limit: bool,
2562) -> Vec<(String, Value)> {
2563    entities.retain(|(key, data)| query_matches_entity(query, key, data));
2564    if !preordered {
2565        entities.sort_by(|left, right| {
2566            let left_seq = left.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2567            let right_seq = right.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2568            let order = if query.after.is_some() {
2569                cmp_seq(left_seq, right_seq)
2570            } else {
2571                cmp_seq(right_seq, left_seq)
2572            };
2573            order.then_with(|| left.0.cmp(&right.0))
2574        });
2575    }
2576
2577    let skip = query.skip.unwrap_or(0);
2578    let take = query.take.unwrap_or(usize::MAX);
2579    let mut selected: Vec<_> = entities.into_iter().skip(skip).take(take).collect();
2580    if apply_snapshot_limit {
2581        if let Some(limit) = query.snapshot_limit {
2582            selected.truncate(limit);
2583        }
2584    }
2585    selected
2586}
2587
2588fn query_matches_entity(query: &SubscriptionQuery, key: &str, data: &Value) -> bool {
2589    if !query.matches_key(key) {
2590        return false;
2591    }
2592    if let Some(partition) = &query.partition {
2593        if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2594            return false;
2595        }
2596    }
2597    if let Some(after) = &query.after {
2598        let Some(seq) = data.get("_seq").and_then(Value::as_str) else {
2599            return false;
2600        };
2601        if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
2602            return false;
2603        }
2604    }
2605    query
2606        .filters
2607        .iter()
2608        .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2609}
2610
2611fn value_at_dot_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
2612    path.split('.')
2613        .try_fold(value, |current, segment| current.get(segment))
2614}
2615
2616#[cfg(test)]
2617mod tests {
2618    use super::*;
2619    use crate::cache::EntityCacheConfig;
2620    use crate::view::{Delivery, Filters, Projection};
2621    use serde_json::json;
2622    use tokio::sync::oneshot;
2623
2624    fn list_spec() -> ViewSpec {
2625        ViewSpec {
2626            id: "Thing/list".to_string(),
2627            export: "Thing".to_string(),
2628            mode: Mode::List,
2629            wire_format: Default::default(),
2630            projection: Projection::all(),
2631            filters: Filters::all(),
2632            delivery: Delivery::default(),
2633            pipeline: None,
2634            source_view: None,
2635        }
2636    }
2637
2638    /// Plan the whole window the way `emit_collection_delta` does, so a test
2639    /// can assert on what a subscriber is actually sent.
2640    fn plan_window(
2641        current_keys: &[&str],
2642        next_keys: &[&str],
2643        envelope_key: &str,
2644        is_derived: bool,
2645        op: &str,
2646    ) -> Vec<(String, MemberAction)> {
2647        next_keys
2648            .iter()
2649            .map(|key| {
2650                let was_member = current_keys.contains(key);
2651                (
2652                    (*key).to_string(),
2653                    member_action(was_member, *key == envelope_key, is_derived, op),
2654                )
2655            })
2656            .collect()
2657    }
2658
2659    #[test]
2660    fn a_reordered_window_sends_only_the_mutated_entity() {
2661        // A `_seq`-descending list: mutating "1" moves it to the front and
2662        // shifts every other key down one. Only "1" changed, so only "1" is
2663        // sent — and it rides its own patch, not a full entity.
2664        let current = ["4", "3", "2", "1"];
2665        let next = ["1", "4", "3", "2"];
2666        let plan = plan_window(&current, &next, "1", false, "patch");
2667
2668        assert_eq!(
2669            plan,
2670            vec![
2671                ("1".to_string(), MemberAction::ForwardPatch),
2672                ("4".to_string(), MemberAction::Skip),
2673                ("3".to_string(), MemberAction::Skip),
2674                ("2".to_string(), MemberAction::Skip),
2675            ]
2676        );
2677    }
2678
2679    #[test]
2680    fn a_key_entering_the_window_gets_the_whole_entity() {
2681        // "5" has no local state for the subscriber to merge a patch into.
2682        let plan = plan_window(&["4", "3"], &["5", "4", "3"], "5", false, "patch");
2683        assert_eq!(
2684            plan,
2685            vec![
2686                ("5".to_string(), MemberAction::Upsert),
2687                ("4".to_string(), MemberAction::Skip),
2688                ("3".to_string(), MemberAction::Skip),
2689            ]
2690        );
2691    }
2692
2693    #[test]
2694    fn derived_views_still_send_whole_entities() {
2695        // The patch on the bus is scoped to the source view, so a derived
2696        // subscription cannot forward it verbatim (A4-150).
2697        let plan = plan_window(&["1", "2"], &["1", "2"], "1", true, "patch");
2698        assert_eq!(
2699            plan,
2700            vec![
2701                ("1".to_string(), MemberAction::Upsert),
2702                ("2".to_string(), MemberAction::Skip),
2703            ]
2704        );
2705    }
2706
2707    #[test]
2708    fn a_delete_envelope_never_forwards_a_patch() {
2709        // A surviving key on a delete envelope carries no mergeable patch.
2710        let plan = plan_window(&["1", "2"], &["1", "2"], "1", false, "delete");
2711        assert_eq!(
2712            plan,
2713            vec![
2714                ("1".to_string(), MemberAction::Upsert),
2715                ("2".to_string(), MemberAction::Skip),
2716            ]
2717        );
2718    }
2719
2720    #[test]
2721    fn an_unchanged_window_sends_one_frame_not_a_broadcast() {
2722        // The regression this guards: 500 members used to mean 500 full
2723        // entities on the wire for a single mutation.
2724        let keys: Vec<String> = (0..500).map(|index| index.to_string()).collect();
2725        let refs: Vec<&str> = keys.iter().map(String::as_str).collect();
2726        let plan = plan_window(&refs, &refs, "250", false, "patch");
2727
2728        let sent = plan
2729            .iter()
2730            .filter(|(_, action)| *action != MemberAction::Skip)
2731            .count();
2732        assert_eq!(sent, 1);
2733        assert_eq!(plan[250].1, MemberAction::ForwardPatch);
2734    }
2735
2736    fn list_message(key: &str, op: &str, seq: &str) -> Arc<BusMessage> {
2737        Arc::new(BusMessage {
2738            key: key.to_string(),
2739            entity: "Thing/list".to_string(),
2740            payload: Arc::new(Bytes::from(
2741                serde_json::to_vec(&json!({
2742                    "entity": "Thing/list",
2743                    "op": op,
2744                    "key": key,
2745                    "seq": seq,
2746                    "data": {},
2747                }))
2748                .unwrap(),
2749            )),
2750        })
2751    }
2752
2753    #[test]
2754    fn coalescing_emits_only_the_final_full_state_per_changed_key() {
2755        let current = vec![
2756            ("a".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2757            ("b".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2758            (
2759                "stable".to_string(),
2760                json!({"count": 1, "_seq": "10:000001"}),
2761            ),
2762        ];
2763        let next = vec![
2764            ("a".to_string(), json!({"count": 3, "_seq": "10:000004"})),
2765            ("c".to_string(), json!({"count": 1, "_seq": "10:000003"})),
2766            (
2767                "stable".to_string(),
2768                json!({"count": 1, "_seq": "10:000001"}),
2769            ),
2770        ];
2771        let pending = HashMap::from([
2772            ("a".to_string(), list_message("a", "patch", "10:000004")),
2773            ("b".to_string(), list_message("b", "delete", "10:000002")),
2774            ("c".to_string(), list_message("c", "patch", "10:000003")),
2775        ]);
2776
2777        assert_eq!(
2778            plan_coalesced_collection_delta(&current, &next, &pending),
2779            vec![
2780                CollectionChange {
2781                    op: "delete",
2782                    key: "b".to_string(),
2783                    data: Value::Null,
2784                    seq: Some("10:000002".to_string()),
2785                },
2786                CollectionChange {
2787                    op: "upsert",
2788                    key: "a".to_string(),
2789                    data: json!({"count": 3, "_seq": "10:000004"}),
2790                    seq: Some("10:000004".to_string()),
2791                },
2792                CollectionChange {
2793                    op: "upsert",
2794                    key: "c".to_string(),
2795                    data: json!({"count": 1, "_seq": "10:000003"}),
2796                    seq: Some("10:000003".to_string()),
2797                },
2798            ]
2799        );
2800    }
2801
2802    #[test]
2803    fn a_delayed_delete_cannot_erase_a_newer_recreated_entity() {
2804        let recreated = json!({"balance": 2, "_seq": "10:000004"});
2805        assert!(source_delete_is_stale(Some(&recreated), Some("10:000002")));
2806        assert!(!source_delete_is_stale(Some(&recreated), Some("10:000004")));
2807        assert!(!source_delete_is_stale(Some(&recreated), None));
2808    }
2809
2810    #[test]
2811    fn a_lagged_snapshotless_subscription_gets_a_fatal_retryable_error() {
2812        let issue = SocketIssueMessage::subscription_lagged("balances".to_string(), 42);
2813        assert_eq!(issue.code, "subscription-lagged");
2814        assert!(issue.retryable);
2815        assert!(issue.fatal);
2816        assert!(issue.message.contains("42"));
2817        assert!(issue
2818            .suggested_action
2819            .unwrap()
2820            .contains("snapshots enabled"));
2821
2822        let append = SocketIssueMessage::append_subscription_lagged("trades".to_string(), 9);
2823        assert!(append.fatal);
2824        assert!(append.suggested_action.unwrap().contains("retained replay"));
2825    }
2826
2827    #[test]
2828    fn snapshot_batches_share_identity_and_completion() {
2829        let entities = ["one", "two", "three"].map(|key| SnapshotEntity {
2830            key: key.to_string(),
2831            data: json!({"key": key}),
2832        });
2833        let batches = create_snapshot_batches(
2834            &entities,
2835            SnapshotMetadata {
2836                subscription_id: "sub-1",
2837                snapshot_id: "snapshot-1",
2838                authoritative: true,
2839                mode: Mode::List,
2840                view_id: "Thing/list",
2841                key: None,
2842            },
2843            &SnapshotBatchConfig {
2844                initial_batch_size: 2,
2845                subsequent_batch_size: 1,
2846            },
2847        );
2848        assert_eq!(batches.len(), 2);
2849        assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
2850        assert!(batches
2851            .iter()
2852            .all(|batch| batch.snapshot_id == "snapshot-1"));
2853        assert!(!batches[0].complete);
2854        assert!(batches[1].complete);
2855        assert!(batches.iter().all(|batch| batch.authoritative));
2856    }
2857
2858    #[test]
2859    fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
2860        let batches = create_snapshot_batches(
2861            &[],
2862            SnapshotMetadata {
2863                subscription_id: "sub-1",
2864                snapshot_id: "snapshot-1",
2865                authoritative: false,
2866                mode: Mode::State,
2867                view_id: "Thing/state",
2868                key: Some("missing"),
2869            },
2870            &SnapshotBatchConfig {
2871                initial_batch_size: 1,
2872                subsequent_batch_size: 1,
2873            },
2874        );
2875        assert_eq!(batches.len(), 1);
2876        assert!(!batches[0].authoritative);
2877        assert!(batches[0].complete);
2878        assert_eq!(batches[0].key.as_deref(), Some("missing"));
2879    }
2880
2881    #[test]
2882    fn lag_recovery_is_authoritative_for_an_after_query() {
2883        let subscription = Subscription {
2884            protocol_version: PROTOCOL_VERSION,
2885            subscription_id: "sub-1".to_string(),
2886            query: SubscriptionQuery {
2887                view: "Thing/list".to_string(),
2888                after: Some("40:000000000010".to_string()),
2889                ..Default::default()
2890            },
2891            snapshot: Default::default(),
2892        };
2893
2894        assert!(!SnapshotPurpose::Initial.authoritative(&subscription));
2895        assert!(SnapshotPurpose::Recovery.authoritative(&subscription));
2896    }
2897
2898    #[test]
2899    fn dot_path_filters_are_exact_and_type_sensitive() {
2900        let mut query = SubscriptionQuery {
2901            view: "Thing/list".to_string(),
2902            ..Default::default()
2903        };
2904        query
2905            .filters
2906            .insert("state.status".to_string(), json!("open"));
2907        query.filters.insert("metrics.count".to_string(), json!(2));
2908        assert!(query_matches_entity(
2909            &query,
2910            "one",
2911            &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
2912        ));
2913        assert!(!query_matches_entity(
2914            &query,
2915            "one",
2916            &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
2917        ));
2918    }
2919
2920    #[test]
2921    fn take_and_skip_define_independent_deterministic_windows() {
2922        let entities: Vec<_> = (1..=6)
2923            .map(|id| {
2924                (
2925                    id.to_string(),
2926                    json!({"id": id, "_seq": format!("10:{id:012}")}),
2927                )
2928            })
2929            .collect();
2930        let first = SubscriptionQuery {
2931            view: "Thing/list".to_string(),
2932            take: Some(2),
2933            skip: Some(0),
2934            ..Default::default()
2935        };
2936        let second = SubscriptionQuery {
2937            skip: Some(2),
2938            ..first.clone()
2939        };
2940        let first_keys: Vec<_> = select_query_entities(entities.clone(), &first, false, false)
2941            .into_iter()
2942            .map(|(key, _)| key)
2943            .collect();
2944        let second_keys: Vec<_> = select_query_entities(entities, &second, false, false)
2945            .into_iter()
2946            .map(|(key, _)| key)
2947            .collect();
2948        assert_eq!(first_keys, ["6", "5"]);
2949        assert_eq!(second_keys, ["4", "3"]);
2950    }
2951
2952    #[tokio::test]
2953    async fn state_receiver_is_installed_before_snapshot_awaits() {
2954        let bus = BusManager::new();
2955        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
2956        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
2957        let bus_for_task = bus.clone();
2958        let task = tokio::spawn(async move {
2959            subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
2960                snapshot_started_tx.send(()).unwrap();
2961                release_snapshot_rx.await.unwrap();
2962            })
2963            .await
2964            .0
2965        });
2966        snapshot_started_rx.await.unwrap();
2967        bus.publish_state(
2968            "Thing/state",
2969            "one",
2970            Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
2971        )
2972        .await;
2973        release_snapshot_tx.send(()).unwrap();
2974        let mut receiver = task.await.unwrap();
2975        receiver.changed().await.unwrap();
2976        assert!(!receiver.borrow().is_empty());
2977    }
2978
2979    async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
2980        let bus = BusManager::new();
2981        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
2982        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
2983        let bus_for_task = bus.clone();
2984        let task = tokio::spawn(async move {
2985            subscribe_list_then_snapshot(&bus_for_task, view, || async move {
2986                snapshot_started_tx.send(()).unwrap();
2987                release_snapshot_rx.await.unwrap();
2988            })
2989            .await
2990            .0
2991        });
2992        snapshot_started_rx.await.unwrap();
2993        bus.publish_list(
2994            view,
2995            Arc::new(BusMessage {
2996                key: "one".to_string(),
2997                entity: view.to_string(),
2998                payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
2999            }),
3000        )
3001        .await;
3002        release_snapshot_tx.send(()).unwrap();
3003        let mut receiver = task.await.unwrap();
3004        assert_eq!(receiver.recv().await.unwrap().key, "one");
3005    }
3006
3007    #[tokio::test]
3008    async fn list_receiver_is_installed_before_snapshot() {
3009        assert_list_receiver_precedes_snapshot("Thing/list").await;
3010    }
3011
3012    #[tokio::test]
3013    async fn append_receiver_is_installed_before_snapshot() {
3014        assert_list_receiver_precedes_snapshot("Thing/append").await;
3015    }
3016
3017    #[tokio::test]
3018    async fn derived_source_receiver_is_installed_before_snapshot() {
3019        assert_list_receiver_precedes_snapshot("Thing/list-source").await;
3020    }
3021
3022    #[tokio::test]
3023    async fn snapshot_limit_does_not_change_live_take_skip_membership() {
3024        let cache = EntityCache::with_config(EntityCacheConfig {
3025            max_entities_per_view: 10,
3026            ..Default::default()
3027        });
3028        for id in 1..=4 {
3029            cache
3030                .upsert(
3031                    "Thing/list",
3032                    &id.to_string(),
3033                    json!({"_seq": format!("10:{id:012}")}),
3034                )
3035                .await;
3036        }
3037        let query = SubscriptionQuery {
3038            view: "Thing/list".to_string(),
3039            take: Some(3),
3040            skip: Some(1),
3041            snapshot_limit: Some(1),
3042            ..Default::default()
3043        };
3044        let live = load_query_entities(&cache, None, &list_spec(), &query, false).await;
3045        let snapshot = load_query_entities(&cache, None, &list_spec(), &query, true).await;
3046        assert_eq!(live.len(), 3);
3047        assert_eq!(snapshot.len(), 1);
3048    }
3049
3050    #[test]
3051    fn fixture_manifest_covers_required_conformance_cases() {
3052        let manifest: Value = serde_json::from_str(include_str!(
3053            "../../../../tests/fixtures/websocket-v2/manifest.json"
3054        ))
3055        .unwrap();
3056        let names: HashSet<_> = manifest["fixtures"]
3057            .as_array()
3058            .unwrap()
3059            .iter()
3060            .filter_map(Value::as_str)
3061            .collect();
3062        for required in [
3063            "keyed-state.json",
3064            "list-windows.json",
3065            "filters.json",
3066            "multi-batch-authoritative.json",
3067            "empty-snapshot.json",
3068            "remove.json",
3069            "delete.json",
3070            "incremental-snapshot.json",
3071            "reconnect-replacement.json",
3072            "errors.json",
3073        ] {
3074            assert!(names.contains(required), "missing fixture {required}");
3075        }
3076
3077        for document in [
3078            include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
3079            include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
3080            include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
3081            include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
3082            include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
3083            include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
3084            include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
3085            include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
3086            include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
3087            include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
3088        ] {
3089            let fixture: Value = serde_json::from_str(document).unwrap();
3090            assert!(fixture["name"].is_string());
3091        }
3092    }
3093
3094    fn append_frame(key: &str, data: Value) -> Arc<Bytes> {
3095        let frame = json!({
3096            "entity": "Trade/append",
3097            "op": "patch",
3098            "key": key,
3099            "offset": 7,
3100            "data": data,
3101        });
3102        Arc::new(Bytes::from(serde_json::to_vec(&frame).unwrap()))
3103    }
3104
3105    /// A replayable subscription must honour the same predicates a live
3106    /// collection subscription does; otherwise a filtered consumer receives
3107    /// events outside the query it asked for.
3108    #[test]
3109    fn replay_delivery_applies_key_partition_and_filters() {
3110        let matching = append_frame("pool1", json!({"_partition": "us", "side": "buy"}));
3111        let other_partition = append_frame("pool1", json!({"_partition": "eu", "side": "buy"}));
3112        let other_side = append_frame("pool1", json!({"_partition": "us", "side": "sell"}));
3113
3114        let unfiltered = SubscriptionQuery {
3115            view: "Trade/append".to_string(),
3116            ..Default::default()
3117        };
3118        assert!(live_frame_matches(&unfiltered, "pool1", &matching));
3119        assert!(live_frame_matches(&unfiltered, "pool9", &matching));
3120
3121        let keyed = SubscriptionQuery {
3122            view: "Trade/append".to_string(),
3123            key: Some("pool1".to_string()),
3124            ..Default::default()
3125        };
3126        assert!(live_frame_matches(&keyed, "pool1", &matching));
3127        assert!(!live_frame_matches(&keyed, "pool2", &matching));
3128
3129        let partitioned = SubscriptionQuery {
3130            view: "Trade/append".to_string(),
3131            partition: Some("us".to_string()),
3132            ..Default::default()
3133        };
3134        assert!(live_frame_matches(&partitioned, "pool1", &matching));
3135        assert!(!live_frame_matches(&partitioned, "pool1", &other_partition));
3136
3137        let filtered = SubscriptionQuery {
3138            view: "Trade/append".to_string(),
3139            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3140            ..Default::default()
3141        };
3142        assert!(live_frame_matches(&filtered, "pool1", &matching));
3143        assert!(!live_frame_matches(&filtered, "pool1", &other_side));
3144    }
3145
3146    #[test]
3147    fn a_frame_without_decodable_data_does_not_satisfy_a_filter() {
3148        let filtered = SubscriptionQuery {
3149            view: "Trade/append".to_string(),
3150            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3151            ..Default::default()
3152        };
3153        let garbage = Arc::new(Bytes::from_static(b"not json"));
3154        assert!(!live_frame_matches(&filtered, "pool1", &garbage));
3155    }
3156
3157    /// The recovery cursor must be the last offset delivered *before* the
3158    /// gap. Reporting the newest offset seen would step the consumer over
3159    /// the skipped records permanently.
3160    #[test]
3161    fn replay_lagged_recovers_from_before_the_gap() {
3162        let epoch = crate::journal::JournalEpoch::new();
3163        let issue = SocketIssueMessage::replay_lagged(
3164            Some("trades".to_string()),
3165            37,
3166            Some(crate::journal::Cursor {
3167                epoch: epoch.clone(),
3168                offset: 4180,
3169            }),
3170        );
3171        assert_eq!(issue.code, "replay-lagged");
3172        assert_eq!(issue.recover_from, Some(format!("{epoch}:4180")));
3173        assert!(
3174            issue.suggested_action.unwrap().contains("4180"),
3175            "the consumer is told exactly which cursor recovers the gap"
3176        );
3177
3178        // Nothing delivered yet: there is no pre-gap offset, so the whole
3179        // retained window is the recovery.
3180        let from_scratch = SocketIssueMessage::replay_lagged(Some("trades".to_string()), 9, None);
3181        assert_eq!(from_scratch.recover_from, None);
3182        assert!(from_scratch
3183            .suggested_action
3184            .unwrap()
3185            .contains("without `after`"));
3186    }
3187
3188    fn bus_message(key: &str) -> Arc<BusMessage> {
3189        Arc::new(BusMessage {
3190            key: key.to_string(),
3191            entity: "Trade/append".to_string(),
3192            payload: Arc::new(Bytes::from_static(b"{}")),
3193        })
3194    }
3195
3196    /// A long replay must keep the bus drained. The broadcast buffer is
3197    /// bounded, so a busy view would otherwise lap the replay and the first
3198    /// live `recv` would return `Lagged` — telling the client to resubscribe,
3199    /// starting another long replay, which laps again.
3200    #[tokio::test]
3201    async fn draining_during_a_replay_keeps_a_busy_view_from_lapping_it() {
3202        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(16);
3203        let mut pending = VecDeque::new();
3204        let mut lagged = None;
3205
3206        // Publish more than the channel holds, draining as a replay would
3207        // between sends.
3208        for index in 0..48 {
3209            sender.send(bus_message(&format!("k{index}"))).unwrap();
3210            drain_available(&mut receiver, &mut pending, &mut lagged);
3211        }
3212
3213        assert_eq!(lagged, None, "draining as we go means nothing is dropped");
3214        assert_eq!(pending.len(), 48, "every published frame is buffered");
3215    }
3216
3217    #[tokio::test]
3218    async fn a_replay_that_never_drains_is_reported_as_a_gap() {
3219        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3220        for index in 0..32 {
3221            sender.send(bus_message(&format!("k{index}"))).unwrap();
3222        }
3223
3224        let mut pending = VecDeque::new();
3225        let mut lagged = None;
3226        drain_available(&mut receiver, &mut pending, &mut lagged);
3227
3228        assert!(
3229            lagged.is_some(),
3230            "overflowing the bus is a gap, not silent truncation"
3231        );
3232    }
3233
3234    /// Everything still on the bus after a gap is on the far side of it.
3235    /// Buffering it would put those frames in front of the lag report and
3236    /// advance the recovery cursor past the records it is meant to recover.
3237    #[tokio::test]
3238    async fn nothing_after_a_gap_is_buffered_ahead_of_the_report() {
3239        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3240        let mut pending = VecDeque::new();
3241        let mut lagged = None;
3242
3243        // Delivered and buffered normally.
3244        sender.send(bus_message("before")).unwrap();
3245        drain_available(&mut receiver, &mut pending, &mut lagged);
3246        assert_eq!(pending.len(), 1);
3247
3248        // Overflow the bus: everything published from here is past the gap.
3249        for index in 0..32 {
3250            sender.send(bus_message(&format!("lost{index}"))).unwrap();
3251        }
3252        drain_available(&mut receiver, &mut pending, &mut lagged);
3253        assert!(lagged.is_some());
3254
3255        // The bus still holds what survived the overflow, all of it past the
3256        // gap. A later iteration of the replay loop drains again.
3257        let buffered_at_gap = pending.len();
3258        drain_available(&mut receiver, &mut pending, &mut lagged);
3259        assert_eq!(
3260            pending.len(),
3261            buffered_at_gap,
3262            "post-gap frames must not join the pre-gap flush"
3263        );
3264        assert_eq!(pending.front().unwrap().key, "before");
3265    }
3266    /// The bus is subscribed before the tape is read, so a record published
3267    /// in that window arrives on both paths.
3268    #[test]
3269    fn the_seam_between_replay_and_live_neither_repeats_nor_skips() {
3270        let mut last_sent = Some(4211);
3271
3272        assert!(
3273            already_delivered(Some(4211), &mut last_sent),
3274            "the record the replay ended on must not be sent twice"
3275        );
3276        assert!(already_delivered(Some(4100), &mut last_sent));
3277        assert_eq!(last_sent, Some(4211), "a duplicate never moves the mark");
3278
3279        assert!(
3280            !already_delivered(Some(4212), &mut last_sent),
3281            "the next record is new"
3282        );
3283        assert_eq!(last_sent, Some(4212));
3284
3285        // A frame with no offset comes from a view with no tape; it cannot
3286        // have been replayed, and must not disturb the mark.
3287        assert!(!already_delivered(None, &mut last_sent));
3288        assert_eq!(last_sent, Some(4212));
3289    }
3290
3291    /// A subscription with no cursor has delivered nothing, so the first live
3292    /// frame is not a duplicate.
3293    #[test]
3294    fn a_fresh_subscription_delivers_its_first_live_frame() {
3295        let mut last_sent = None;
3296        assert!(!already_delivered(Some(0), &mut last_sent));
3297        assert_eq!(last_sent, Some(0));
3298    }
3299
3300    /// End-to-end over a real socket: the pieces above are unit-tested
3301    /// individually, but the thing a consumer actually does — reconnect with
3302    /// a stored cursor and keep reading — only exists once a subscription is
3303    /// attached to a connection.
3304    mod over_a_socket {
3305        use super::*;
3306        use crate::journal::{EventJournal, JournalConfig};
3307        use crate::projector::Projector;
3308        use crate::{MutationBatch, SlotContext};
3309        use arete_interpreter::Mutation;
3310        use futures_util::{SinkExt, StreamExt};
3311        use std::time::Duration;
3312        use tokio::net::{TcpListener, TcpStream};
3313        use tokio::sync::mpsc;
3314        use tokio_tungstenite::tungstenite::Message;
3315        use tokio_tungstenite::{client_async, WebSocketStream};
3316
3317        const RETAINED: u64 = 600;
3318
3319        fn append_index() -> ViewIndex {
3320            let mut index = ViewIndex::new();
3321            index.add_spec(ViewSpec {
3322                id: "Trade/append".to_string(),
3323                export: "Trade".to_string(),
3324                mode: Mode::Append,
3325                wire_format: Default::default(),
3326                projection: Projection::all(),
3327                filters: Filters::all(),
3328                delivery: Delivery::default(),
3329                pipeline: None,
3330                source_view: None,
3331            });
3332            index
3333        }
3334
3335        fn trade(index: u64) -> MutationBatch {
3336            MutationBatch::with_slot_context(
3337                vec![Mutation {
3338                    export: "Trade".to_string(),
3339                    key: json!(format!("pool{}", index % 4)),
3340                    patch: json!({"trade": index}),
3341                    append: vec![],
3342                }]
3343                .into_iter()
3344                .collect(),
3345                SlotContext::new(100 + index / 3, index % 3),
3346            )
3347        }
3348
3349        struct Harness {
3350            addr: SocketAddr,
3351            journal: Arc<EventJournal>,
3352            tx: mpsc::Sender<MutationBatch>,
3353        }
3354
3355        impl Harness {
3356            async fn start() -> Self {
3357                let view_index = Arc::new(append_index());
3358                let entity_cache = EntityCache::new();
3359                let bus_manager = BusManager::new();
3360                let journal = Arc::new(EventJournal::new(JournalConfig {
3361                    enabled: true,
3362                    max_bytes_per_view: u64::MAX,
3363                    max_records_per_view: 10_000,
3364                    max_age: Duration::from_secs(3_600),
3365                }));
3366
3367                let (tx, rx) = mpsc::channel::<MutationBatch>(256);
3368                tokio::spawn(
3369                    Projector::new(
3370                        view_index.clone(),
3371                        bus_manager.clone(),
3372                        entity_cache.clone(),
3373                        rx,
3374                        #[cfg(feature = "otel")]
3375                        None,
3376                    )
3377                    .with_journal(journal.clone())
3378                    .run(),
3379                );
3380
3381                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3382                let addr = listener.local_addr().unwrap();
3383                let server = WebSocketServer::new(
3384                    addr,
3385                    bus_manager,
3386                    entity_cache,
3387                    view_index,
3388                    #[cfg(feature = "otel")]
3389                    None,
3390                )
3391                .with_journal(journal.clone());
3392                let (acceptor, _cleanup) = server.into_acceptor();
3393                tokio::spawn(async move { acceptor.serve_listener(listener).await });
3394
3395                Self { addr, journal, tx }
3396            }
3397
3398            async fn publish(&self, range: std::ops::Range<u64>) {
3399                for index in range {
3400                    self.tx.send(trade(index)).await.unwrap();
3401                }
3402                let (ack, wait) = oneshot::channel();
3403                self.tx
3404                    .send(MutationBatch::flush_marker(ack))
3405                    .await
3406                    .unwrap();
3407                wait.await.unwrap();
3408            }
3409
3410            async fn connect(&self) -> WebSocketStream<TcpStream> {
3411                let stream = TcpStream::connect(self.addr).await.unwrap();
3412                client_async(format!("ws://{}/", self.addr), stream)
3413                    .await
3414                    .unwrap()
3415                    .0
3416            }
3417        }
3418
3419        async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
3420            loop {
3421                let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
3422                    .await
3423                    .expect("the server answers within the timeout")
3424                    .expect("the stream stays open")
3425                    .expect("a readable frame");
3426                // Control and data frames arrive as binary; issue frames as
3427                // text. Both are JSON.
3428                let bytes = match &message {
3429                    Message::Text(text) => text.as_bytes(),
3430                    Message::Binary(bytes) => bytes.as_ref(),
3431                    _ => continue,
3432                };
3433                return serde_json::from_slice(bytes).expect("frames are JSON");
3434            }
3435        }
3436
3437        /// Collect `count` event frames, ignoring anything else on the wire.
3438        async fn collect_trades(socket: &mut WebSocketStream<TcpStream>, count: usize) -> Vec<u64> {
3439            let mut offsets = Vec::with_capacity(count);
3440            while offsets.len() < count {
3441                let frame = next_frame(socket).await;
3442                assert_ne!(
3443                    frame["type"], "error",
3444                    "no error frame should interrupt delivery: {frame}"
3445                );
3446                if let Some(offset) = frame["offset"].as_u64() {
3447                    offsets.push(offset);
3448                }
3449            }
3450            offsets
3451        }
3452
3453        /// The headline claim: reconnecting with a stored cursor delivers every
3454        /// event published since it, in order, and then continues live without
3455        /// a duplicate or a hole at the seam.
3456        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3457        async fn a_reconnect_replays_from_a_cursor_and_continues_live() {
3458            let harness = Harness::start().await;
3459            harness.publish(0..RETAINED).await;
3460
3461            let cursor = harness.journal.window("Trade/append").await;
3462            let stored = format!("{}:{}", cursor.epoch, 99);
3463
3464            let mut socket = harness.connect().await;
3465            socket
3466                .send(Message::Text(
3467                    json!({
3468                        "type": "subscribe",
3469                        "protocolVersion": 2,
3470                        "subscriptionId": "trades",
3471                        "query": {"view": "Trade/append", "after": stored},
3472                    })
3473                    .to_string()
3474                    .into(),
3475                ))
3476                .await
3477                .unwrap();
3478
3479            let ack = next_frame(&mut socket).await;
3480            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3481            assert_eq!(ack["replayWindow"]["next"], json!(RETAINED));
3482
3483            // Well over the 500 a single read or buffer would cover.
3484            let replayed = collect_trades(&mut socket, (RETAINED - 100) as usize).await;
3485            assert_eq!(
3486                replayed,
3487                (100..RETAINED).collect::<Vec<_>>(),
3488                "every event after the cursor, in order, exactly once"
3489            );
3490
3491            // Published only now, so these can only arrive over the live path.
3492            harness.publish(RETAINED..RETAINED + 40).await;
3493            let live = collect_trades(&mut socket, 40).await;
3494            assert_eq!(
3495                live,
3496                (RETAINED..RETAINED + 40).collect::<Vec<_>>(),
3497                "the live stream resumes exactly where the replay stopped"
3498            );
3499
3500            socket.close(None).await.ok();
3501        }
3502
3503        /// Events published *during* the replay must still arrive. The replay
3504        /// and the live subscription are separate reads of the same tape, and
3505        /// the seam between them is where a naive implementation drops or
3506        /// repeats.
3507        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3508        async fn events_published_during_a_replay_are_not_lost() {
3509            let harness = Harness::start().await;
3510            harness.publish(0..RETAINED).await;
3511
3512            let epoch = harness.journal.window("Trade/append").await.epoch;
3513            let mut socket = harness.connect().await;
3514            socket
3515                .send(Message::Text(
3516                    json!({
3517                        "type": "subscribe",
3518                        "protocolVersion": 2,
3519                        "subscriptionId": "trades",
3520                        "query": {"view": "Trade/append", "after": format!("{epoch}:0")},
3521                    })
3522                    .to_string()
3523                    .into(),
3524                ))
3525                .await
3526                .unwrap();
3527            let ack = next_frame(&mut socket).await;
3528            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3529
3530            // Keep publishing while the replay is still draining.
3531            harness.publish(RETAINED..RETAINED + 200).await;
3532
3533            let total = (RETAINED + 200 - 1) as usize;
3534            let delivered = collect_trades(&mut socket, total).await;
3535            assert_eq!(
3536                delivered,
3537                (1..RETAINED + 200).collect::<Vec<_>>(),
3538                "replay and live output join without a gap or a repeat"
3539            );
3540
3541            socket.close(None).await.ok();
3542        }
3543
3544        /// A cursor from another tape lifetime is refused rather than served
3545        /// as a continuation, and the refusal releases the subscription id.
3546        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3547        async fn a_stale_epoch_is_refused_and_frees_the_subscription_id() {
3548            let harness = Harness::start().await;
3549            harness.publish(0..50).await;
3550
3551            let mut socket = harness.connect().await;
3552            socket
3553                .send(Message::Text(
3554                    json!({
3555                        "type": "subscribe",
3556                        "protocolVersion": 2,
3557                        "subscriptionId": "trades",
3558                        "query": {
3559                            "view": "Trade/append",
3560                            "after": format!("{}:10", crate::journal::JournalEpoch::new()),
3561                        },
3562                    })
3563                    .to_string()
3564                    .into(),
3565                ))
3566                .await
3567                .unwrap();
3568
3569            let error = next_frame(&mut socket).await;
3570            assert_eq!(error["type"], "error", "unexpected frame: {error}");
3571            assert_eq!(error["code"], "cursor-epoch-changed");
3572
3573            // The documented recovery is to resubscribe without a cursor. That
3574            // only works if the refused attempt released the id.
3575            socket
3576                .send(Message::Text(
3577                    json!({
3578                        "type": "subscribe",
3579                        "protocolVersion": 2,
3580                        "subscriptionId": "trades",
3581                        "query": {"view": "Trade/append"},
3582                    })
3583                    .to_string()
3584                    .into(),
3585                ))
3586                .await
3587                .unwrap();
3588            let ack = next_frame(&mut socket).await;
3589            assert_eq!(
3590                ack["op"], "subscribed",
3591                "the refused id must be reusable: {ack}"
3592            );
3593
3594            assert_eq!(collect_trades(&mut socket, 50).await.len(), 50);
3595            socket.close(None).await.ok();
3596        }
3597    }
3598
3599    /// Session tokens that expire while their socket is open.
3600    mod session_expiry {
3601        use super::*;
3602        use crate::websocket::auth::SignedSessionAuthPlugin;
3603        use arete_auth::{KeyClass, SessionClaims, SigningKey, TokenSigner, TokenVerifier};
3604        use futures_util::{SinkExt, StreamExt};
3605        use std::time::Duration;
3606        use tokio::net::{TcpListener, TcpStream};
3607        use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
3608        use tokio_tungstenite::tungstenite::protocol::CloseFrame;
3609        use tokio_tungstenite::tungstenite::Message;
3610        use tokio_tungstenite::{client_async, WebSocketStream};
3611
3612        struct Server {
3613            addr: SocketAddr,
3614            signer: TokenSigner,
3615        }
3616
3617        impl Server {
3618            async fn start() -> Self {
3619                let signing_key = SigningKey::generate();
3620                let verifier =
3621                    TokenVerifier::new(signing_key.verifying_key(), "test-issuer", "test-audience");
3622                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3623                let addr = listener.local_addr().unwrap();
3624                let server = WebSocketServer::new(
3625                    addr,
3626                    BusManager::new(),
3627                    EntityCache::new(),
3628                    Arc::new(ViewIndex::new()),
3629                    #[cfg(feature = "otel")]
3630                    None,
3631                )
3632                .with_auth_plugin(Arc::new(SignedSessionAuthPlugin::new(verifier)));
3633                let (acceptor, _cleanup) = server.into_acceptor();
3634                tokio::spawn(async move { acceptor.serve_listener(listener).await });
3635                Self {
3636                    addr,
3637                    signer: TokenSigner::new(signing_key, "test-issuer"),
3638                }
3639            }
3640
3641            fn token(&self, ttl_seconds: u64) -> String {
3642                let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
3643                    .with_scope("read")
3644                    .with_key_class(KeyClass::Secret)
3645                    .with_ttl(ttl_seconds)
3646                    .build();
3647                self.signer.sign(claims).unwrap()
3648            }
3649
3650            async fn connect(&self, token: &str) -> WebSocketStream<TcpStream> {
3651                let stream = TcpStream::connect(self.addr).await.unwrap();
3652                client_async(format!("ws://{}/?hs_token={token}", self.addr), stream)
3653                    .await
3654                    .unwrap()
3655                    .0
3656            }
3657        }
3658
3659        async fn send_json(socket: &mut WebSocketStream<TcpStream>, message: Value) {
3660            socket
3661                .send(Message::Text(message.to_string().into()))
3662                .await
3663                .unwrap();
3664        }
3665
3666        /// The close frame, if the server closes the socket within `wait`.
3667        async fn close_within(
3668            socket: &mut WebSocketStream<TcpStream>,
3669            wait: Duration,
3670        ) -> Option<Option<CloseFrame>> {
3671            tokio::time::timeout(wait, async {
3672                while let Some(Ok(message)) = socket.next().await {
3673                    if let Message::Close(frame) = message {
3674                        return Some(frame);
3675                    }
3676                }
3677                Some(None)
3678            })
3679            .await
3680            .ok()
3681            .flatten()
3682        }
3683
3684        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3685        async fn an_expired_session_is_closed_with_the_reason() {
3686            let server = Server::start().await;
3687            let mut socket = server.connect(&server.token(2)).await;
3688
3689            tokio::time::sleep(Duration::from_secs(3)).await;
3690            send_json(&mut socket, json!({"type": "ping"})).await;
3691
3692            let frame = close_within(&mut socket, Duration::from_secs(5))
3693                .await
3694                .expect("the server closes the expired session")
3695                .expect("the close frame carries a reason");
3696            assert_eq!(frame.code, CloseCode::Policy);
3697            assert_eq!(
3698                frame.reason.as_str(),
3699                "token-expired: Authentication token expired"
3700            );
3701        }
3702
3703        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3704        async fn a_session_refreshed_in_band_outlives_its_first_token() {
3705            let server = Server::start().await;
3706            let mut socket = server.connect(&server.token(2)).await;
3707
3708            send_json(
3709                &mut socket,
3710                json!({"type": "refresh_auth", "token": server.token(3_600)}),
3711            )
3712            .await;
3713            let reply = tokio::time::timeout(Duration::from_secs(5), async {
3714                while let Some(Ok(message)) = socket.next().await {
3715                    if let Message::Text(text) = message {
3716                        return serde_json::from_str::<Value>(text.as_str()).ok();
3717                    }
3718                }
3719                None
3720            })
3721            .await
3722            .expect("the server answers the refresh")
3723            .expect("the answer is JSON");
3724            assert_eq!(reply["success"], true, "refresh accepted: {reply}");
3725
3726            tokio::time::sleep(Duration::from_secs(3)).await;
3727            send_json(&mut socket, json!({"type": "ping"})).await;
3728            assert!(
3729                close_within(&mut socket, Duration::from_millis(1_500))
3730                    .await
3731                    .is_none(),
3732                "the socket stays open on the refreshed token"
3733            );
3734        }
3735    }
3736}