Skip to main content

arete_server/websocket/
server.rs

1use crate::bus::{BusManager, BusMessage, StateUpdate};
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<StateUpdate>, u64, 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    let published = receiver.borrow_and_update().published;
1301    let snapshot = snapshot().await;
1302    (receiver, published, 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, mut seen, 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, published) = {
1418                            let update = receiver.borrow_and_update();
1419                            (update.payload.clone(), update.published)
1420                        };
1421                        // The bus keeps only the latest frame. If more than one
1422                        // was published since the last read, the earlier ones
1423                        // were overwritten and the latest patch alone would
1424                        // drop their fields, so send the cached entity instead.
1425                        let missed = published > seen + 1;
1426                        seen = published;
1427                        let metadata = source_frame_metadata(&payload);
1428                        if metadata.op == "delete" {
1429                            task_context.entity_cache.remove(&query.view, &key).await;
1430                            if member && send_membership_frame(
1431                                &task_context,
1432                                &subscription_id,
1433                                &view_spec_task,
1434                                "delete",
1435                                &key,
1436                                Value::Null,
1437                                metadata.seq,
1438                            ).is_err() {
1439                                break;
1440                            }
1441                            member = false;
1442                            continue;
1443                        }
1444
1445                        let selected = load_query_entities(
1446                            &task_context.entity_cache,
1447                            None,
1448                            &view_spec_task,
1449                            &query,
1450                            false,
1451                        ).await;
1452                        let is_member = !selected.is_empty();
1453                        let result = match (member, selected.into_iter().next()) {
1454                            (true, Some(_)) if !missed => send_scoped_source_payload(
1455                                &task_context,
1456                                &subscription_id,
1457                                &query.view,
1458                                payload,
1459                            ),
1460                            // The whole cached entity, as a patch: it merges into
1461                            // what the client holds, so a field the cache lost
1462                            // to eviction is kept, not replaced. It carries no
1463                            // seq, because it is newer than anything this
1464                            // subscriber was sent and seqs are not ordered within
1465                            // a slot: account updates and instructions number
1466                            // themselves differently, so the latest patch's seq
1467                            // can sort below one already delivered.
1468                            (true, Some((entity_key, data))) => send_membership_frame(
1469                                &task_context,
1470                                &subscription_id,
1471                                &view_spec_task,
1472                                "patch",
1473                                &entity_key,
1474                                data,
1475                                None,
1476                            ),
1477                            (false, Some((entity_key, data))) => send_membership_frame(
1478                                &task_context,
1479                                &subscription_id,
1480                                &view_spec_task,
1481                                "upsert",
1482                                &entity_key,
1483                                data,
1484                                metadata.seq,
1485                            ),
1486                            (true, None) => send_membership_frame(
1487                                &task_context,
1488                                &subscription_id,
1489                                &view_spec_task,
1490                                "remove",
1491                                &key,
1492                                Value::Null,
1493                                metadata.seq,
1494                            ),
1495                            (false, None) => Ok(()),
1496                        };
1497                        if result.is_err() {
1498                            break;
1499                        }
1500                        member = is_member;
1501                    }
1502                }
1503            }
1504        }
1505        .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
1506    );
1507    Ok(())
1508}
1509
1510async fn attach_collection_subscription(
1511    context: &SubscriptionContext,
1512    subscription: Subscription,
1513    view_spec: ViewSpec,
1514    cancel_token: CancellationToken,
1515) -> Result<()> {
1516    let view_id = subscription.query.view.clone();
1517    let source_view_id = view_spec
1518        .source_view
1519        .clone()
1520        .unwrap_or_else(|| view_id.clone());
1521    let (mut receiver, initial_membership) = subscribe_collection_then_snapshot(
1522        context,
1523        &source_view_id,
1524        &view_spec,
1525        &subscription.query,
1526    )
1527    .await;
1528
1529    let mut snapshot_entities = initial_membership.clone();
1530    if let Some(limit) = subscription.query.snapshot_limit {
1531        snapshot_entities.truncate(limit);
1532    }
1533    enforce_snapshot_limit(context, snapshot_entities.len())?;
1534    send_subscribed_frame(context, &subscription, &view_spec)?;
1535    if subscription.snapshot.enabled {
1536        send_snapshot_batches(
1537            context,
1538            &subscription,
1539            &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1540            view_spec.mode,
1541            SnapshotPurpose::Initial,
1542            &context.entity_cache.snapshot_config(),
1543        )
1544        .await?;
1545    }
1546
1547    let task_context = context.clone();
1548    let task_subscription = subscription.clone();
1549    let subscription_id = subscription.subscription_id.clone();
1550    let query = subscription.query.clone();
1551    let view_spec_task = view_spec.clone();
1552    let span_view = view_id.clone();
1553    tokio::spawn(
1554        async move {
1555            let mut current = initial_membership;
1556            // Append views are event tapes: collapsing two records would lose
1557            // observable history. Coalescing is only valid for latest-state
1558            // list membership, where a full final entity preserves meaning.
1559            let coalesce_ms = (view_spec_task.mode == Mode::List)
1560                .then(|| {
1561                    view_spec_task
1562                        .delivery
1563                        .coalesce_ms
1564                        .or(task_context.delivery.collection_coalesce_ms)
1565                        .filter(|milliseconds| *milliseconds > 0)
1566                })
1567                .flatten();
1568            let mut flush_interval = coalesce_ms.map(|milliseconds| {
1569                let period = Duration::from_millis(milliseconds);
1570                let mut interval = tokio::time::interval_at(
1571                    tokio::time::Instant::now() + period,
1572                    period,
1573                );
1574                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
1575                interval
1576            });
1577            let mut pending = HashMap::<String, Arc<BusMessage>>::new();
1578            let mut pending_updates = 0_u64;
1579
1580            loop {
1581                tokio::select! {
1582                    _ = cancel_token.cancelled() => break,
1583                    _ = async {
1584                        flush_interval
1585                            .as_mut()
1586                            .expect("coalescing interval is guarded")
1587                            .tick()
1588                            .await;
1589                    }, if flush_interval.is_some() => {
1590                        if pending.is_empty() {
1591                            continue;
1592                        }
1593                        let sorted_caches = view_spec_task
1594                            .is_derived()
1595                            .then(|| task_context.view_index.sorted_caches());
1596                        let next = load_query_entities(
1597                            &task_context.entity_cache,
1598                            sorted_caches,
1599                            &view_spec_task,
1600                            &query,
1601                            false,
1602                        ).await;
1603                        if emit_coalesced_collection_delta(
1604                            &task_context,
1605                            &subscription_id,
1606                            &view_spec_task,
1607                            &current,
1608                            &next,
1609                            &pending,
1610                        ).is_err() {
1611                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
1612                            break;
1613                        }
1614                        task_context
1615                            .metrics
1616                            .collection_coalesced(&view_id, pending_updates);
1617                        current = next;
1618                        pending.clear();
1619                        pending_updates = 0;
1620                    }
1621                    received = receiver.recv() => {
1622                        let envelope = match received {
1623                            Ok(envelope) => envelope,
1624                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
1625                                task_context.metrics.subscription_lagged(&view_id, skipped);
1626                                warn!(
1627                                    "Subscription {} lagged by {} updates",
1628                                    subscription_id, skipped
1629                                );
1630                                if view_spec_task.mode == Mode::Append {
1631                                    let _ = send_control_frame(
1632                                        &task_context,
1633                                        &SocketIssueMessage::append_subscription_lagged(
1634                                            subscription_id.clone(),
1635                                            skipped,
1636                                        ),
1637                                        &view_id,
1638                                    );
1639                                    task_context.metrics.delivery_stopped(&view_id, "append-lagged-without-replay");
1640                                    break;
1641                                }
1642                                if !task_subscription.snapshot.enabled {
1643                                    let _ = send_control_frame(
1644                                        &task_context,
1645                                        &SocketIssueMessage::subscription_lagged(
1646                                            subscription_id.clone(),
1647                                            skipped,
1648                                        ),
1649                                        &view_id,
1650                                    );
1651                                    task_context.metrics.delivery_stopped(&view_id, "lagged-without-snapshot");
1652                                    break;
1653                                }
1654                                info!(
1655                                    "Subscription {} is recovering from an authoritative snapshot",
1656                                    subscription_id
1657                                );
1658                                match recover_collection_subscription(
1659                                    &task_context,
1660                                    &task_subscription,
1661                                    &view_spec_task,
1662                                    &source_view_id,
1663                                ).await {
1664                                    Ok((next_receiver, recovered)) => {
1665                                        receiver = next_receiver;
1666                                        current = recovered;
1667                                        pending.clear();
1668                                        pending_updates = 0;
1669                                        task_context.metrics.subscription_resnapshot(&view_id);
1670                                        continue;
1671                                    }
1672                                    Err(error) => {
1673                                        warn!(
1674                                            "Subscription {} failed to recover from lag: {error:#}",
1675                                            subscription_id
1676                                        );
1677                                        task_context.metrics.delivery_stopped(&view_id, "resnapshot-failed");
1678                                        break;
1679                                    }
1680                                }
1681                            }
1682                            Err(broadcast::error::RecvError::Closed) => break,
1683                        };
1684
1685                        apply_collection_source_event(
1686                            &task_context,
1687                            &source_view_id,
1688                            &view_spec_task,
1689                            &query,
1690                            &envelope,
1691                        ).await;
1692
1693                        if flush_interval.is_some() {
1694                            pending_updates = pending_updates.saturating_add(1);
1695                            pending.insert(envelope.key.clone(), envelope);
1696                            continue;
1697                        }
1698
1699                        let metadata = source_frame_metadata(&envelope.payload);
1700                        let sorted_caches = view_spec_task
1701                            .is_derived()
1702                            .then(|| task_context.view_index.sorted_caches());
1703                        let next = load_query_entities(
1704                            &task_context.entity_cache,
1705                            sorted_caches,
1706                            &view_spec_task,
1707                            &query,
1708                            false,
1709                        ).await;
1710                        if emit_collection_delta(
1711                            &task_context,
1712                            &subscription_id,
1713                            &view_spec_task,
1714                            &current,
1715                            &next,
1716                            &envelope,
1717                            &metadata,
1718                        ).is_err() {
1719                            task_context.metrics.delivery_stopped(&view_id, "send-failed");
1720                            break;
1721                        }
1722                        current = next;
1723                    }
1724                }
1725            }
1726        }
1727        .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
1728    );
1729    Ok(())
1730}
1731
1732async fn subscribe_collection_then_snapshot(
1733    context: &SubscriptionContext,
1734    source_view_id: &str,
1735    view_spec: &ViewSpec,
1736    query: &SubscriptionQuery,
1737) -> (broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>) {
1738    let cache = context.entity_cache.clone();
1739    let sorted_caches = view_spec
1740        .is_derived()
1741        .then(|| context.view_index.sorted_caches());
1742    let view_spec = view_spec.clone();
1743    let query = query.clone();
1744    subscribe_list_then_snapshot(&context.bus_manager, source_view_id, move || async move {
1745        load_query_entities(&cache, sorted_caches, &view_spec, &query, false).await
1746    })
1747    .await
1748}
1749
1750async fn recover_collection_subscription(
1751    context: &SubscriptionContext,
1752    subscription: &Subscription,
1753    view_spec: &ViewSpec,
1754    source_view_id: &str,
1755) -> Result<(broadcast::Receiver<Arc<BusMessage>>, Vec<(String, Value)>)> {
1756    let (receiver, membership) =
1757        subscribe_collection_then_snapshot(context, source_view_id, view_spec, &subscription.query)
1758            .await;
1759    // `snapshotLimit` caps only the initial transfer. Recovery has to replace
1760    // the complete live membership; truncating it would make the replacement
1761    // authoritative while immediately omitting members the server still
1762    // considers current.
1763    enforce_snapshot_limit(context, membership.len())?;
1764    send_snapshot_batches(
1765        context,
1766        subscription,
1767        &to_wire_snapshot_entities(membership.clone(), view_spec),
1768        view_spec.mode,
1769        SnapshotPurpose::Recovery,
1770        &context.entity_cache.snapshot_config(),
1771    )
1772    .await?;
1773    Ok((receiver, membership))
1774}
1775
1776async fn apply_collection_source_event(
1777    context: &SubscriptionContext,
1778    source_view_id: &str,
1779    view_spec: &ViewSpec,
1780    query: &SubscriptionQuery,
1781    envelope: &BusMessage,
1782) {
1783    let metadata = source_frame_metadata(&envelope.payload);
1784    if metadata.op != "delete" {
1785        return;
1786    }
1787    // A slow subscription can observe an old delete after the projector has
1788    // already recreated the key. Never let subscriber-local lag erase newer
1789    // shared cache state.
1790    let current = context
1791        .entity_cache
1792        .get(source_view_id, &envelope.key)
1793        .await;
1794    if source_delete_is_stale(current.as_ref(), metadata.seq.as_deref()) {
1795        return;
1796    }
1797    context
1798        .entity_cache
1799        .remove(source_view_id, &envelope.key)
1800        .await;
1801    if view_spec.is_derived() {
1802        let caches = context.view_index.sorted_caches();
1803        let mut guard = caches.write().await;
1804        if let Some(cache) = guard.get_mut(&query.view) {
1805            cache.remove(&envelope.key);
1806        }
1807    }
1808}
1809
1810fn source_delete_is_stale(current: Option<&Value>, delete_seq: Option<&str>) -> bool {
1811    match (current, delete_seq) {
1812        (Some(current), Some(delete_seq)) => current
1813            .get("_seq")
1814            .and_then(Value::as_str)
1815            .is_some_and(|current_seq| {
1816                cmp_seq(current_seq, delete_seq) == std::cmp::Ordering::Greater
1817            }),
1818        _ => false,
1819    }
1820}
1821
1822/// A subscription refused for a reason the client needs spelled out.
1823///
1824/// Attach paths that return this get the registration released by the
1825/// connection loop, the same as any other failure, while the client still
1826/// receives the specific error rather than a generic `subscription-rejected`.
1827/// Sending the frame and returning `Ok` instead would leave a registered
1828/// subscription with nothing attached: it would hold a slot against the
1829/// client's limit and make the advertised "resubscribe" remediation fail with
1830/// `duplicate-subscription-id`.
1831#[derive(Debug)]
1832pub(crate) struct RejectedSubscription(pub SocketIssueMessage);
1833
1834impl RejectedSubscription {
1835    fn into_error(issue: SocketIssueMessage) -> anyhow::Error {
1836        anyhow::Error::new(Self(issue))
1837    }
1838}
1839
1840impl std::fmt::Display for RejectedSubscription {
1841    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1842        formatter.write_str(&self.0.message)
1843    }
1844}
1845
1846impl std::error::Error for RejectedSubscription {}
1847
1848/// Deliver an append view as an event tape: replay the retained records after
1849/// the cursor, then forward live frames.
1850///
1851/// This deliberately does not recompute membership from the entity cache the
1852/// way [`attach_collection_subscription`] does. The cache folds each patch
1853/// into the resident entity, so a membership diff cannot express "these three
1854/// events happened"; the retained records can.
1855async fn attach_journal_subscription(
1856    context: &SubscriptionContext,
1857    subscription: Subscription,
1858    view_spec: ViewSpec,
1859    journal: Arc<crate::journal::EventJournal>,
1860    cancel_token: CancellationToken,
1861) -> Result<()> {
1862    let view_id = view_spec.id.clone();
1863    let subscription_id = subscription.subscription_id.clone();
1864
1865    // A tape has no membership window, so `take`/`skip` cannot mean what they
1866    // mean on a list view. Refuse them rather than accept and ignore them.
1867    if subscription.query.take.is_some()
1868        || subscription.query.skip.is_some()
1869        || subscription.query.snapshot_limit.is_some()
1870    {
1871        return Err(RejectedSubscription::into_error(SocketIssueMessage::protocol(
1872            Some(subscription_id),
1873            "invalid-subscription",
1874            format!(
1875                "take, skip and snapshotLimit are window options and do not apply to the replayable view {view_id}"
1876            ),
1877        )));
1878    }
1879
1880    // Reject a malformed cursor rather than silently replaying from the start,
1881    // which would look like success and duplicate everything.
1882    let cursor = match subscription.query.after.as_deref() {
1883        Some(raw) => match crate::journal::Cursor::parse(raw) {
1884            Some(cursor) => Some(cursor),
1885            None => {
1886                return Err(RejectedSubscription::into_error(
1887                    SocketIssueMessage::protocol(
1888                        Some(subscription_id),
1889                        "invalid-cursor",
1890                        format!(
1891                            "`after` must be an {{epoch}}:{{offset}} replay cursor for view {view_id}, got {raw:?}"
1892                        ),
1893                    ),
1894                ));
1895            }
1896        },
1897        None => None,
1898    };
1899
1900    // Subscribe before reading the journal so anything published during the
1901    // replay is still delivered; the offset filter below drops the overlap.
1902    let mut receiver = context.bus_manager.get_or_create_list_bus(&view_id).await;
1903
1904    let replayed = match journal.replay_after(&view_id, cursor.as_ref()).await {
1905        Ok(records) => records,
1906        Err(error) => {
1907            return Err(RejectedSubscription::into_error(
1908                SocketIssueMessage::replay_refused(Some(subscription_id), &error),
1909            ));
1910        }
1911    };
1912
1913    let frame = SubscribedFrame::new(
1914        subscription.subscription_id.clone(),
1915        subscription.query.clone(),
1916        view_spec.mode,
1917        extract_sort_config(&view_spec),
1918    )
1919    .with_replay_window(journal.window(&view_id).await);
1920    send_control_frame(context, &frame, &view_id)?;
1921
1922    // Everything past the acknowledgement runs on its own task. The replay can
1923    // be long and applies real backpressure, and `attach_client_to_bus` is
1924    // awaited directly on the connection's inbound loop — doing it there would
1925    // block unsubscribe, auth refresh and pong for the whole replay.
1926    let task_context = context.clone();
1927    let task_subscription_id = subscription.subscription_id.clone();
1928    let task_query = subscription.query.clone();
1929    let task_epoch = journal.epoch().await;
1930    let span_view = view_id.clone();
1931    tokio::spawn(
1932        async move {
1933            let mut last_sent = cursor.map(|cursor| cursor.offset);
1934            // Frames that published while the replay was still running. The
1935            // bus is a bounded broadcast, so it has to be drained as we go or
1936            // a busy view laps us before the replay finishes.
1937            let mut pending: VecDeque<Arc<BusMessage>> = VecDeque::new();
1938            let mut lagged: Option<u64> = None;
1939
1940            for record in replayed {
1941                if cancel_token.is_cancelled() {
1942                    return;
1943                }
1944                drain_available(&mut receiver, &mut pending, &mut lagged);
1945                if journal_record_matches(&task_query, &record)
1946                    && send_scoped_source_payload_async(
1947                        &task_context,
1948                        &task_subscription_id,
1949                        &span_view,
1950                        record.payload,
1951                    )
1952                    .await
1953                    .is_err()
1954                {
1955                    return;
1956                }
1957                // Advance past filtered records too: they were considered.
1958                last_sent = Some(record.offset);
1959            }
1960
1961            // Flush what arrived during the replay before going live, so the
1962            // handover keeps offset order.
1963            while let Some(envelope) = pending.pop_front() {
1964                if !forward_live_frame(
1965                    &task_context,
1966                    &task_subscription_id,
1967                    &span_view,
1968                    &task_query,
1969                    &envelope,
1970                    &mut last_sent,
1971                )
1972                .await
1973                {
1974                    return;
1975                }
1976            }
1977
1978            if let Some(skipped) = lagged {
1979                report_replay_gap(
1980                    &task_context,
1981                    &task_subscription_id,
1982                    &span_view,
1983                    &task_epoch,
1984                    skipped,
1985                    last_sent,
1986                );
1987                return;
1988            }
1989
1990            loop {
1991                tokio::select! {
1992                    _ = cancel_token.cancelled() => break,
1993                    received = receiver.recv() => {
1994                        let envelope = match received {
1995                            Ok(envelope) => envelope,
1996                            // A lagged tape is a gap. Report it with the last
1997                            // offset delivered *before* the gap and stop
1998                            // delivering on this subscription.
1999                            //
2000                            // Continuing would hand the consumer frames from
2001                            // after the gap, advancing its checkpoint past
2002                            // the skipped records so they could never be
2003                            // replayed. Stopping is not a silent stall: the
2004                            // consumer has an explicit error and a cursor
2005                            // that recovers exactly what it missed.
2006                            //
2007                            // The registration is deliberately left alone.
2008                            // Its lifecycle belongs to the connection loop,
2009                            // which holds the only handle to
2010                            // `active_subscriptions`; releasing half of it
2011                            // here would desynchronise unsubscribe, the
2012                            // duplicate-ID gate and close-time usage.
2013                            Err(broadcast::error::RecvError::Lagged(skipped)) => {
2014                                report_replay_gap(
2015                                    &task_context,
2016                                    &task_subscription_id,
2017                                    &span_view,
2018                                    &task_epoch,
2019                                    skipped,
2020                                    last_sent,
2021                                );
2022                                break;
2023                            }
2024                            Err(broadcast::error::RecvError::Closed) => break,
2025                        };
2026
2027                        if !forward_live_frame(
2028                            &task_context,
2029                            &task_subscription_id,
2030                            &span_view,
2031                            &task_query,
2032                            &envelope,
2033                            &mut last_sent,
2034                        )
2035                        .await
2036                        {
2037                            break;
2038                        }
2039                    }
2040                }
2041            }
2042        }
2043        .instrument(info_span!(
2044            "ws.subscribe.replay",
2045            client_id = %context.client_id,
2046            view = %view_id
2047        )),
2048    );
2049    Ok(())
2050}
2051
2052/// Take whatever the bus already has without waiting, so a long replay cannot
2053/// be lapped by a busy view.
2054fn drain_available(
2055    receiver: &mut broadcast::Receiver<Arc<BusMessage>>,
2056    pending: &mut VecDeque<Arc<BusMessage>>,
2057    lagged: &mut Option<u64>,
2058) {
2059    // Once a gap is known, everything still on the bus is on the far side of
2060    // it. Buffering it would put post-gap frames in front of the lag report,
2061    // advancing `last_sent` past the hole and making `recoverFrom` point
2062    // after the very records it is supposed to recover.
2063    if lagged.is_some() {
2064        return;
2065    }
2066    // Bounded so a view publishing faster than the client drains cannot turn
2067    // the buffer into an unbounded queue. Filling it is not itself a gap:
2068    // nothing has been skipped at that instant, the buffer simply stopped
2069    // accepting. Stop buffering and let the bus report the loss, with the
2070    // count it actually measures, when delivery reaches it.
2071    const MAX_PENDING: usize = 8_192;
2072    loop {
2073        if pending.len() >= MAX_PENDING {
2074            return;
2075        }
2076        match receiver.try_recv() {
2077            Ok(envelope) => pending.push_back(envelope),
2078            Err(broadcast::error::TryRecvError::Empty)
2079            | Err(broadcast::error::TryRecvError::Closed) => return,
2080            Err(broadcast::error::TryRecvError::Lagged(skipped)) => {
2081                *lagged = Some(skipped);
2082                return;
2083            }
2084        }
2085    }
2086}
2087
2088/// Whether a live frame was already delivered by the replay that preceded it.
2089///
2090/// The bus is subscribed before the tape is read, so a record published in
2091/// between appears on both paths; without this the consumer sees it twice.
2092/// Advances the high-water mark as a side effect.
2093fn already_delivered(offset: Option<u64>, last_sent: &mut Option<u64>) -> bool {
2094    let Some(offset) = offset else {
2095        // A frame with no offset predates the tape, so it cannot have been
2096        // replayed and must not move the mark.
2097        return false;
2098    };
2099    if last_sent.is_some_and(|last| offset <= last) {
2100        return true;
2101    }
2102    *last_sent = Some(offset);
2103    false
2104}
2105
2106/// Deliver one live frame, skipping anything the replay already sent.
2107///
2108/// Returns false when the subscription should end.
2109async fn forward_live_frame(
2110    context: &SubscriptionContext,
2111    subscription_id: &str,
2112    view_id: &str,
2113    query: &SubscriptionQuery,
2114    envelope: &Arc<BusMessage>,
2115    last_sent: &mut Option<u64>,
2116) -> bool {
2117    let metadata = source_frame_metadata(&envelope.payload);
2118    if already_delivered(metadata.offset, last_sent) {
2119        return true;
2120    }
2121    if !live_frame_matches(query, &envelope.key, &envelope.payload) {
2122        return true;
2123    }
2124    send_scoped_source_payload(context, subscription_id, view_id, envelope.payload.clone()).is_ok()
2125}
2126
2127fn report_replay_gap(
2128    context: &SubscriptionContext,
2129    subscription_id: &str,
2130    view_id: &str,
2131    epoch: &crate::journal::JournalEpoch,
2132    skipped: u64,
2133    last_sent: Option<u64>,
2134) {
2135    warn!(
2136        "Replay subscription {} lagged past {} records; stopping with a recovery cursor",
2137        subscription_id, skipped
2138    );
2139    let recover_from = last_sent.map(|offset| crate::journal::Cursor {
2140        epoch: epoch.clone(),
2141        offset,
2142    });
2143    let _ = send_control_frame(
2144        context,
2145        &SocketIssueMessage::replay_lagged(
2146            Some(subscription_id.to_string()),
2147            skipped,
2148            recover_from,
2149        ),
2150        view_id,
2151    );
2152}
2153
2154/// Apply the subscription's `key`, `partition` and `filters` to a retained
2155/// record. A replay must honour the same predicates a live subscription does.
2156fn journal_record_matches(
2157    query: &SubscriptionQuery,
2158    record: &crate::journal::JournalRecord,
2159) -> bool {
2160    live_frame_matches(query, &record.key, &record.payload)
2161}
2162
2163fn live_frame_matches(query: &SubscriptionQuery, key: &str, payload: &[u8]) -> bool {
2164    if !query.matches_key(key) {
2165        return false;
2166    }
2167    if query.partition.is_none() && query.filters.is_empty() {
2168        return true;
2169    }
2170    let Ok(frame) = serde_json::from_slice::<Value>(payload) else {
2171        return false;
2172    };
2173    let Some(data) = frame.get("data") else {
2174        return false;
2175    };
2176    if let Some(partition) = &query.partition {
2177        if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2178            return false;
2179        }
2180    }
2181    query
2182        .filters
2183        .iter()
2184        .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2185}
2186
2187/// Awaiting variant of [`send_scoped_source_payload`], for replays that can
2188/// exceed the client's send queue.
2189async fn send_scoped_source_payload_async(
2190    context: &SubscriptionContext,
2191    subscription_id: &str,
2192    view_id: &str,
2193    payload: Arc<Bytes>,
2194) -> Result<()> {
2195    let mut value: Value = serde_json::from_slice(&payload)?;
2196    let object = value
2197        .as_object_mut()
2198        .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2199    object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2200    object.insert(
2201        "subscriptionId".to_string(),
2202        Value::String(subscription_id.to_string()),
2203    );
2204    let json = serde_json::to_vec(&value)?;
2205    let compressed = maybe_compress(&json);
2206    let bytes = compressed.as_bytes().len();
2207    context
2208        .client_manager
2209        .send_compressed_async(context.client_id, compressed)
2210        .await
2211        .map_err(|error| anyhow::anyhow!("failed to send replayed frame: {error}"))?;
2212    context.metrics.message_sent();
2213    emit_update_sent_for_client(
2214        &context.usage_emitter,
2215        &context.client_manager,
2216        context.client_id,
2217        view_id,
2218        bytes,
2219    );
2220    Ok(())
2221}
2222
2223#[derive(Default)]
2224struct SourceFrameMetadata {
2225    op: String,
2226    seq: Option<String>,
2227    offset: Option<u64>,
2228}
2229
2230fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
2231    serde_json::from_slice::<Value>(payload)
2232        .ok()
2233        .map(|value| SourceFrameMetadata {
2234            op: value
2235                .get("op")
2236                .and_then(Value::as_str)
2237                .unwrap_or_default()
2238                .to_string(),
2239            seq: value.get("seq").and_then(Value::as_str).map(str::to_string),
2240            offset: value.get("offset").and_then(Value::as_u64),
2241        })
2242        .unwrap_or_default()
2243}
2244
2245fn send_scoped_source_payload(
2246    context: &SubscriptionContext,
2247    subscription_id: &str,
2248    view_id: &str,
2249    payload: Arc<Bytes>,
2250) -> Result<()> {
2251    let mut value: Value = serde_json::from_slice(&payload)?;
2252    let object = value
2253        .as_object_mut()
2254        .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
2255    object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
2256    object.insert(
2257        "subscriptionId".to_string(),
2258        Value::String(subscription_id.to_string()),
2259    );
2260    let encoded = Arc::new(Bytes::from(serde_json::to_vec(&value)?));
2261    let bytes = encoded.len();
2262    context
2263        .client_manager
2264        .send_to_client(context.client_id, encoded)
2265        .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
2266    context.metrics.message_sent();
2267    emit_update_sent_for_client(
2268        &context.usage_emitter,
2269        &context.client_manager,
2270        context.client_id,
2271        view_id,
2272        bytes,
2273    );
2274    Ok(())
2275}
2276
2277fn send_membership_frame(
2278    context: &SubscriptionContext,
2279    subscription_id: &str,
2280    view_spec: &ViewSpec,
2281    op: &str,
2282    key: &str,
2283    mut data: Value,
2284    seq: Option<String>,
2285) -> Result<()> {
2286    apply_wire_format(&mut data, &view_spec.wire_format);
2287    let frame = Frame::scoped(
2288        subscription_id,
2289        view_spec.mode,
2290        &view_spec.id,
2291        op,
2292        key,
2293        data,
2294        seq,
2295    );
2296    let encoded = Arc::new(Bytes::from(serde_json::to_vec(&frame)?));
2297    let bytes = encoded.len();
2298    context
2299        .client_manager
2300        .send_to_client(context.client_id, encoded)
2301        .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
2302    context.metrics.message_sent();
2303    emit_update_sent_for_client(
2304        &context.usage_emitter,
2305        &context.client_manager,
2306        context.client_id,
2307        &view_spec.id,
2308        bytes,
2309    );
2310    Ok(())
2311}
2312
2313fn emit_collection_delta(
2314    context: &SubscriptionContext,
2315    subscription_id: &str,
2316    view_spec: &ViewSpec,
2317    current: &[(String, Value)],
2318    next: &[(String, Value)],
2319    envelope: &BusMessage,
2320    metadata: &SourceFrameMetadata,
2321) -> Result<()> {
2322    let current_keys: Vec<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
2323    let next_keys: Vec<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2324    let next_set: HashSet<&str> = next_keys.iter().copied().collect();
2325
2326    for key in current_keys
2327        .iter()
2328        .copied()
2329        .filter(|key| !next_set.contains(key))
2330    {
2331        let op = if metadata.op == "delete" && key == envelope.key {
2332            "delete"
2333        } else {
2334            "remove"
2335        };
2336        send_membership_frame(
2337            context,
2338            subscription_id,
2339            view_spec,
2340            op,
2341            key,
2342            Value::Null,
2343            metadata.seq.clone(),
2344        )?;
2345    }
2346
2347    for (key, data) in next.iter() {
2348        let was_member = current_keys.iter().any(|candidate| *candidate == key);
2349        match member_action(
2350            was_member,
2351            key == &envelope.key,
2352            view_spec.is_derived(),
2353            &metadata.op,
2354        ) {
2355            MemberAction::Skip => {}
2356            MemberAction::ForwardPatch => send_scoped_source_payload(
2357                context,
2358                subscription_id,
2359                &view_spec.id,
2360                envelope.payload.clone(),
2361            )?,
2362            MemberAction::Upsert => {
2363                let seq = metadata
2364                    .seq
2365                    .clone()
2366                    .or_else(|| data.get("_seq").and_then(Value::as_str).map(str::to_string));
2367                send_membership_frame(
2368                    context,
2369                    subscription_id,
2370                    view_spec,
2371                    "upsert",
2372                    key,
2373                    data.clone(),
2374                    seq,
2375                )?;
2376            }
2377        }
2378    }
2379    Ok(())
2380}
2381
2382/// Emit the net effect of every source mutation seen during one coalescing
2383/// interval. Full entities are used for changed members because intermediate
2384/// sparse patches were intentionally discarded and can no longer be merged
2385/// safely by a client.
2386fn emit_coalesced_collection_delta(
2387    context: &SubscriptionContext,
2388    subscription_id: &str,
2389    view_spec: &ViewSpec,
2390    current: &[(String, Value)],
2391    next: &[(String, Value)],
2392    pending: &HashMap<String, Arc<BusMessage>>,
2393) -> Result<()> {
2394    for change in plan_coalesced_collection_delta(current, next, pending) {
2395        send_membership_frame(
2396            context,
2397            subscription_id,
2398            view_spec,
2399            change.op,
2400            &change.key,
2401            change.data,
2402            change.seq,
2403        )?;
2404    }
2405    Ok(())
2406}
2407
2408#[derive(Debug, PartialEq)]
2409struct CollectionChange {
2410    op: &'static str,
2411    key: String,
2412    data: Value,
2413    seq: Option<String>,
2414}
2415
2416fn plan_coalesced_collection_delta(
2417    current: &[(String, Value)],
2418    next: &[(String, Value)],
2419    pending: &HashMap<String, Arc<BusMessage>>,
2420) -> Vec<CollectionChange> {
2421    let current_by_key: HashMap<&str, &Value> = current
2422        .iter()
2423        .map(|(key, data)| (key.as_str(), data))
2424        .collect();
2425    let next_keys: HashSet<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
2426    let latest_seq = pending
2427        .values()
2428        .filter_map(|envelope| source_frame_metadata(&envelope.payload).seq)
2429        .max_by(|left, right| cmp_seq(left, right));
2430    let mut changes = Vec::new();
2431
2432    for (key, _) in current
2433        .iter()
2434        .filter(|(key, _)| !next_keys.contains(key.as_str()))
2435    {
2436        let metadata = pending
2437            .get(key)
2438            .map(|envelope| source_frame_metadata(&envelope.payload));
2439        let op = if metadata
2440            .as_ref()
2441            .is_some_and(|metadata| metadata.op == "delete")
2442        {
2443            "delete"
2444        } else {
2445            "remove"
2446        };
2447        let seq = metadata
2448            .and_then(|metadata| metadata.seq)
2449            .or_else(|| latest_seq.clone());
2450        changes.push(CollectionChange {
2451            op,
2452            key: key.clone(),
2453            data: Value::Null,
2454            seq,
2455        });
2456    }
2457
2458    for (key, data) in next {
2459        let changed = current_by_key
2460            .get(key.as_str())
2461            .is_none_or(|previous| *previous != data);
2462        if !changed && !pending.contains_key(key) {
2463            continue;
2464        }
2465        let seq = data
2466            .get("_seq")
2467            .and_then(Value::as_str)
2468            .map(str::to_string)
2469            .or_else(|| {
2470                pending
2471                    .get(key)
2472                    .and_then(|envelope| source_frame_metadata(&envelope.payload).seq)
2473            })
2474            .or_else(|| latest_seq.clone());
2475        changes.push(CollectionChange {
2476            op: "upsert",
2477            key: key.clone(),
2478            data: data.clone(),
2479            seq,
2480        });
2481    }
2482    changes
2483}
2484
2485/// What one in-window key owes a subscriber after a source mutation.
2486#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2487enum MemberAction {
2488    /// Send nothing: the subscriber already holds this entity and its data did
2489    /// not change.
2490    Skip,
2491    /// Forward the source patch verbatim.
2492    ForwardPatch,
2493    /// Send the whole entity.
2494    Upsert,
2495}
2496
2497/// Decide what to send for one key in the query window.
2498///
2499/// Only the mutated key's data changed, so a key that was already a member has
2500/// at most moved index — and index is not the server's to communicate.
2501/// `subscribed` announces the window's sort (see [`extract_sort_config`], which
2502/// is `_seq` descending for a plain list) and every SDK re-sorts locally from
2503/// it, so resending an unchanged entity to convey its new position is pure
2504/// waste. This matters because on a `_seq`-ordered list the mutated entity
2505/// jumps to the front on *every* mutation: keying the decision on position
2506/// change meant rebroadcasting the whole window each time, and the mutated
2507/// entity itself never kept its position long enough to have its patch
2508/// forwarded.
2509fn member_action(
2510    was_member: bool,
2511    is_mutated_key: bool,
2512    is_derived: bool,
2513    op: &str,
2514) -> MemberAction {
2515    if !is_mutated_key {
2516        // A key entering the window has no local state to merge into, so it
2517        // needs the whole entity; one already held is unchanged.
2518        return if was_member {
2519            MemberAction::Skip
2520        } else {
2521            MemberAction::Upsert
2522        };
2523    }
2524    // The mutated entity rides its own patch through untouched, but only when
2525    // the subscriber already holds a copy. Derived views still send whole
2526    // entities: the patch on the bus is scoped to the source view, not this
2527    // one (see A4-150).
2528    if was_member && !is_derived && op != "delete" {
2529        MemberAction::ForwardPatch
2530    } else {
2531        MemberAction::Upsert
2532    }
2533}
2534
2535fn to_wire_snapshot_entities(
2536    entities: Vec<(String, Value)>,
2537    view_spec: &ViewSpec,
2538) -> Vec<SnapshotEntity> {
2539    entities
2540        .into_iter()
2541        .map(|(key, mut data)| {
2542            apply_wire_format(&mut data, &view_spec.wire_format);
2543            SnapshotEntity { key, data }
2544        })
2545        .collect()
2546}
2547
2548async fn load_query_entities(
2549    entity_cache: &EntityCache,
2550    sorted_caches: Option<
2551        Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
2552    >,
2553    view_spec: &ViewSpec,
2554    query: &SubscriptionQuery,
2555    apply_snapshot_limit: bool,
2556) -> Vec<(String, Value)> {
2557    let (entities, preordered) = if let Some(sorted_caches) = sorted_caches {
2558        let mut caches = sorted_caches.write().await;
2559        let entities = caches
2560            .get_mut(&view_spec.id)
2561            .map(|cache| cache.get_all_ordered())
2562            .unwrap_or_default();
2563        (entities, true)
2564    } else if view_spec.mode == Mode::State {
2565        let entity = match query.key.as_deref() {
2566            Some(key) => entity_cache
2567                .get(&view_spec.id, key)
2568                .await
2569                .map(|data| vec![(key.to_string(), data)])
2570                .unwrap_or_default(),
2571            None => vec![],
2572        };
2573        (entity, true)
2574    } else {
2575        (entity_cache.get_all(&view_spec.id).await, false)
2576    };
2577    select_query_entities(entities, query, preordered, apply_snapshot_limit)
2578}
2579
2580fn select_query_entities(
2581    mut entities: Vec<(String, Value)>,
2582    query: &SubscriptionQuery,
2583    preordered: bool,
2584    apply_snapshot_limit: bool,
2585) -> Vec<(String, Value)> {
2586    entities.retain(|(key, data)| query_matches_entity(query, key, data));
2587    if !preordered {
2588        entities.sort_by(|left, right| {
2589            let left_seq = left.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2590            let right_seq = right.1.get("_seq").and_then(Value::as_str).unwrap_or("");
2591            let order = if query.after.is_some() {
2592                cmp_seq(left_seq, right_seq)
2593            } else {
2594                cmp_seq(right_seq, left_seq)
2595            };
2596            order.then_with(|| left.0.cmp(&right.0))
2597        });
2598    }
2599
2600    let skip = query.skip.unwrap_or(0);
2601    let take = query.take.unwrap_or(usize::MAX);
2602    let mut selected: Vec<_> = entities.into_iter().skip(skip).take(take).collect();
2603    if apply_snapshot_limit {
2604        if let Some(limit) = query.snapshot_limit {
2605            selected.truncate(limit);
2606        }
2607    }
2608    selected
2609}
2610
2611fn query_matches_entity(query: &SubscriptionQuery, key: &str, data: &Value) -> bool {
2612    if !query.matches_key(key) {
2613        return false;
2614    }
2615    if let Some(partition) = &query.partition {
2616        if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
2617            return false;
2618        }
2619    }
2620    if let Some(after) = &query.after {
2621        let Some(seq) = data.get("_seq").and_then(Value::as_str) else {
2622            return false;
2623        };
2624        if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
2625            return false;
2626        }
2627    }
2628    query
2629        .filters
2630        .iter()
2631        .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
2632}
2633
2634fn value_at_dot_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
2635    path.split('.')
2636        .try_fold(value, |current, segment| current.get(segment))
2637}
2638
2639#[cfg(test)]
2640mod tests {
2641    use super::*;
2642    use crate::cache::EntityCacheConfig;
2643    use crate::view::{Delivery, Filters, Projection};
2644    use serde_json::json;
2645    use tokio::sync::oneshot;
2646
2647    fn list_spec() -> ViewSpec {
2648        ViewSpec {
2649            id: "Thing/list".to_string(),
2650            export: "Thing".to_string(),
2651            mode: Mode::List,
2652            wire_format: Default::default(),
2653            projection: Projection::all(),
2654            filters: Filters::all(),
2655            delivery: Delivery::default(),
2656            pipeline: None,
2657            source_view: None,
2658        }
2659    }
2660
2661    /// Plan the whole window the way `emit_collection_delta` does, so a test
2662    /// can assert on what a subscriber is actually sent.
2663    fn plan_window(
2664        current_keys: &[&str],
2665        next_keys: &[&str],
2666        envelope_key: &str,
2667        is_derived: bool,
2668        op: &str,
2669    ) -> Vec<(String, MemberAction)> {
2670        next_keys
2671            .iter()
2672            .map(|key| {
2673                let was_member = current_keys.contains(key);
2674                (
2675                    (*key).to_string(),
2676                    member_action(was_member, *key == envelope_key, is_derived, op),
2677                )
2678            })
2679            .collect()
2680    }
2681
2682    #[test]
2683    fn a_reordered_window_sends_only_the_mutated_entity() {
2684        // A `_seq`-descending list: mutating "1" moves it to the front and
2685        // shifts every other key down one. Only "1" changed, so only "1" is
2686        // sent — and it rides its own patch, not a full entity.
2687        let current = ["4", "3", "2", "1"];
2688        let next = ["1", "4", "3", "2"];
2689        let plan = plan_window(&current, &next, "1", false, "patch");
2690
2691        assert_eq!(
2692            plan,
2693            vec![
2694                ("1".to_string(), MemberAction::ForwardPatch),
2695                ("4".to_string(), MemberAction::Skip),
2696                ("3".to_string(), MemberAction::Skip),
2697                ("2".to_string(), MemberAction::Skip),
2698            ]
2699        );
2700    }
2701
2702    #[test]
2703    fn a_key_entering_the_window_gets_the_whole_entity() {
2704        // "5" has no local state for the subscriber to merge a patch into.
2705        let plan = plan_window(&["4", "3"], &["5", "4", "3"], "5", false, "patch");
2706        assert_eq!(
2707            plan,
2708            vec![
2709                ("5".to_string(), MemberAction::Upsert),
2710                ("4".to_string(), MemberAction::Skip),
2711                ("3".to_string(), MemberAction::Skip),
2712            ]
2713        );
2714    }
2715
2716    #[test]
2717    fn derived_views_still_send_whole_entities() {
2718        // The patch on the bus is scoped to the source view, so a derived
2719        // subscription cannot forward it verbatim (A4-150).
2720        let plan = plan_window(&["1", "2"], &["1", "2"], "1", true, "patch");
2721        assert_eq!(
2722            plan,
2723            vec![
2724                ("1".to_string(), MemberAction::Upsert),
2725                ("2".to_string(), MemberAction::Skip),
2726            ]
2727        );
2728    }
2729
2730    #[test]
2731    fn a_delete_envelope_never_forwards_a_patch() {
2732        // A surviving key on a delete envelope carries no mergeable patch.
2733        let plan = plan_window(&["1", "2"], &["1", "2"], "1", false, "delete");
2734        assert_eq!(
2735            plan,
2736            vec![
2737                ("1".to_string(), MemberAction::Upsert),
2738                ("2".to_string(), MemberAction::Skip),
2739            ]
2740        );
2741    }
2742
2743    #[test]
2744    fn an_unchanged_window_sends_one_frame_not_a_broadcast() {
2745        // The regression this guards: 500 members used to mean 500 full
2746        // entities on the wire for a single mutation.
2747        let keys: Vec<String> = (0..500).map(|index| index.to_string()).collect();
2748        let refs: Vec<&str> = keys.iter().map(String::as_str).collect();
2749        let plan = plan_window(&refs, &refs, "250", false, "patch");
2750
2751        let sent = plan
2752            .iter()
2753            .filter(|(_, action)| *action != MemberAction::Skip)
2754            .count();
2755        assert_eq!(sent, 1);
2756        assert_eq!(plan[250].1, MemberAction::ForwardPatch);
2757    }
2758
2759    fn list_message(key: &str, op: &str, seq: &str) -> Arc<BusMessage> {
2760        Arc::new(BusMessage {
2761            key: key.to_string(),
2762            entity: "Thing/list".to_string(),
2763            payload: Arc::new(Bytes::from(
2764                serde_json::to_vec(&json!({
2765                    "entity": "Thing/list",
2766                    "op": op,
2767                    "key": key,
2768                    "seq": seq,
2769                    "data": {},
2770                }))
2771                .unwrap(),
2772            )),
2773        })
2774    }
2775
2776    #[test]
2777    fn coalescing_emits_only_the_final_full_state_per_changed_key() {
2778        let current = vec![
2779            ("a".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2780            ("b".to_string(), json!({"count": 1, "_seq": "10:000001"})),
2781            (
2782                "stable".to_string(),
2783                json!({"count": 1, "_seq": "10:000001"}),
2784            ),
2785        ];
2786        let next = vec![
2787            ("a".to_string(), json!({"count": 3, "_seq": "10:000004"})),
2788            ("c".to_string(), json!({"count": 1, "_seq": "10:000003"})),
2789            (
2790                "stable".to_string(),
2791                json!({"count": 1, "_seq": "10:000001"}),
2792            ),
2793        ];
2794        let pending = HashMap::from([
2795            ("a".to_string(), list_message("a", "patch", "10:000004")),
2796            ("b".to_string(), list_message("b", "delete", "10:000002")),
2797            ("c".to_string(), list_message("c", "patch", "10:000003")),
2798        ]);
2799
2800        assert_eq!(
2801            plan_coalesced_collection_delta(&current, &next, &pending),
2802            vec![
2803                CollectionChange {
2804                    op: "delete",
2805                    key: "b".to_string(),
2806                    data: Value::Null,
2807                    seq: Some("10:000002".to_string()),
2808                },
2809                CollectionChange {
2810                    op: "upsert",
2811                    key: "a".to_string(),
2812                    data: json!({"count": 3, "_seq": "10:000004"}),
2813                    seq: Some("10:000004".to_string()),
2814                },
2815                CollectionChange {
2816                    op: "upsert",
2817                    key: "c".to_string(),
2818                    data: json!({"count": 1, "_seq": "10:000003"}),
2819                    seq: Some("10:000003".to_string()),
2820                },
2821            ]
2822        );
2823    }
2824
2825    #[test]
2826    fn a_delayed_delete_cannot_erase_a_newer_recreated_entity() {
2827        let recreated = json!({"balance": 2, "_seq": "10:000004"});
2828        assert!(source_delete_is_stale(Some(&recreated), Some("10:000002")));
2829        assert!(!source_delete_is_stale(Some(&recreated), Some("10:000004")));
2830        assert!(!source_delete_is_stale(Some(&recreated), None));
2831    }
2832
2833    #[test]
2834    fn a_lagged_snapshotless_subscription_gets_a_fatal_retryable_error() {
2835        let issue = SocketIssueMessage::subscription_lagged("balances".to_string(), 42);
2836        assert_eq!(issue.code, "subscription-lagged");
2837        assert!(issue.retryable);
2838        assert!(issue.fatal);
2839        assert!(issue.message.contains("42"));
2840        assert!(issue
2841            .suggested_action
2842            .unwrap()
2843            .contains("snapshots enabled"));
2844
2845        let append = SocketIssueMessage::append_subscription_lagged("trades".to_string(), 9);
2846        assert!(append.fatal);
2847        assert!(append.suggested_action.unwrap().contains("retained replay"));
2848    }
2849
2850    #[test]
2851    fn snapshot_batches_share_identity_and_completion() {
2852        let entities = ["one", "two", "three"].map(|key| SnapshotEntity {
2853            key: key.to_string(),
2854            data: json!({"key": key}),
2855        });
2856        let batches = create_snapshot_batches(
2857            &entities,
2858            SnapshotMetadata {
2859                subscription_id: "sub-1",
2860                snapshot_id: "snapshot-1",
2861                authoritative: true,
2862                mode: Mode::List,
2863                view_id: "Thing/list",
2864                key: None,
2865            },
2866            &SnapshotBatchConfig {
2867                initial_batch_size: 2,
2868                subsequent_batch_size: 1,
2869            },
2870        );
2871        assert_eq!(batches.len(), 2);
2872        assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
2873        assert!(batches
2874            .iter()
2875            .all(|batch| batch.snapshot_id == "snapshot-1"));
2876        assert!(!batches[0].complete);
2877        assert!(batches[1].complete);
2878        assert!(batches.iter().all(|batch| batch.authoritative));
2879    }
2880
2881    #[test]
2882    fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
2883        let batches = create_snapshot_batches(
2884            &[],
2885            SnapshotMetadata {
2886                subscription_id: "sub-1",
2887                snapshot_id: "snapshot-1",
2888                authoritative: false,
2889                mode: Mode::State,
2890                view_id: "Thing/state",
2891                key: Some("missing"),
2892            },
2893            &SnapshotBatchConfig {
2894                initial_batch_size: 1,
2895                subsequent_batch_size: 1,
2896            },
2897        );
2898        assert_eq!(batches.len(), 1);
2899        assert!(!batches[0].authoritative);
2900        assert!(batches[0].complete);
2901        assert_eq!(batches[0].key.as_deref(), Some("missing"));
2902    }
2903
2904    #[test]
2905    fn lag_recovery_is_authoritative_for_an_after_query() {
2906        let subscription = Subscription {
2907            protocol_version: PROTOCOL_VERSION,
2908            subscription_id: "sub-1".to_string(),
2909            query: SubscriptionQuery {
2910                view: "Thing/list".to_string(),
2911                after: Some("40:000000000010".to_string()),
2912                ..Default::default()
2913            },
2914            snapshot: Default::default(),
2915        };
2916
2917        assert!(!SnapshotPurpose::Initial.authoritative(&subscription));
2918        assert!(SnapshotPurpose::Recovery.authoritative(&subscription));
2919    }
2920
2921    #[test]
2922    fn dot_path_filters_are_exact_and_type_sensitive() {
2923        let mut query = SubscriptionQuery {
2924            view: "Thing/list".to_string(),
2925            ..Default::default()
2926        };
2927        query
2928            .filters
2929            .insert("state.status".to_string(), json!("open"));
2930        query.filters.insert("metrics.count".to_string(), json!(2));
2931        assert!(query_matches_entity(
2932            &query,
2933            "one",
2934            &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
2935        ));
2936        assert!(!query_matches_entity(
2937            &query,
2938            "one",
2939            &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
2940        ));
2941    }
2942
2943    #[test]
2944    fn take_and_skip_define_independent_deterministic_windows() {
2945        let entities: Vec<_> = (1..=6)
2946            .map(|id| {
2947                (
2948                    id.to_string(),
2949                    json!({"id": id, "_seq": format!("10:{id:012}")}),
2950                )
2951            })
2952            .collect();
2953        let first = SubscriptionQuery {
2954            view: "Thing/list".to_string(),
2955            take: Some(2),
2956            skip: Some(0),
2957            ..Default::default()
2958        };
2959        let second = SubscriptionQuery {
2960            skip: Some(2),
2961            ..first.clone()
2962        };
2963        let first_keys: Vec<_> = select_query_entities(entities.clone(), &first, false, false)
2964            .into_iter()
2965            .map(|(key, _)| key)
2966            .collect();
2967        let second_keys: Vec<_> = select_query_entities(entities, &second, false, false)
2968            .into_iter()
2969            .map(|(key, _)| key)
2970            .collect();
2971        assert_eq!(first_keys, ["6", "5"]);
2972        assert_eq!(second_keys, ["4", "3"]);
2973    }
2974
2975    #[tokio::test]
2976    async fn state_receiver_is_installed_before_snapshot_awaits() {
2977        let bus = BusManager::new();
2978        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
2979        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
2980        let bus_for_task = bus.clone();
2981        let task = tokio::spawn(async move {
2982            subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
2983                snapshot_started_tx.send(()).unwrap();
2984                release_snapshot_rx.await.unwrap();
2985            })
2986            .await
2987            .0
2988        });
2989        snapshot_started_rx.await.unwrap();
2990        bus.publish_state(
2991            "Thing/state",
2992            "one",
2993            Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
2994        )
2995        .await;
2996        release_snapshot_tx.send(()).unwrap();
2997        let mut receiver = task.await.unwrap();
2998        receiver.changed().await.unwrap();
2999        assert!(!receiver.borrow().payload.is_empty());
3000    }
3001
3002    async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
3003        let bus = BusManager::new();
3004        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
3005        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
3006        let bus_for_task = bus.clone();
3007        let task = tokio::spawn(async move {
3008            subscribe_list_then_snapshot(&bus_for_task, view, || async move {
3009                snapshot_started_tx.send(()).unwrap();
3010                release_snapshot_rx.await.unwrap();
3011            })
3012            .await
3013            .0
3014        });
3015        snapshot_started_rx.await.unwrap();
3016        bus.publish_list(
3017            view,
3018            Arc::new(BusMessage {
3019                key: "one".to_string(),
3020                entity: view.to_string(),
3021                payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
3022            }),
3023        )
3024        .await;
3025        release_snapshot_tx.send(()).unwrap();
3026        let mut receiver = task.await.unwrap();
3027        assert_eq!(receiver.recv().await.unwrap().key, "one");
3028    }
3029
3030    #[tokio::test]
3031    async fn list_receiver_is_installed_before_snapshot() {
3032        assert_list_receiver_precedes_snapshot("Thing/list").await;
3033    }
3034
3035    #[tokio::test]
3036    async fn append_receiver_is_installed_before_snapshot() {
3037        assert_list_receiver_precedes_snapshot("Thing/append").await;
3038    }
3039
3040    #[tokio::test]
3041    async fn derived_source_receiver_is_installed_before_snapshot() {
3042        assert_list_receiver_precedes_snapshot("Thing/list-source").await;
3043    }
3044
3045    #[tokio::test]
3046    async fn snapshot_limit_does_not_change_live_take_skip_membership() {
3047        let cache = EntityCache::with_config(EntityCacheConfig {
3048            max_entities_per_view: 10,
3049            ..Default::default()
3050        });
3051        for id in 1..=4 {
3052            cache
3053                .upsert(
3054                    "Thing/list",
3055                    &id.to_string(),
3056                    json!({"_seq": format!("10:{id:012}")}),
3057                )
3058                .await;
3059        }
3060        let query = SubscriptionQuery {
3061            view: "Thing/list".to_string(),
3062            take: Some(3),
3063            skip: Some(1),
3064            snapshot_limit: Some(1),
3065            ..Default::default()
3066        };
3067        let live = load_query_entities(&cache, None, &list_spec(), &query, false).await;
3068        let snapshot = load_query_entities(&cache, None, &list_spec(), &query, true).await;
3069        assert_eq!(live.len(), 3);
3070        assert_eq!(snapshot.len(), 1);
3071    }
3072
3073    #[test]
3074    fn fixture_manifest_covers_required_conformance_cases() {
3075        let manifest: Value = serde_json::from_str(include_str!(
3076            "../../../../tests/fixtures/websocket-v2/manifest.json"
3077        ))
3078        .unwrap();
3079        let names: HashSet<_> = manifest["fixtures"]
3080            .as_array()
3081            .unwrap()
3082            .iter()
3083            .filter_map(Value::as_str)
3084            .collect();
3085        for required in [
3086            "keyed-state.json",
3087            "list-windows.json",
3088            "filters.json",
3089            "multi-batch-authoritative.json",
3090            "empty-snapshot.json",
3091            "remove.json",
3092            "delete.json",
3093            "incremental-snapshot.json",
3094            "reconnect-replacement.json",
3095            "errors.json",
3096        ] {
3097            assert!(names.contains(required), "missing fixture {required}");
3098        }
3099
3100        for document in [
3101            include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
3102            include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
3103            include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
3104            include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
3105            include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
3106            include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
3107            include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
3108            include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
3109            include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
3110            include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
3111        ] {
3112            let fixture: Value = serde_json::from_str(document).unwrap();
3113            assert!(fixture["name"].is_string());
3114        }
3115    }
3116
3117    fn append_frame(key: &str, data: Value) -> Arc<Bytes> {
3118        let frame = json!({
3119            "entity": "Trade/append",
3120            "op": "patch",
3121            "key": key,
3122            "offset": 7,
3123            "data": data,
3124        });
3125        Arc::new(Bytes::from(serde_json::to_vec(&frame).unwrap()))
3126    }
3127
3128    /// A replayable subscription must honour the same predicates a live
3129    /// collection subscription does; otherwise a filtered consumer receives
3130    /// events outside the query it asked for.
3131    #[test]
3132    fn replay_delivery_applies_key_partition_and_filters() {
3133        let matching = append_frame("pool1", json!({"_partition": "us", "side": "buy"}));
3134        let other_partition = append_frame("pool1", json!({"_partition": "eu", "side": "buy"}));
3135        let other_side = append_frame("pool1", json!({"_partition": "us", "side": "sell"}));
3136
3137        let unfiltered = SubscriptionQuery {
3138            view: "Trade/append".to_string(),
3139            ..Default::default()
3140        };
3141        assert!(live_frame_matches(&unfiltered, "pool1", &matching));
3142        assert!(live_frame_matches(&unfiltered, "pool9", &matching));
3143
3144        let keyed = SubscriptionQuery {
3145            view: "Trade/append".to_string(),
3146            key: Some("pool1".to_string()),
3147            ..Default::default()
3148        };
3149        assert!(live_frame_matches(&keyed, "pool1", &matching));
3150        assert!(!live_frame_matches(&keyed, "pool2", &matching));
3151
3152        let partitioned = SubscriptionQuery {
3153            view: "Trade/append".to_string(),
3154            partition: Some("us".to_string()),
3155            ..Default::default()
3156        };
3157        assert!(live_frame_matches(&partitioned, "pool1", &matching));
3158        assert!(!live_frame_matches(&partitioned, "pool1", &other_partition));
3159
3160        let filtered = SubscriptionQuery {
3161            view: "Trade/append".to_string(),
3162            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3163            ..Default::default()
3164        };
3165        assert!(live_frame_matches(&filtered, "pool1", &matching));
3166        assert!(!live_frame_matches(&filtered, "pool1", &other_side));
3167    }
3168
3169    #[test]
3170    fn a_frame_without_decodable_data_does_not_satisfy_a_filter() {
3171        let filtered = SubscriptionQuery {
3172            view: "Trade/append".to_string(),
3173            filters: [("side".to_string(), json!("buy"))].into_iter().collect(),
3174            ..Default::default()
3175        };
3176        let garbage = Arc::new(Bytes::from_static(b"not json"));
3177        assert!(!live_frame_matches(&filtered, "pool1", &garbage));
3178    }
3179
3180    /// The recovery cursor must be the last offset delivered *before* the
3181    /// gap. Reporting the newest offset seen would step the consumer over
3182    /// the skipped records permanently.
3183    #[test]
3184    fn replay_lagged_recovers_from_before_the_gap() {
3185        let epoch = crate::journal::JournalEpoch::new();
3186        let issue = SocketIssueMessage::replay_lagged(
3187            Some("trades".to_string()),
3188            37,
3189            Some(crate::journal::Cursor {
3190                epoch: epoch.clone(),
3191                offset: 4180,
3192            }),
3193        );
3194        assert_eq!(issue.code, "replay-lagged");
3195        assert_eq!(issue.recover_from, Some(format!("{epoch}:4180")));
3196        assert!(
3197            issue.suggested_action.unwrap().contains("4180"),
3198            "the consumer is told exactly which cursor recovers the gap"
3199        );
3200
3201        // Nothing delivered yet: there is no pre-gap offset, so the whole
3202        // retained window is the recovery.
3203        let from_scratch = SocketIssueMessage::replay_lagged(Some("trades".to_string()), 9, None);
3204        assert_eq!(from_scratch.recover_from, None);
3205        assert!(from_scratch
3206            .suggested_action
3207            .unwrap()
3208            .contains("without `after`"));
3209    }
3210
3211    fn bus_message(key: &str) -> Arc<BusMessage> {
3212        Arc::new(BusMessage {
3213            key: key.to_string(),
3214            entity: "Trade/append".to_string(),
3215            payload: Arc::new(Bytes::from_static(b"{}")),
3216        })
3217    }
3218
3219    /// A long replay must keep the bus drained. The broadcast buffer is
3220    /// bounded, so a busy view would otherwise lap the replay and the first
3221    /// live `recv` would return `Lagged` — telling the client to resubscribe,
3222    /// starting another long replay, which laps again.
3223    #[tokio::test]
3224    async fn draining_during_a_replay_keeps_a_busy_view_from_lapping_it() {
3225        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(16);
3226        let mut pending = VecDeque::new();
3227        let mut lagged = None;
3228
3229        // Publish more than the channel holds, draining as a replay would
3230        // between sends.
3231        for index in 0..48 {
3232            sender.send(bus_message(&format!("k{index}"))).unwrap();
3233            drain_available(&mut receiver, &mut pending, &mut lagged);
3234        }
3235
3236        assert_eq!(lagged, None, "draining as we go means nothing is dropped");
3237        assert_eq!(pending.len(), 48, "every published frame is buffered");
3238    }
3239
3240    #[tokio::test]
3241    async fn a_replay_that_never_drains_is_reported_as_a_gap() {
3242        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3243        for index in 0..32 {
3244            sender.send(bus_message(&format!("k{index}"))).unwrap();
3245        }
3246
3247        let mut pending = VecDeque::new();
3248        let mut lagged = None;
3249        drain_available(&mut receiver, &mut pending, &mut lagged);
3250
3251        assert!(
3252            lagged.is_some(),
3253            "overflowing the bus is a gap, not silent truncation"
3254        );
3255    }
3256
3257    /// Everything still on the bus after a gap is on the far side of it.
3258    /// Buffering it would put those frames in front of the lag report and
3259    /// advance the recovery cursor past the records it is meant to recover.
3260    #[tokio::test]
3261    async fn nothing_after_a_gap_is_buffered_ahead_of_the_report() {
3262        let (sender, mut receiver) = broadcast::channel::<Arc<BusMessage>>(8);
3263        let mut pending = VecDeque::new();
3264        let mut lagged = None;
3265
3266        // Delivered and buffered normally.
3267        sender.send(bus_message("before")).unwrap();
3268        drain_available(&mut receiver, &mut pending, &mut lagged);
3269        assert_eq!(pending.len(), 1);
3270
3271        // Overflow the bus: everything published from here is past the gap.
3272        for index in 0..32 {
3273            sender.send(bus_message(&format!("lost{index}"))).unwrap();
3274        }
3275        drain_available(&mut receiver, &mut pending, &mut lagged);
3276        assert!(lagged.is_some());
3277
3278        // The bus still holds what survived the overflow, all of it past the
3279        // gap. A later iteration of the replay loop drains again.
3280        let buffered_at_gap = pending.len();
3281        drain_available(&mut receiver, &mut pending, &mut lagged);
3282        assert_eq!(
3283            pending.len(),
3284            buffered_at_gap,
3285            "post-gap frames must not join the pre-gap flush"
3286        );
3287        assert_eq!(pending.front().unwrap().key, "before");
3288    }
3289    /// The bus is subscribed before the tape is read, so a record published
3290    /// in that window arrives on both paths.
3291    #[test]
3292    fn the_seam_between_replay_and_live_neither_repeats_nor_skips() {
3293        let mut last_sent = Some(4211);
3294
3295        assert!(
3296            already_delivered(Some(4211), &mut last_sent),
3297            "the record the replay ended on must not be sent twice"
3298        );
3299        assert!(already_delivered(Some(4100), &mut last_sent));
3300        assert_eq!(last_sent, Some(4211), "a duplicate never moves the mark");
3301
3302        assert!(
3303            !already_delivered(Some(4212), &mut last_sent),
3304            "the next record is new"
3305        );
3306        assert_eq!(last_sent, Some(4212));
3307
3308        // A frame with no offset comes from a view with no tape; it cannot
3309        // have been replayed, and must not disturb the mark.
3310        assert!(!already_delivered(None, &mut last_sent));
3311        assert_eq!(last_sent, Some(4212));
3312    }
3313
3314    /// A subscription with no cursor has delivered nothing, so the first live
3315    /// frame is not a duplicate.
3316    #[test]
3317    fn a_fresh_subscription_delivers_its_first_live_frame() {
3318        let mut last_sent = None;
3319        assert!(!already_delivered(Some(0), &mut last_sent));
3320        assert_eq!(last_sent, Some(0));
3321    }
3322
3323    /// End-to-end over a real socket: the pieces above are unit-tested
3324    /// individually, but the thing a consumer actually does — reconnect with
3325    /// a stored cursor and keep reading — only exists once a subscription is
3326    /// attached to a connection.
3327    mod over_a_socket {
3328        use super::*;
3329        use crate::journal::{EventJournal, JournalConfig};
3330        use crate::projector::Projector;
3331        use crate::{MutationBatch, SlotContext};
3332        use arete_interpreter::Mutation;
3333        use futures_util::{SinkExt, StreamExt};
3334        use std::time::Duration;
3335        use tokio::net::{TcpListener, TcpStream};
3336        use tokio::sync::mpsc;
3337        use tokio_tungstenite::tungstenite::Message;
3338        use tokio_tungstenite::{client_async, WebSocketStream};
3339
3340        const RETAINED: u64 = 600;
3341
3342        fn append_index() -> ViewIndex {
3343            let mut index = ViewIndex::new();
3344            index.add_spec(ViewSpec {
3345                id: "Trade/append".to_string(),
3346                export: "Trade".to_string(),
3347                mode: Mode::Append,
3348                wire_format: Default::default(),
3349                projection: Projection::all(),
3350                filters: Filters::all(),
3351                delivery: Delivery::default(),
3352                pipeline: None,
3353                source_view: None,
3354            });
3355            index
3356        }
3357
3358        fn trade(index: u64) -> MutationBatch {
3359            MutationBatch::with_slot_context(
3360                vec![Mutation {
3361                    export: "Trade".to_string(),
3362                    key: json!(format!("pool{}", index % 4)),
3363                    patch: json!({"trade": index}),
3364                    append: vec![],
3365                }]
3366                .into_iter()
3367                .collect(),
3368                SlotContext::new(100 + index / 3, index % 3),
3369            )
3370        }
3371
3372        struct Harness {
3373            addr: SocketAddr,
3374            journal: Arc<EventJournal>,
3375            tx: mpsc::Sender<MutationBatch>,
3376        }
3377
3378        impl Harness {
3379            async fn start() -> Self {
3380                let view_index = Arc::new(append_index());
3381                let entity_cache = EntityCache::new();
3382                let bus_manager = BusManager::new();
3383                let journal = Arc::new(EventJournal::new(JournalConfig {
3384                    enabled: true,
3385                    max_bytes_per_view: u64::MAX,
3386                    max_records_per_view: 10_000,
3387                    max_age: Duration::from_secs(3_600),
3388                }));
3389
3390                let (tx, rx) = mpsc::channel::<MutationBatch>(256);
3391                tokio::spawn(
3392                    Projector::new(
3393                        view_index.clone(),
3394                        bus_manager.clone(),
3395                        entity_cache.clone(),
3396                        rx,
3397                        #[cfg(feature = "otel")]
3398                        None,
3399                    )
3400                    .with_journal(journal.clone())
3401                    .run(),
3402                );
3403
3404                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3405                let addr = listener.local_addr().unwrap();
3406                let server = WebSocketServer::new(
3407                    addr,
3408                    bus_manager,
3409                    entity_cache,
3410                    view_index,
3411                    #[cfg(feature = "otel")]
3412                    None,
3413                )
3414                .with_journal(journal.clone());
3415                let (acceptor, _cleanup) = server.into_acceptor();
3416                tokio::spawn(async move { acceptor.serve_listener(listener).await });
3417
3418                Self { addr, journal, tx }
3419            }
3420
3421            async fn publish(&self, range: std::ops::Range<u64>) {
3422                for index in range {
3423                    self.tx.send(trade(index)).await.unwrap();
3424                }
3425                let (ack, wait) = oneshot::channel();
3426                self.tx
3427                    .send(MutationBatch::flush_marker(ack))
3428                    .await
3429                    .unwrap();
3430                wait.await.unwrap();
3431            }
3432
3433            async fn connect(&self) -> WebSocketStream<TcpStream> {
3434                let stream = TcpStream::connect(self.addr).await.unwrap();
3435                client_async(format!("ws://{}/", self.addr), stream)
3436                    .await
3437                    .unwrap()
3438                    .0
3439            }
3440        }
3441
3442        async fn next_frame(socket: &mut WebSocketStream<TcpStream>) -> Value {
3443            loop {
3444                let message = tokio::time::timeout(Duration::from_secs(10), socket.next())
3445                    .await
3446                    .expect("the server answers within the timeout")
3447                    .expect("the stream stays open")
3448                    .expect("a readable frame");
3449                // Control and data frames arrive as binary; issue frames as
3450                // text. Both are JSON.
3451                let bytes = match &message {
3452                    Message::Text(text) => text.as_bytes(),
3453                    Message::Binary(bytes) => bytes.as_ref(),
3454                    _ => continue,
3455                };
3456                return serde_json::from_slice(bytes).expect("frames are JSON");
3457            }
3458        }
3459
3460        /// Collect `count` event frames, ignoring anything else on the wire.
3461        async fn collect_trades(socket: &mut WebSocketStream<TcpStream>, count: usize) -> Vec<u64> {
3462            let mut offsets = Vec::with_capacity(count);
3463            while offsets.len() < count {
3464                let frame = next_frame(socket).await;
3465                assert_ne!(
3466                    frame["type"], "error",
3467                    "no error frame should interrupt delivery: {frame}"
3468                );
3469                if let Some(offset) = frame["offset"].as_u64() {
3470                    offsets.push(offset);
3471                }
3472            }
3473            offsets
3474        }
3475
3476        /// The headline claim: reconnecting with a stored cursor delivers every
3477        /// event published since it, in order, and then continues live without
3478        /// a duplicate or a hole at the seam.
3479        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3480        async fn a_reconnect_replays_from_a_cursor_and_continues_live() {
3481            let harness = Harness::start().await;
3482            harness.publish(0..RETAINED).await;
3483
3484            let cursor = harness.journal.window("Trade/append").await;
3485            let stored = format!("{}:{}", cursor.epoch, 99);
3486
3487            let mut socket = harness.connect().await;
3488            socket
3489                .send(Message::Text(
3490                    json!({
3491                        "type": "subscribe",
3492                        "protocolVersion": 2,
3493                        "subscriptionId": "trades",
3494                        "query": {"view": "Trade/append", "after": stored},
3495                    })
3496                    .to_string()
3497                    .into(),
3498                ))
3499                .await
3500                .unwrap();
3501
3502            let ack = next_frame(&mut socket).await;
3503            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3504            assert_eq!(ack["replayWindow"]["next"], json!(RETAINED));
3505
3506            // Well over the 500 a single read or buffer would cover.
3507            let replayed = collect_trades(&mut socket, (RETAINED - 100) as usize).await;
3508            assert_eq!(
3509                replayed,
3510                (100..RETAINED).collect::<Vec<_>>(),
3511                "every event after the cursor, in order, exactly once"
3512            );
3513
3514            // Published only now, so these can only arrive over the live path.
3515            harness.publish(RETAINED..RETAINED + 40).await;
3516            let live = collect_trades(&mut socket, 40).await;
3517            assert_eq!(
3518                live,
3519                (RETAINED..RETAINED + 40).collect::<Vec<_>>(),
3520                "the live stream resumes exactly where the replay stopped"
3521            );
3522
3523            socket.close(None).await.ok();
3524        }
3525
3526        /// Events published *during* the replay must still arrive. The replay
3527        /// and the live subscription are separate reads of the same tape, and
3528        /// the seam between them is where a naive implementation drops or
3529        /// repeats.
3530        #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
3531        async fn events_published_during_a_replay_are_not_lost() {
3532            let harness = Harness::start().await;
3533            harness.publish(0..RETAINED).await;
3534
3535            let epoch = harness.journal.window("Trade/append").await.epoch;
3536            let mut socket = harness.connect().await;
3537            socket
3538                .send(Message::Text(
3539                    json!({
3540                        "type": "subscribe",
3541                        "protocolVersion": 2,
3542                        "subscriptionId": "trades",
3543                        "query": {"view": "Trade/append", "after": format!("{epoch}:0")},
3544                    })
3545                    .to_string()
3546                    .into(),
3547                ))
3548                .await
3549                .unwrap();
3550            let ack = next_frame(&mut socket).await;
3551            assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3552
3553            // Keep publishing while the replay is still draining.
3554            harness.publish(RETAINED..RETAINED + 200).await;
3555
3556            let total = (RETAINED + 200 - 1) as usize;
3557            let delivered = collect_trades(&mut socket, total).await;
3558            assert_eq!(
3559                delivered,
3560                (1..RETAINED + 200).collect::<Vec<_>>(),
3561                "replay and live output join without a gap or a repeat"
3562            );
3563
3564            socket.close(None).await.ok();
3565        }
3566
3567        /// A cursor from another tape lifetime is refused rather than served
3568        /// as a continuation, and the refusal releases the subscription id.
3569        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3570        async fn a_stale_epoch_is_refused_and_frees_the_subscription_id() {
3571            let harness = Harness::start().await;
3572            harness.publish(0..50).await;
3573
3574            let mut socket = harness.connect().await;
3575            socket
3576                .send(Message::Text(
3577                    json!({
3578                        "type": "subscribe",
3579                        "protocolVersion": 2,
3580                        "subscriptionId": "trades",
3581                        "query": {
3582                            "view": "Trade/append",
3583                            "after": format!("{}:10", crate::journal::JournalEpoch::new()),
3584                        },
3585                    })
3586                    .to_string()
3587                    .into(),
3588                ))
3589                .await
3590                .unwrap();
3591
3592            let error = next_frame(&mut socket).await;
3593            assert_eq!(error["type"], "error", "unexpected frame: {error}");
3594            assert_eq!(error["code"], "cursor-epoch-changed");
3595
3596            // The documented recovery is to resubscribe without a cursor. That
3597            // only works if the refused attempt released the id.
3598            socket
3599                .send(Message::Text(
3600                    json!({
3601                        "type": "subscribe",
3602                        "protocolVersion": 2,
3603                        "subscriptionId": "trades",
3604                        "query": {"view": "Trade/append"},
3605                    })
3606                    .to_string()
3607                    .into(),
3608                ))
3609                .await
3610                .unwrap();
3611            let ack = next_frame(&mut socket).await;
3612            assert_eq!(
3613                ack["op"], "subscribed",
3614                "the refused id must be reusable: {ack}"
3615            );
3616
3617            assert_eq!(collect_trades(&mut socket, 50).await.len(), 50);
3618            socket.close(None).await.ok();
3619        }
3620
3621        const ROUND: &str = "Round/state";
3622
3623        /// A state view with no projector: the test writes the cache and the
3624        /// bus itself, so it controls exactly when the subscriber can run.
3625        #[derive(Clone)]
3626        struct StateServer {
3627            addr: SocketAddr,
3628            bus_manager: BusManager,
3629            entity_cache: EntityCache,
3630        }
3631
3632        impl StateServer {
3633            async fn start() -> Self {
3634                let mut index = ViewIndex::new();
3635                index.add_spec(ViewSpec {
3636                    id: ROUND.to_string(),
3637                    export: "Round".to_string(),
3638                    mode: Mode::State,
3639                    wire_format: Default::default(),
3640                    projection: Projection::all(),
3641                    filters: Filters::all(),
3642                    delivery: Delivery::default(),
3643                    pipeline: None,
3644                    source_view: None,
3645                });
3646                let bus_manager = BusManager::new();
3647                let entity_cache = EntityCache::new();
3648                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3649                let addr = listener.local_addr().unwrap();
3650                let server = WebSocketServer::new(
3651                    addr,
3652                    bus_manager.clone(),
3653                    entity_cache.clone(),
3654                    Arc::new(index),
3655                    #[cfg(feature = "otel")]
3656                    None,
3657                );
3658                let (acceptor, _cleanup) = server.into_acceptor();
3659                tokio::spawn(async move { acceptor.serve_listener(listener).await });
3660                Self {
3661                    addr,
3662                    bus_manager,
3663                    entity_cache,
3664                }
3665            }
3666
3667            /// Publish a patch the way the projector does: stamp its `_seq`,
3668            /// write the cache, then the key's bus.
3669            async fn publish(&self, key: &str, mut patch: Value, seq: &str) {
3670                patch["_seq"] = Value::String(seq.to_string());
3671                self.entity_cache
3672                    .upsert_with_append(ROUND, key, patch.clone(), &[])
3673                    .await;
3674                let frame = json!({
3675                    "mode": "state",
3676                    "entity": ROUND,
3677                    "op": "patch",
3678                    "key": key,
3679                    "data": patch,
3680                    "seq": seq,
3681                });
3682                self.bus_manager
3683                    .publish_state(ROUND, key, Arc::new(Bytes::from(frame.to_string())))
3684                    .await;
3685            }
3686
3687            /// Subscribe to one key and read through its snapshot.
3688            async fn subscribe(&self, key: &str) -> WebSocketStream<TcpStream> {
3689                let stream = TcpStream::connect(self.addr).await.unwrap();
3690                let mut socket = client_async(format!("ws://{}/", self.addr), stream)
3691                    .await
3692                    .unwrap()
3693                    .0;
3694                socket
3695                    .send(Message::Text(
3696                        json!({
3697                            "type": "subscribe",
3698                            "protocolVersion": 2,
3699                            "subscriptionId": "round",
3700                            "query": {"view": ROUND, "key": key},
3701                        })
3702                        .to_string()
3703                        .into(),
3704                    ))
3705                    .await
3706                    .unwrap();
3707                let ack = next_frame(&mut socket).await;
3708                assert_eq!(ack["op"], "subscribed", "unexpected ack: {ack}");
3709                loop {
3710                    let frame = next_frame(&mut socket).await;
3711                    assert_eq!(frame["op"], "snapshot", "unexpected frame: {frame}");
3712                    if frame["complete"] == true {
3713                        return socket;
3714                    }
3715                }
3716            }
3717        }
3718
3719        /// Two patches to one key, published before the subscriber reads the
3720        /// first, must both reach it. The bus holds only the latest frame, so
3721        /// forwarding it alone drops the fields of the one it replaced.
3722        // One worker, so tasks run one at a time. The handshake needs the
3723        // multi-threaded runtime.
3724        #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
3725        async fn a_state_subscriber_keeps_the_fields_of_an_overwritten_patch() {
3726            let server = StateServer::start().await;
3727            server
3728                .publish("7", json!({"id": 7}), "100:000000000001")
3729                .await;
3730            let mut socket = server.subscribe("7").await;
3731
3732            // A lone patch is forwarded as it was published. This one is an
3733            // account update, numbered by its write version.
3734            server
3735                .publish(
3736                    "7",
3737                    json!({"results": {"slot_hash": "abc"}}),
3738                    "140:003836292257",
3739                )
3740                .await;
3741            let frame = next_frame(&mut socket).await;
3742            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3743            assert_eq!(frame["seq"], "140:003836292257");
3744
3745            // Two instruction patches from later in the same slot, numbered by
3746            // transaction index. Both writes happen in one task that never
3747            // yields to the subscriber, so the second replaces the first on
3748            // the bus before the subscriber wakes.
3749            let writer = server.clone();
3750            tokio::spawn(async move {
3751                writer
3752                    .publish("7", json!({"entropy": {"seed": "def"}}), "140:000000001200")
3753                    .await;
3754                writer
3755                    .publish("7", json!({"total": 1}), "140:000000001200")
3756                    .await;
3757            })
3758            .await
3759            .unwrap();
3760            let frame = next_frame(&mut socket).await;
3761            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3762            assert_eq!(
3763                frame["data"],
3764                json!({
3765                    "id": 7,
3766                    "results": {"slot_hash": "abc"},
3767                    "entropy": {"seed": "def"},
3768                    "total": 1,
3769                    "_seq": "140:000000001200",
3770                })
3771            );
3772            // The latest seq sorts below the one already delivered, and a
3773            // client would drop the frame as stale.
3774            assert!(frame.get("seq").is_none(), "unexpected seq: {frame}");
3775
3776            // Having caught up, the subscriber forwards patches again.
3777            server
3778                .publish("7", json!({"total": 2}), "141:000000000002")
3779                .await;
3780            let frame = next_frame(&mut socket).await;
3781            assert_eq!(frame["op"], "patch", "unexpected frame: {frame}");
3782            assert_eq!(
3783                frame["data"],
3784                json!({"total": 2, "_seq": "141:000000000002"})
3785            );
3786
3787            socket.close(None).await.ok();
3788        }
3789    }
3790
3791    /// Session tokens that expire while their socket is open.
3792    mod session_expiry {
3793        use super::*;
3794        use crate::websocket::auth::SignedSessionAuthPlugin;
3795        use arete_auth::{KeyClass, SessionClaims, SigningKey, TokenSigner, TokenVerifier};
3796        use futures_util::{SinkExt, StreamExt};
3797        use std::time::Duration;
3798        use tokio::net::{TcpListener, TcpStream};
3799        use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
3800        use tokio_tungstenite::tungstenite::protocol::CloseFrame;
3801        use tokio_tungstenite::tungstenite::Message;
3802        use tokio_tungstenite::{client_async, WebSocketStream};
3803
3804        struct Server {
3805            addr: SocketAddr,
3806            signer: TokenSigner,
3807        }
3808
3809        impl Server {
3810            async fn start() -> Self {
3811                let signing_key = SigningKey::generate();
3812                let verifier =
3813                    TokenVerifier::new(signing_key.verifying_key(), "test-issuer", "test-audience");
3814                let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3815                let addr = listener.local_addr().unwrap();
3816                let server = WebSocketServer::new(
3817                    addr,
3818                    BusManager::new(),
3819                    EntityCache::new(),
3820                    Arc::new(ViewIndex::new()),
3821                    #[cfg(feature = "otel")]
3822                    None,
3823                )
3824                .with_auth_plugin(Arc::new(SignedSessionAuthPlugin::new(verifier)));
3825                let (acceptor, _cleanup) = server.into_acceptor();
3826                tokio::spawn(async move { acceptor.serve_listener(listener).await });
3827                Self {
3828                    addr,
3829                    signer: TokenSigner::new(signing_key, "test-issuer"),
3830                }
3831            }
3832
3833            fn token(&self, ttl_seconds: u64) -> String {
3834                let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
3835                    .with_scope("read")
3836                    .with_key_class(KeyClass::Secret)
3837                    .with_ttl(ttl_seconds)
3838                    .build();
3839                self.signer.sign(claims).unwrap()
3840            }
3841
3842            async fn connect(&self, token: &str) -> WebSocketStream<TcpStream> {
3843                let stream = TcpStream::connect(self.addr).await.unwrap();
3844                client_async(format!("ws://{}/?hs_token={token}", self.addr), stream)
3845                    .await
3846                    .unwrap()
3847                    .0
3848            }
3849        }
3850
3851        async fn send_json(socket: &mut WebSocketStream<TcpStream>, message: Value) {
3852            socket
3853                .send(Message::Text(message.to_string().into()))
3854                .await
3855                .unwrap();
3856        }
3857
3858        /// The close frame, if the server closes the socket within `wait`.
3859        async fn close_within(
3860            socket: &mut WebSocketStream<TcpStream>,
3861            wait: Duration,
3862        ) -> Option<Option<CloseFrame>> {
3863            tokio::time::timeout(wait, async {
3864                while let Some(Ok(message)) = socket.next().await {
3865                    if let Message::Close(frame) = message {
3866                        return Some(frame);
3867                    }
3868                }
3869                Some(None)
3870            })
3871            .await
3872            .ok()
3873            .flatten()
3874        }
3875
3876        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3877        async fn an_expired_session_is_closed_with_the_reason() {
3878            let server = Server::start().await;
3879            let mut socket = server.connect(&server.token(2)).await;
3880
3881            tokio::time::sleep(Duration::from_secs(3)).await;
3882            send_json(&mut socket, json!({"type": "ping"})).await;
3883
3884            let frame = close_within(&mut socket, Duration::from_secs(5))
3885                .await
3886                .expect("the server closes the expired session")
3887                .expect("the close frame carries a reason");
3888            assert_eq!(frame.code, CloseCode::Policy);
3889            assert_eq!(
3890                frame.reason.as_str(),
3891                "token-expired: Authentication token expired"
3892            );
3893        }
3894
3895        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3896        async fn a_session_refreshed_in_band_outlives_its_first_token() {
3897            let server = Server::start().await;
3898            let mut socket = server.connect(&server.token(2)).await;
3899
3900            send_json(
3901                &mut socket,
3902                json!({"type": "refresh_auth", "token": server.token(3_600)}),
3903            )
3904            .await;
3905            let reply = tokio::time::timeout(Duration::from_secs(5), async {
3906                while let Some(Ok(message)) = socket.next().await {
3907                    if let Message::Text(text) = message {
3908                        return serde_json::from_str::<Value>(text.as_str()).ok();
3909                    }
3910                }
3911                None
3912            })
3913            .await
3914            .expect("the server answers the refresh")
3915            .expect("the answer is JSON");
3916            assert_eq!(reply["success"], true, "refresh accepted: {reply}");
3917
3918            tokio::time::sleep(Duration::from_secs(3)).await;
3919            send_json(&mut socket, json!({"type": "ping"})).await;
3920            assert!(
3921                close_within(&mut socket, Duration::from_millis(1_500))
3922                    .await
3923                    .is_none(),
3924                "the socket stays open on the refreshed token"
3925            );
3926        }
3927    }
3928}