Skip to main content

arete_server/websocket/
server.rs

1use crate::bus::{BusManager, BusMessage};
2use crate::cache::{cmp_seq, EntityCache, SnapshotBatchConfig};
3use crate::compression::maybe_compress;
4use crate::view::{ViewIndex, ViewSpec};
5use crate::websocket::auth::{
6    AuthContext, AuthDecision, AuthDeny, ConnectionAuthRequest, WebSocketAuthPlugin,
7};
8use crate::websocket::client_manager::{ClientManager, RateLimitConfig};
9use crate::websocket::frame::{
10    apply_wire_format, Frame, Mode, SnapshotEntity, SnapshotFrame, SortConfig, SortOrder,
11    SubscribedFrame, UnsubscribedFrame,
12};
13use crate::websocket::subscription::{
14    ClientMessage, RefreshAuthRequest, RefreshAuthResponse, SocketIssueMessage, Subscription,
15    SubscriptionQuery, Unsubscription, PROTOCOL_VERSION,
16};
17use crate::websocket::usage::{WebSocketUsageEmitter, WebSocketUsageEvent};
18use anyhow::Result;
19use bytes::Bytes;
20use futures_util::StreamExt;
21use serde::Serialize;
22use serde_json::Value;
23use std::collections::{HashMap, HashSet};
24use std::future::Future;
25use std::net::SocketAddr;
26use std::sync::Arc;
27use std::time::Instant;
28use tokio::net::{TcpListener, TcpStream};
29use tokio::sync::{broadcast, watch};
30use tokio_tungstenite::{
31    accept_hdr_async,
32    tungstenite::{
33        handshake::server::{ErrorResponse as HandshakeErrorResponse, Request, Response},
34        http::{header::CONTENT_TYPE, StatusCode},
35        Error as WsError,
36    },
37};
38use tokio_util::sync::CancellationToken;
39use tracing::{debug, error, info, info_span, warn, Instrument};
40use uuid::Uuid;
41
42#[cfg(feature = "otel")]
43use crate::metrics::Metrics;
44
45#[derive(Clone, Default)]
46struct WsMetrics {
47    #[cfg(feature = "otel")]
48    inner: Option<Arc<Metrics>>,
49}
50
51impl WsMetrics {
52    #[cfg(feature = "otel")]
53    fn new(inner: Option<Arc<Metrics>>) -> Self {
54        Self { inner }
55    }
56
57    fn connection_opened(&self, metering_key: Option<&str>) {
58        #[cfg(not(feature = "otel"))]
59        let _ = metering_key;
60        #[cfg(feature = "otel")]
61        if let Some(metrics) = &self.inner {
62            if let Some(metering_key) = metering_key {
63                metrics.record_ws_connection_with_metering(metering_key);
64            } else {
65                metrics.record_ws_connection();
66            }
67        }
68    }
69
70    fn connection_closed(&self, duration_secs: f64, metering_key: Option<&str>) {
71        #[cfg(not(feature = "otel"))]
72        let _ = (duration_secs, metering_key);
73        #[cfg(feature = "otel")]
74        if let Some(metrics) = &self.inner {
75            if let Some(metering_key) = metering_key {
76                metrics.record_ws_disconnection_with_metering(duration_secs, metering_key);
77            } else {
78                metrics.record_ws_disconnection(duration_secs);
79            }
80        }
81    }
82
83    fn message_received(&self, metering_key: Option<&str>) {
84        #[cfg(not(feature = "otel"))]
85        let _ = metering_key;
86        #[cfg(feature = "otel")]
87        if let Some(metrics) = &self.inner {
88            if let Some(metering_key) = metering_key {
89                metrics.record_ws_message_received_with_metering(metering_key);
90            } else {
91                metrics.record_ws_message_received();
92            }
93        }
94    }
95
96    fn message_sent(&self) {
97        #[cfg(feature = "otel")]
98        if let Some(metrics) = &self.inner {
99            metrics.record_ws_message_sent();
100        }
101    }
102
103    fn subscription_created(&self, view: &str, metering_key: Option<&str>) {
104        #[cfg(not(feature = "otel"))]
105        let _ = (view, metering_key);
106        #[cfg(feature = "otel")]
107        if let Some(metrics) = &self.inner {
108            if let Some(metering_key) = metering_key {
109                metrics.record_subscription_created_with_metering(view, metering_key);
110            } else {
111                metrics.record_subscription_created(view);
112            }
113        }
114    }
115
116    fn subscription_removed(&self, view: &str, metering_key: Option<&str>) {
117        #[cfg(not(feature = "otel"))]
118        let _ = (view, metering_key);
119        #[cfg(feature = "otel")]
120        if let Some(metrics) = &self.inner {
121            if let Some(metering_key) = metering_key {
122                metrics.record_subscription_removed_with_metering(view, metering_key);
123            } else {
124                metrics.record_subscription_removed(view);
125            }
126        }
127    }
128
129    fn protocol_error(&self, code: &str) {
130        #[cfg(not(feature = "otel"))]
131        let _ = code;
132        #[cfg(feature = "otel")]
133        if let Some(metrics) = &self.inner {
134            metrics.record_ws_protocol_error(code);
135        }
136    }
137}
138
139async fn handle_refresh_auth(
140    client_id: Uuid,
141    refresh_req: &RefreshAuthRequest,
142    client_manager: &ClientManager,
143    auth_plugin: &Arc<dyn WebSocketAuthPlugin>,
144) {
145    let refresh_result: Result<AuthContext, String> = if let Some(signed_plugin) = auth_plugin
146        .as_any()
147        .downcast_ref::<crate::websocket::auth::SignedSessionAuthPlugin>()
148    {
149        signed_plugin
150            .verify_refresh_token(&refresh_req.token)
151            .await
152            .map_err(|error| error.reason)
153    } else {
154        Err("In-band auth refresh not supported with current auth plugin".to_string())
155    };
156
157    let response = match refresh_result {
158        Ok(new_context) => {
159            let expires_at = new_context.expires_at;
160            if client_manager.update_client_auth(client_id, new_context) {
161                RefreshAuthResponse {
162                    success: true,
163                    error: None,
164                    expires_at: Some(expires_at),
165                }
166            } else {
167                RefreshAuthResponse {
168                    success: false,
169                    error: Some("client-not-found".to_string()),
170                    expires_at: None,
171                }
172            }
173        }
174        Err(error) => {
175            let code = if error.contains("expired") {
176                "token-expired"
177            } else if error.contains("signature") {
178                "token-invalid-signature"
179            } else if error.contains("issuer") {
180                "token-invalid-issuer"
181            } else if error.contains("audience") {
182                "token-invalid-audience"
183            } else {
184                "token-invalid"
185            };
186            RefreshAuthResponse {
187                success: false,
188                error: Some(code.to_string()),
189                expires_at: None,
190            }
191        }
192    };
193
194    if let Ok(json) = serde_json::to_string(&response) {
195        let _ = client_manager.send_text_to_client(client_id, json).await;
196    }
197}
198
199async fn send_socket_issue(
200    client_id: Uuid,
201    client_manager: &ClientManager,
202    deny: &AuthDeny,
203    fatal: bool,
204    subscription_id: Option<String>,
205) {
206    let message = SocketIssueMessage::from_auth_deny(deny, fatal, subscription_id);
207    if let Ok(json) = serde_json::to_string(&message) {
208        let _ = client_manager.send_text_to_client(client_id, json).await;
209    }
210}
211
212async fn send_protocol_issue(
213    client_id: Uuid,
214    client_manager: &ClientManager,
215    metrics: &WsMetrics,
216    subscription_id: Option<String>,
217    code: &str,
218    message: impl Into<String>,
219) {
220    metrics.protocol_error(code);
221    let issue = SocketIssueMessage::protocol(subscription_id, code, message);
222    if let Ok(json) = serde_json::to_string(&issue) {
223        let _ = client_manager.send_text_to_client(client_id, json).await;
224    }
225}
226
227fn key_class_label(key_class: arete_auth::KeyClass) -> &'static str {
228    match key_class {
229        arete_auth::KeyClass::Secret => "secret",
230        arete_auth::KeyClass::Publishable => "publishable",
231    }
232}
233
234fn emit_usage_event(
235    usage_emitter: &Option<Arc<dyn WebSocketUsageEmitter>>,
236    event: WebSocketUsageEvent,
237) {
238    if let Some(emitter) = usage_emitter.clone() {
239        tokio::spawn(async move {
240            emitter.emit(event).await;
241        });
242    }
243}
244
245fn usage_identity(
246    auth_context: Option<&AuthContext>,
247) -> (
248    Option<String>,
249    Option<String>,
250    Option<String>,
251    Option<String>,
252) {
253    match auth_context {
254        Some(context) => (
255            Some(context.metering_key.clone()),
256            Some(context.subject.clone()),
257            Some(key_class_label(context.key_class).to_string()),
258            context.deployment_id.clone(),
259        ),
260        None => (None, None, None, None),
261    }
262}
263
264fn emit_update_sent_for_client(
265    usage_emitter: &Option<Arc<dyn WebSocketUsageEmitter>>,
266    client_manager: &ClientManager,
267    client_id: Uuid,
268    view_id: &str,
269    bytes: usize,
270) {
271    let auth_context = client_manager.get_auth_context(client_id);
272    let (metering_key, subject, _, deployment_id) = usage_identity(auth_context.as_ref());
273    emit_usage_event(
274        usage_emitter,
275        WebSocketUsageEvent::UpdateSent {
276            client_id: client_id.to_string(),
277            deployment_id,
278            metering_key,
279            subject,
280            view_id: view_id.to_string(),
281            messages: 1,
282            bytes: bytes as u64,
283        },
284    );
285}
286
287#[derive(Clone)]
288struct SubscriptionContext {
289    client_id: Uuid,
290    client_manager: ClientManager,
291    bus_manager: BusManager,
292    entity_cache: EntityCache,
293    view_index: Arc<ViewIndex>,
294    usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
295    metrics: WsMetrics,
296}
297
298pub struct WebSocketServer {
299    bind_addr: SocketAddr,
300    client_manager: ClientManager,
301    bus_manager: BusManager,
302    entity_cache: EntityCache,
303    view_index: Arc<ViewIndex>,
304    max_clients: usize,
305    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
306    usage_emitter: Option<Arc<dyn WebSocketUsageEmitter>>,
307    rate_limit_config: Option<RateLimitConfig>,
308    #[cfg(feature = "otel")]
309    metrics: Option<Arc<Metrics>>,
310}
311
312impl WebSocketServer {
313    #[cfg(feature = "otel")]
314    pub fn new(
315        bind_addr: SocketAddr,
316        bus_manager: BusManager,
317        entity_cache: EntityCache,
318        view_index: Arc<ViewIndex>,
319        metrics: Option<Arc<Metrics>>,
320    ) -> Self {
321        Self {
322            bind_addr,
323            client_manager: ClientManager::new(),
324            bus_manager,
325            entity_cache,
326            view_index,
327            max_clients: 10_000,
328            auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
329            usage_emitter: None,
330            rate_limit_config: None,
331            metrics,
332        }
333    }
334
335    #[cfg(not(feature = "otel"))]
336    pub fn new(
337        bind_addr: SocketAddr,
338        bus_manager: BusManager,
339        entity_cache: EntityCache,
340        view_index: Arc<ViewIndex>,
341    ) -> Self {
342        Self {
343            bind_addr,
344            client_manager: ClientManager::new(),
345            bus_manager,
346            entity_cache,
347            view_index,
348            max_clients: 10_000,
349            auth_plugin: Arc::new(crate::websocket::auth::AllowAllAuthPlugin),
350            usage_emitter: None,
351            rate_limit_config: None,
352        }
353    }
354
355    pub fn with_max_clients(mut self, max_clients: usize) -> Self {
356        self.max_clients = max_clients;
357        self
358    }
359
360    pub fn with_auth_plugin(mut self, auth_plugin: Arc<dyn WebSocketAuthPlugin>) -> Self {
361        self.auth_plugin = auth_plugin;
362        self
363    }
364
365    pub fn with_usage_emitter(mut self, usage_emitter: Arc<dyn WebSocketUsageEmitter>) -> Self {
366        self.usage_emitter = Some(usage_emitter);
367        self
368    }
369
370    pub fn with_rate_limit_config(mut self, config: RateLimitConfig) -> Self {
371        self.rate_limit_config = Some(config);
372        self
373    }
374
375    pub async fn start(self) -> Result<()> {
376        info!(
377            "Starting WebSocket server on {} (max_clients: {})",
378            self.bind_addr, self.max_clients
379        );
380        let listener = TcpListener::bind(&self.bind_addr).await?;
381        let client_manager = self
382            .rate_limit_config
383            .map(ClientManager::with_config)
384            .unwrap_or(self.client_manager);
385        client_manager.start_cleanup_task();
386
387        #[cfg(feature = "otel")]
388        let metrics = WsMetrics::new(self.metrics.clone());
389        #[cfg(not(feature = "otel"))]
390        let metrics = WsMetrics::default();
391
392        loop {
393            match listener.accept().await {
394                Ok((stream, addr)) => {
395                    if client_manager.client_count() >= self.max_clients {
396                        warn!("Rejecting connection from {}: max clients reached", addr);
397                        continue;
398                    }
399
400                    let context = SubscriptionContext {
401                        client_id: Uuid::nil(),
402                        client_manager: client_manager.clone(),
403                        bus_manager: self.bus_manager.clone(),
404                        entity_cache: self.entity_cache.clone(),
405                        view_index: self.view_index.clone(),
406                        usage_emitter: self.usage_emitter.clone(),
407                        metrics: metrics.clone(),
408                    };
409                    let auth_plugin = self.auth_plugin.clone();
410                    tokio::spawn(
411                        async move {
412                            if let Err(error) =
413                                handle_connection(stream, context, addr, auth_plugin).await
414                            {
415                                error!("WebSocket connection error: {}", error);
416                            }
417                        }
418                        .instrument(info_span!("ws.connection", %addr)),
419                    );
420                }
421                Err(error) => error!("Failed to accept connection: {}", error),
422            }
423        }
424    }
425}
426
427#[derive(Debug, Clone)]
428struct HandshakeReject {
429    status: StatusCode,
430    body: crate::websocket::auth::ErrorResponse,
431    error_code: String,
432    retry_after_secs: Option<u64>,
433}
434
435impl HandshakeReject {
436    fn from_deny(deny: &AuthDeny) -> Self {
437        let retry_after_secs = match deny.retry_policy {
438            crate::websocket::auth::RetryPolicy::RetryAfter(duration) => Some(duration.as_secs()),
439            _ => None,
440        };
441        Self {
442            status: StatusCode::from_u16(deny.http_status).unwrap_or(StatusCode::UNAUTHORIZED),
443            body: deny.to_error_response(),
444            error_code: deny.code.to_string(),
445            retry_after_secs,
446        }
447    }
448}
449
450fn build_handshake_error_response(
451    response: &Response,
452    reject: &HandshakeReject,
453) -> HandshakeErrorResponse {
454    let mut builder = Response::builder()
455        .status(reject.status)
456        .version(response.version())
457        .header(CONTENT_TYPE, "application/json; charset=utf-8")
458        .header("X-Error-Code", &reject.error_code)
459        .header("Cache-Control", "no-store");
460    if let Some(retry_after_secs) = reject.retry_after_secs {
461        builder = builder.header("Retry-After", retry_after_secs.to_string());
462    }
463    let body = serde_json::to_string(&reject.body).unwrap_or_else(|_| {
464        format!(
465            r#"{{"error":"{}","message":"{}","code":"{}","retryable":false}}"#,
466            reject.body.error, reject.body.message, reject.body.code
467        )
468    });
469    builder
470        .body(Some(body))
471        .expect("handshake rejection response should build")
472}
473
474#[allow(clippy::result_large_err)]
475async fn accept_authorized_connection(
476    stream: TcpStream,
477    remote_addr: SocketAddr,
478    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
479    client_manager: ClientManager,
480) -> Result<Option<(tokio_tungstenite::WebSocketStream<TcpStream>, AuthContext)>> {
481    use std::sync::Mutex;
482
483    let capture: Arc<Mutex<Option<Result<AuthContext, HandshakeReject>>>> =
484        Arc::new(Mutex::new(None));
485    let capture_ref = capture.clone();
486    let auth_plugin_ref = auth_plugin.clone();
487    let manager_ref = client_manager.clone();
488
489    let handshake_result = accept_hdr_async(stream, move |request: &Request, response| {
490        let request = ConnectionAuthRequest::from_http_request(remote_addr, request);
491        let result = tokio::task::block_in_place(|| {
492            tokio::runtime::Handle::current().block_on(async {
493                match auth_plugin_ref.authorize(&request).await {
494                    AuthDecision::Allow(context) => manager_ref
495                        .check_connection_allowed(remote_addr, &Some(context.clone()))
496                        .await
497                        .map(|()| context)
498                        .map_err(|deny| HandshakeReject::from_deny(&deny)),
499                    AuthDecision::Deny(deny) => Err(HandshakeReject::from_deny(&deny)),
500                }
501            })
502        });
503        *capture_ref.lock().expect("capture lock poisoned") = Some(result.clone());
504        match result {
505            Ok(_) => Ok(response),
506            Err(reject) => Err(build_handshake_error_response(&response, &reject)),
507        }
508    })
509    .await;
510
511    let auth_result = capture.lock().expect("capture lock poisoned").take();
512    match handshake_result {
513        Ok(stream) => match auth_result {
514            Some(Ok(context)) => Ok(Some((stream, context))),
515            Some(Err(reject)) => Err(anyhow::anyhow!(
516                "handshake unexpectedly succeeded after rejection: {}",
517                reject.body.message
518            )),
519            None => Err(anyhow::anyhow!("no auth result captured during handshake")),
520        },
521        Err(WsError::Http(_)) => Ok(None),
522        Err(error) => Err(error.into()),
523    }
524}
525
526async fn handle_connection(
527    stream: TcpStream,
528    mut context: SubscriptionContext,
529    remote_addr: SocketAddr,
530    auth_plugin: Arc<dyn WebSocketAuthPlugin>,
531) -> Result<()> {
532    let Some((ws_stream, auth_context)) = accept_authorized_connection(
533        stream,
534        remote_addr,
535        auth_plugin.clone(),
536        context.client_manager.clone(),
537    )
538    .await?
539    else {
540        return Ok(());
541    };
542
543    let client_id = Uuid::new_v4();
544    context.client_id = client_id;
545    let connection_start = Instant::now();
546    let (metering_key, subject, key_class, deployment_id) = usage_identity(Some(&auth_context));
547    context.metrics.connection_opened(metering_key.as_deref());
548
549    let (ws_sender, mut ws_receiver) = ws_stream.split();
550    context
551        .client_manager
552        .add_client(client_id, ws_sender, Some(auth_context), remote_addr);
553    emit_usage_event(
554        &context.usage_emitter,
555        WebSocketUsageEvent::ConnectionEstablished {
556            client_id: client_id.to_string(),
557            remote_addr: remote_addr.to_string(),
558            deployment_id: deployment_id.clone(),
559            metering_key: metering_key.clone(),
560            subject: subject.clone(),
561            key_class,
562        },
563    );
564
565    let mut active_subscriptions: HashMap<String, String> = HashMap::new();
566    while let Some(message) = ws_receiver.next().await {
567        let message = match message {
568            Ok(message) => message,
569            Err(error) => {
570                warn!("WebSocket error for client {}: {}", client_id, error);
571                break;
572            }
573        };
574        if message.is_close() {
575            break;
576        }
577        context.client_manager.update_client_last_seen(client_id);
578        if !message.is_text() {
579            continue;
580        }
581        if let Err(deny) = context
582            .client_manager
583            .check_inbound_message_allowed(client_id)
584        {
585            send_socket_issue(client_id, &context.client_manager, &deny, true, None).await;
586            break;
587        }
588        context.metrics.message_received(metering_key.as_deref());
589
590        let text = match message.to_text() {
591            Ok(text) => text,
592            Err(_) => continue,
593        };
594        let client_message = match serde_json::from_str::<ClientMessage>(text) {
595            Ok(message) => message,
596            Err(parse_error) => {
597                let subscription_id = extract_subscription_id(text);
598                send_protocol_issue(
599                    client_id,
600                    &context.client_manager,
601                    &context.metrics,
602                    subscription_id,
603                    "malformed-message",
604                    format!("invalid protocol v2 message: {parse_error}"),
605                )
606                .await;
607                continue;
608            }
609        };
610
611        match client_message {
612            ClientMessage::Subscribe(subscription) => {
613                let subscription_id = subscription.subscription_id.clone();
614                if let Err(message) = subscription.validate() {
615                    send_protocol_issue(
616                        client_id,
617                        &context.client_manager,
618                        &context.metrics,
619                        Some(subscription_id),
620                        "invalid-subscription",
621                        message,
622                    )
623                    .await;
624                    continue;
625                }
626                if let Err(deny) = context
627                    .client_manager
628                    .check_subscription_allowed(client_id)
629                    .await
630                {
631                    send_socket_issue(
632                        client_id,
633                        &context.client_manager,
634                        &deny,
635                        false,
636                        Some(subscription_id),
637                    )
638                    .await;
639                    continue;
640                }
641
642                let cancel_token = CancellationToken::new();
643                if !context
644                    .client_manager
645                    .add_client_subscription(
646                        client_id,
647                        subscription_id.clone(),
648                        cancel_token.clone(),
649                    )
650                    .await
651                {
652                    send_protocol_issue(
653                        client_id,
654                        &context.client_manager,
655                        &context.metrics,
656                        Some(subscription_id),
657                        "duplicate-subscription-id",
658                        "subscriptionId is already active on this connection",
659                    )
660                    .await;
661                    continue;
662                }
663
664                let view = subscription.query.view.clone();
665                if let Err(error) = attach_client_to_bus(&context, subscription, cancel_token).await
666                {
667                    context
668                        .client_manager
669                        .remove_client_subscription(client_id, &subscription_id)
670                        .await;
671                    send_protocol_issue(
672                        client_id,
673                        &context.client_manager,
674                        &context.metrics,
675                        Some(subscription_id),
676                        "subscription-rejected",
677                        error.to_string(),
678                    )
679                    .await;
680                    continue;
681                }
682
683                active_subscriptions.insert(subscription_id, view.clone());
684                context
685                    .metrics
686                    .subscription_created(&view, metering_key.as_deref());
687                emit_usage_event(
688                    &context.usage_emitter,
689                    WebSocketUsageEvent::SubscriptionCreated {
690                        client_id: client_id.to_string(),
691                        deployment_id: deployment_id.clone(),
692                        metering_key: metering_key.clone(),
693                        subject: subject.clone(),
694                        view_id: view,
695                    },
696                );
697            }
698            ClientMessage::Unsubscribe(unsubscription) => {
699                handle_unsubscribe(
700                    &context,
701                    unsubscription,
702                    &mut active_subscriptions,
703                    metering_key.as_deref(),
704                    &deployment_id,
705                    &metering_key,
706                    &subject,
707                )
708                .await;
709            }
710            ClientMessage::Ping => debug!("Received ping from client {}", client_id),
711            ClientMessage::RefreshAuth(request) => {
712                handle_refresh_auth(client_id, &request, &context.client_manager, &auth_plugin)
713                    .await;
714            }
715        }
716    }
717
718    context
719        .client_manager
720        .cancel_all_client_subscriptions(client_id)
721        .await;
722    context.client_manager.remove_client(client_id);
723    if let Some(rate_limiter) = context.client_manager.rate_limiter().cloned() {
724        rate_limiter.remove_client_buckets(client_id).await;
725    }
726    for view in active_subscriptions.values() {
727        context
728            .metrics
729            .subscription_removed(view, metering_key.as_deref());
730        emit_usage_event(
731            &context.usage_emitter,
732            WebSocketUsageEvent::SubscriptionRemoved {
733                client_id: client_id.to_string(),
734                deployment_id: deployment_id.clone(),
735                metering_key: metering_key.clone(),
736                subject: subject.clone(),
737                view_id: view.clone(),
738            },
739        );
740    }
741    let duration = connection_start.elapsed().as_secs_f64();
742    context
743        .metrics
744        .connection_closed(duration, metering_key.as_deref());
745    emit_usage_event(
746        &context.usage_emitter,
747        WebSocketUsageEvent::ConnectionClosed {
748            client_id: client_id.to_string(),
749            deployment_id,
750            metering_key,
751            subject,
752            duration_secs: Some(duration),
753            subscription_count: u32::try_from(active_subscriptions.len()).unwrap_or(u32::MAX),
754        },
755    );
756    Ok(())
757}
758
759#[allow(clippy::too_many_arguments)]
760async fn handle_unsubscribe(
761    context: &SubscriptionContext,
762    unsubscription: Unsubscription,
763    active_subscriptions: &mut HashMap<String, String>,
764    metrics_metering_key: Option<&str>,
765    deployment_id: &Option<String>,
766    usage_metering_key: &Option<String>,
767    subject: &Option<String>,
768) {
769    let subscription_id = unsubscription.subscription_id.clone();
770    if let Err(message) = unsubscription.validate() {
771        send_protocol_issue(
772            context.client_id,
773            &context.client_manager,
774            &context.metrics,
775            Some(subscription_id),
776            "invalid-unsubscription",
777            message,
778        )
779        .await;
780        return;
781    }
782
783    if !context
784        .client_manager
785        .remove_client_subscription(context.client_id, &subscription_id)
786        .await
787    {
788        send_protocol_issue(
789            context.client_id,
790            &context.client_manager,
791            &context.metrics,
792            Some(subscription_id),
793            "unknown-subscription-id",
794            "subscriptionId is not active on this connection",
795        )
796        .await;
797        return;
798    }
799
800    let Some(view) = active_subscriptions.remove(&subscription_id) else {
801        return;
802    };
803    let _ = send_control_frame(context, &UnsubscribedFrame::new(subscription_id), &view);
804    context
805        .metrics
806        .subscription_removed(&view, metrics_metering_key);
807    emit_usage_event(
808        &context.usage_emitter,
809        WebSocketUsageEvent::SubscriptionRemoved {
810            client_id: context.client_id.to_string(),
811            deployment_id: deployment_id.clone(),
812            metering_key: usage_metering_key.clone(),
813            subject: subject.clone(),
814            view_id: view,
815        },
816    );
817}
818
819fn extract_subscription_id(text: &str) -> Option<String> {
820    serde_json::from_str::<Value>(text)
821        .ok()?
822        .get("subscriptionId")?
823        .as_str()
824        .map(str::to_string)
825}
826
827struct SnapshotMetadata<'a> {
828    subscription_id: &'a str,
829    snapshot_id: &'a str,
830    authoritative: bool,
831    mode: Mode,
832    view_id: &'a str,
833    key: Option<&'a str>,
834}
835
836fn create_snapshot_batches(
837    entities: &[SnapshotEntity],
838    metadata: SnapshotMetadata<'_>,
839    batch_config: &SnapshotBatchConfig,
840) -> Vec<SnapshotFrame> {
841    if entities.is_empty() {
842        return vec![SnapshotFrame {
843            protocol_version: PROTOCOL_VERSION,
844            subscription_id: metadata.subscription_id.to_string(),
845            snapshot_id: metadata.snapshot_id.to_string(),
846            authoritative: metadata.authoritative,
847            mode: metadata.mode,
848            export: metadata.view_id.to_string(),
849            op: "snapshot",
850            key: metadata.key.map(str::to_string),
851            data: vec![],
852            complete: true,
853        }];
854    }
855
856    let mut batches = Vec::new();
857    let mut offset = 0;
858    while offset < entities.len() {
859        let configured_size = if offset == 0 {
860            batch_config.initial_batch_size
861        } else {
862            batch_config.subsequent_batch_size
863        };
864        let end = (offset + configured_size.max(1)).min(entities.len());
865        batches.push(SnapshotFrame {
866            protocol_version: PROTOCOL_VERSION,
867            subscription_id: metadata.subscription_id.to_string(),
868            snapshot_id: metadata.snapshot_id.to_string(),
869            authoritative: metadata.authoritative,
870            mode: metadata.mode,
871            export: metadata.view_id.to_string(),
872            op: "snapshot",
873            key: metadata.key.map(str::to_string),
874            data: entities[offset..end].to_vec(),
875            complete: end == entities.len(),
876        });
877        offset = end;
878    }
879    batches
880}
881
882async fn send_snapshot_batches(
883    context: &SubscriptionContext,
884    subscription: &Subscription,
885    entities: &[SnapshotEntity],
886    mode: Mode,
887    batch_config: &SnapshotBatchConfig,
888) -> Result<()> {
889    let snapshot_id = Uuid::new_v4().to_string();
890    let authoritative = subscription.query.after.is_none();
891    let frames = create_snapshot_batches(
892        entities,
893        SnapshotMetadata {
894            subscription_id: &subscription.subscription_id,
895            snapshot_id: &snapshot_id,
896            authoritative,
897            mode,
898            view_id: &subscription.query.view,
899            key: subscription.query.key.as_deref(),
900        },
901        batch_config,
902    );
903
904    for frame in frames {
905        let rows = frame.data.len() as u32;
906        let json = serde_json::to_vec(&frame)?;
907        let payload = maybe_compress(&json);
908        let bytes = payload.as_bytes().len() as u64;
909        context
910            .client_manager
911            .send_compressed_async(context.client_id, payload)
912            .await
913            .map_err(|error| anyhow::anyhow!("failed to send snapshot: {error}"))?;
914        context.metrics.message_sent();
915
916        let auth_context = context.client_manager.get_auth_context(context.client_id);
917        let (metering_key, subject, _, deployment_id) = usage_identity(auth_context.as_ref());
918        emit_usage_event(
919            &context.usage_emitter,
920            WebSocketUsageEvent::SnapshotSent {
921                client_id: context.client_id.to_string(),
922                deployment_id,
923                metering_key,
924                subject,
925                view_id: subscription.query.view.clone(),
926                rows,
927                messages: 1,
928                bytes,
929            },
930        );
931    }
932    Ok(())
933}
934
935fn extract_sort_config(view_spec: &ViewSpec) -> Option<SortConfig> {
936    if let Some(sort) = view_spec
937        .pipeline
938        .as_ref()
939        .and_then(|pipeline| pipeline.sort.as_ref())
940    {
941        return Some(SortConfig {
942            field: sort.field_path.clone(),
943            order: match sort.order {
944                crate::materialized_view::SortOrder::Asc => SortOrder::Asc,
945                crate::materialized_view::SortOrder::Desc => SortOrder::Desc,
946            },
947        });
948    }
949    (view_spec.mode == Mode::List).then(|| SortConfig {
950        field: vec!["_seq".to_string()],
951        order: SortOrder::Desc,
952    })
953}
954
955fn send_control_frame<T: Serialize>(
956    context: &SubscriptionContext,
957    frame: &T,
958    view_id: &str,
959) -> Result<()> {
960    let json = serde_json::to_vec(frame)?;
961    let bytes = json.len();
962    context
963        .client_manager
964        .send_to_client(context.client_id, Arc::new(Bytes::from(json)))
965        .map_err(|error| anyhow::anyhow!("failed to send control frame: {error}"))?;
966    context.metrics.message_sent();
967    emit_update_sent_for_client(
968        &context.usage_emitter,
969        &context.client_manager,
970        context.client_id,
971        view_id,
972        bytes,
973    );
974    Ok(())
975}
976
977fn send_subscribed_frame(
978    context: &SubscriptionContext,
979    subscription: &Subscription,
980    view_spec: &ViewSpec,
981) -> Result<()> {
982    let frame = SubscribedFrame::new(
983        subscription.subscription_id.clone(),
984        subscription.query.clone(),
985        view_spec.mode,
986        extract_sort_config(view_spec),
987    );
988    send_control_frame(context, &frame, &subscription.query.view)
989}
990
991fn enforce_snapshot_limit(context: &SubscriptionContext, rows: usize) -> Result<()> {
992    context
993        .client_manager
994        .check_snapshot_allowed(context.client_id, u32::try_from(rows).unwrap_or(u32::MAX))
995        .map_err(|deny| anyhow::anyhow!(deny.reason))
996}
997
998async fn subscribe_state_then_snapshot<F, Fut, T>(
999    bus_manager: &BusManager,
1000    view_id: &str,
1001    key: &str,
1002    snapshot: F,
1003) -> (watch::Receiver<Arc<Bytes>>, T)
1004where
1005    F: FnOnce() -> Fut,
1006    Fut: Future<Output = T>,
1007{
1008    let mut receiver = bus_manager.get_or_create_state_bus(view_id, key).await;
1009    receiver.borrow_and_update();
1010    let snapshot = snapshot().await;
1011    (receiver, snapshot)
1012}
1013
1014async fn subscribe_list_then_snapshot<F, Fut, T>(
1015    bus_manager: &BusManager,
1016    view_id: &str,
1017    snapshot: F,
1018) -> (broadcast::Receiver<Arc<BusMessage>>, T)
1019where
1020    F: FnOnce() -> Fut,
1021    Fut: Future<Output = T>,
1022{
1023    let receiver = bus_manager.get_or_create_list_bus(view_id).await;
1024    let snapshot = snapshot().await;
1025    (receiver, snapshot)
1026}
1027
1028async fn attach_client_to_bus(
1029    context: &SubscriptionContext,
1030    mut subscription: Subscription,
1031    cancel_token: CancellationToken,
1032) -> Result<()> {
1033    let view_spec = context
1034        .view_index
1035        .get_view(&subscription.query.view)
1036        .cloned()
1037        .ok_or_else(|| anyhow::anyhow!("unknown view: {}", subscription.query.view))?;
1038
1039    if view_spec.mode == Mode::State && !view_spec.is_derived() && subscription.query.key.is_none()
1040    {
1041        return Err(anyhow::anyhow!("state subscriptions require query.key"));
1042    }
1043    if view_spec.is_derived() && subscription.query.take.is_none() {
1044        subscription.query.take = view_spec
1045            .pipeline
1046            .as_ref()
1047            .and_then(|pipeline| pipeline.limit);
1048    }
1049
1050    if view_spec.mode == Mode::State && !view_spec.is_derived() {
1051        attach_state_subscription(context, subscription, view_spec, cancel_token).await
1052    } else {
1053        attach_collection_subscription(context, subscription, view_spec, cancel_token).await
1054    }
1055}
1056
1057async fn attach_state_subscription(
1058    context: &SubscriptionContext,
1059    subscription: Subscription,
1060    view_spec: ViewSpec,
1061    cancel_token: CancellationToken,
1062) -> Result<()> {
1063    let view_id = subscription.query.view.clone();
1064    let key = subscription.query.key.clone().unwrap_or_default();
1065    let query = subscription.query.clone();
1066    let cache = context.entity_cache.clone();
1067    let view_spec_for_snapshot = view_spec.clone();
1068    let (mut receiver, initial) =
1069        subscribe_state_then_snapshot(&context.bus_manager, &view_id, &key, move || async move {
1070            load_query_entities(&cache, None, &view_spec_for_snapshot, &query, false).await
1071        })
1072        .await;
1073
1074    let mut snapshot_entities = initial.clone();
1075    if let Some(limit) = subscription.query.snapshot_limit {
1076        snapshot_entities.truncate(limit);
1077    }
1078    enforce_snapshot_limit(context, snapshot_entities.len())?;
1079    send_subscribed_frame(context, &subscription, &view_spec)?;
1080    if subscription.snapshot.enabled {
1081        send_snapshot_batches(
1082            context,
1083            &subscription,
1084            &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1085            view_spec.mode,
1086            &context.entity_cache.snapshot_config(),
1087        )
1088        .await?;
1089    }
1090
1091    let task_context = context.clone();
1092    let subscription_id = subscription.subscription_id.clone();
1093    let query = subscription.query.clone();
1094    let view_spec_task = view_spec.clone();
1095    let span_view = view_id.clone();
1096    let span_key = key.clone();
1097    tokio::spawn(
1098        async move {
1099            let mut member = !initial.is_empty();
1100            loop {
1101                tokio::select! {
1102                    _ = cancel_token.cancelled() => break,
1103                    changed = receiver.changed() => {
1104                        if changed.is_err() {
1105                            break;
1106                        }
1107                        let payload = receiver.borrow().clone();
1108                        let metadata = source_frame_metadata(&payload);
1109                        if metadata.op == "delete" {
1110                            task_context.entity_cache.remove(&query.view, &key).await;
1111                            if member && send_membership_frame(
1112                                &task_context,
1113                                &subscription_id,
1114                                &view_spec_task,
1115                                "delete",
1116                                &key,
1117                                Value::Null,
1118                                metadata.seq,
1119                            ).is_err() {
1120                                break;
1121                            }
1122                            member = false;
1123                            continue;
1124                        }
1125
1126                        let selected = load_query_entities(
1127                            &task_context.entity_cache,
1128                            None,
1129                            &view_spec_task,
1130                            &query,
1131                            false,
1132                        ).await;
1133                        let is_member = !selected.is_empty();
1134                        let result = match (member, is_member) {
1135                            (true, true) => send_scoped_source_payload(
1136                                &task_context,
1137                                &subscription_id,
1138                                &query.view,
1139                                payload,
1140                            ),
1141                            (false, true) => {
1142                                let (entity_key, data) = selected.into_iter().next().unwrap();
1143                                send_membership_frame(
1144                                    &task_context,
1145                                    &subscription_id,
1146                                    &view_spec_task,
1147                                    "upsert",
1148                                    &entity_key,
1149                                    data,
1150                                    metadata.seq,
1151                                )
1152                            }
1153                            (true, false) => send_membership_frame(
1154                                &task_context,
1155                                &subscription_id,
1156                                &view_spec_task,
1157                                "remove",
1158                                &key,
1159                                Value::Null,
1160                                metadata.seq,
1161                            ),
1162                            (false, false) => Ok(()),
1163                        };
1164                        if result.is_err() {
1165                            break;
1166                        }
1167                        member = is_member;
1168                    }
1169                }
1170            }
1171        }
1172        .instrument(info_span!("ws.subscribe.state", client_id = %context.client_id, view = %span_view, key = %span_key)),
1173    );
1174    Ok(())
1175}
1176
1177async fn attach_collection_subscription(
1178    context: &SubscriptionContext,
1179    subscription: Subscription,
1180    view_spec: ViewSpec,
1181    cancel_token: CancellationToken,
1182) -> Result<()> {
1183    let view_id = subscription.query.view.clone();
1184    let source_view_id = view_spec
1185        .source_view
1186        .clone()
1187        .unwrap_or_else(|| view_id.clone());
1188    let query = subscription.query.clone();
1189    let cache = context.entity_cache.clone();
1190    let sorted_caches = view_spec
1191        .is_derived()
1192        .then(|| context.view_index.sorted_caches());
1193    let view_spec_for_snapshot = view_spec.clone();
1194    let (mut receiver, initial_membership) =
1195        subscribe_list_then_snapshot(&context.bus_manager, &source_view_id, move || async move {
1196            load_query_entities(
1197                &cache,
1198                sorted_caches,
1199                &view_spec_for_snapshot,
1200                &query,
1201                false,
1202            )
1203            .await
1204        })
1205        .await;
1206
1207    let mut snapshot_entities = initial_membership.clone();
1208    if let Some(limit) = subscription.query.snapshot_limit {
1209        snapshot_entities.truncate(limit);
1210    }
1211    enforce_snapshot_limit(context, snapshot_entities.len())?;
1212    send_subscribed_frame(context, &subscription, &view_spec)?;
1213    if subscription.snapshot.enabled {
1214        send_snapshot_batches(
1215            context,
1216            &subscription,
1217            &to_wire_snapshot_entities(snapshot_entities, &view_spec),
1218            view_spec.mode,
1219            &context.entity_cache.snapshot_config(),
1220        )
1221        .await?;
1222    }
1223
1224    let task_context = context.clone();
1225    let subscription_id = subscription.subscription_id.clone();
1226    let query = subscription.query.clone();
1227    let view_spec_task = view_spec.clone();
1228    let span_view = view_id.clone();
1229    tokio::spawn(
1230        async move {
1231            let mut current = initial_membership;
1232            loop {
1233                tokio::select! {
1234                    _ = cancel_token.cancelled() => break,
1235                    received = receiver.recv() => {
1236                        let envelope = match received {
1237                            Ok(envelope) => envelope,
1238                            Err(broadcast::error::RecvError::Lagged(_)) => {
1239                                warn!("Subscription {} lagged; closing to preserve membership correctness", subscription_id);
1240                                break;
1241                            }
1242                            Err(broadcast::error::RecvError::Closed) => break,
1243                        };
1244                        let metadata = source_frame_metadata(&envelope.payload);
1245                        if metadata.op == "delete" {
1246                            task_context.entity_cache.remove(&source_view_id, &envelope.key).await;
1247                            if view_spec_task.is_derived() {
1248                                let caches = task_context.view_index.sorted_caches();
1249                                let mut guard = caches.write().await;
1250                                if let Some(cache) = guard.get_mut(&query.view) {
1251                                    cache.remove(&envelope.key);
1252                                }
1253                            }
1254                        }
1255
1256                        let sorted_caches = view_spec_task
1257                            .is_derived()
1258                            .then(|| task_context.view_index.sorted_caches());
1259                        let next = load_query_entities(
1260                            &task_context.entity_cache,
1261                            sorted_caches,
1262                            &view_spec_task,
1263                            &query,
1264                            false,
1265                        ).await;
1266                        if emit_collection_delta(
1267                            &task_context,
1268                            &subscription_id,
1269                            &view_spec_task,
1270                            &current,
1271                            &next,
1272                            &envelope,
1273                            &metadata,
1274                        ).is_err() {
1275                            break;
1276                        }
1277                        current = next;
1278                    }
1279                }
1280            }
1281        }
1282        .instrument(info_span!("ws.subscribe.collection", client_id = %context.client_id, view = %span_view)),
1283    );
1284    Ok(())
1285}
1286
1287#[derive(Default)]
1288struct SourceFrameMetadata {
1289    op: String,
1290    seq: Option<String>,
1291}
1292
1293fn source_frame_metadata(payload: &[u8]) -> SourceFrameMetadata {
1294    serde_json::from_slice::<Value>(payload)
1295        .ok()
1296        .map(|value| SourceFrameMetadata {
1297            op: value
1298                .get("op")
1299                .and_then(Value::as_str)
1300                .unwrap_or_default()
1301                .to_string(),
1302            seq: value.get("seq").and_then(Value::as_str).map(str::to_string),
1303        })
1304        .unwrap_or_default()
1305}
1306
1307fn send_scoped_source_payload(
1308    context: &SubscriptionContext,
1309    subscription_id: &str,
1310    view_id: &str,
1311    payload: Arc<Bytes>,
1312) -> Result<()> {
1313    let mut value: Value = serde_json::from_slice(&payload)?;
1314    let object = value
1315        .as_object_mut()
1316        .ok_or_else(|| anyhow::anyhow!("source frame is not an object"))?;
1317    object.insert("protocolVersion".to_string(), Value::from(PROTOCOL_VERSION));
1318    object.insert(
1319        "subscriptionId".to_string(),
1320        Value::String(subscription_id.to_string()),
1321    );
1322    let encoded = Arc::new(Bytes::from(serde_json::to_vec(&value)?));
1323    let bytes = encoded.len();
1324    context
1325        .client_manager
1326        .send_to_client(context.client_id, encoded)
1327        .map_err(|error| anyhow::anyhow!("failed to send live frame: {error}"))?;
1328    context.metrics.message_sent();
1329    emit_update_sent_for_client(
1330        &context.usage_emitter,
1331        &context.client_manager,
1332        context.client_id,
1333        view_id,
1334        bytes,
1335    );
1336    Ok(())
1337}
1338
1339fn send_membership_frame(
1340    context: &SubscriptionContext,
1341    subscription_id: &str,
1342    view_spec: &ViewSpec,
1343    op: &str,
1344    key: &str,
1345    mut data: Value,
1346    seq: Option<String>,
1347) -> Result<()> {
1348    apply_wire_format(&mut data, &view_spec.wire_format);
1349    let frame = Frame::scoped(
1350        subscription_id,
1351        view_spec.mode,
1352        &view_spec.id,
1353        op,
1354        key,
1355        data,
1356        seq,
1357    );
1358    let encoded = Arc::new(Bytes::from(serde_json::to_vec(&frame)?));
1359    let bytes = encoded.len();
1360    context
1361        .client_manager
1362        .send_to_client(context.client_id, encoded)
1363        .map_err(|error| anyhow::anyhow!("failed to send membership frame: {error}"))?;
1364    context.metrics.message_sent();
1365    emit_update_sent_for_client(
1366        &context.usage_emitter,
1367        &context.client_manager,
1368        context.client_id,
1369        &view_spec.id,
1370        bytes,
1371    );
1372    Ok(())
1373}
1374
1375fn emit_collection_delta(
1376    context: &SubscriptionContext,
1377    subscription_id: &str,
1378    view_spec: &ViewSpec,
1379    current: &[(String, Value)],
1380    next: &[(String, Value)],
1381    envelope: &BusMessage,
1382    metadata: &SourceFrameMetadata,
1383) -> Result<()> {
1384    let current_keys: Vec<&str> = current.iter().map(|(key, _)| key.as_str()).collect();
1385    let next_keys: Vec<&str> = next.iter().map(|(key, _)| key.as_str()).collect();
1386    let next_set: HashSet<&str> = next_keys.iter().copied().collect();
1387
1388    for key in current_keys
1389        .iter()
1390        .copied()
1391        .filter(|key| !next_set.contains(key))
1392    {
1393        let op = if metadata.op == "delete" && key == envelope.key {
1394            "delete"
1395        } else {
1396            "remove"
1397        };
1398        send_membership_frame(
1399            context,
1400            subscription_id,
1401            view_spec,
1402            op,
1403            key,
1404            Value::Null,
1405            metadata.seq.clone(),
1406        )?;
1407    }
1408
1409    for (position, (key, data)) in next.iter().enumerate() {
1410        let previous_position = current_keys.iter().position(|candidate| *candidate == key);
1411        let changed_position = previous_position != Some(position);
1412        if previous_position.is_none() || changed_position || key == &envelope.key {
1413            let can_forward_patch = !view_spec.is_derived()
1414                && previous_position == Some(position)
1415                && key == &envelope.key
1416                && metadata.op != "delete";
1417            if can_forward_patch {
1418                send_scoped_source_payload(
1419                    context,
1420                    subscription_id,
1421                    &view_spec.id,
1422                    envelope.payload.clone(),
1423                )?;
1424            } else {
1425                let seq = metadata
1426                    .seq
1427                    .clone()
1428                    .or_else(|| data.get("_seq").and_then(Value::as_str).map(str::to_string));
1429                send_membership_frame(
1430                    context,
1431                    subscription_id,
1432                    view_spec,
1433                    "upsert",
1434                    key,
1435                    data.clone(),
1436                    seq,
1437                )?;
1438            }
1439        }
1440    }
1441    Ok(())
1442}
1443
1444fn to_wire_snapshot_entities(
1445    entities: Vec<(String, Value)>,
1446    view_spec: &ViewSpec,
1447) -> Vec<SnapshotEntity> {
1448    entities
1449        .into_iter()
1450        .map(|(key, mut data)| {
1451            apply_wire_format(&mut data, &view_spec.wire_format);
1452            SnapshotEntity { key, data }
1453        })
1454        .collect()
1455}
1456
1457async fn load_query_entities(
1458    entity_cache: &EntityCache,
1459    sorted_caches: Option<
1460        Arc<tokio::sync::RwLock<HashMap<String, crate::sorted_cache::SortedViewCache>>>,
1461    >,
1462    view_spec: &ViewSpec,
1463    query: &SubscriptionQuery,
1464    apply_snapshot_limit: bool,
1465) -> Vec<(String, Value)> {
1466    let (entities, preordered) = if let Some(sorted_caches) = sorted_caches {
1467        let mut caches = sorted_caches.write().await;
1468        let entities = caches
1469            .get_mut(&view_spec.id)
1470            .map(|cache| cache.get_all_ordered())
1471            .unwrap_or_default();
1472        (entities, true)
1473    } else if view_spec.mode == Mode::State {
1474        let entity = match query.key.as_deref() {
1475            Some(key) => entity_cache
1476                .get(&view_spec.id, key)
1477                .await
1478                .map(|data| vec![(key.to_string(), data)])
1479                .unwrap_or_default(),
1480            None => vec![],
1481        };
1482        (entity, true)
1483    } else {
1484        (entity_cache.get_all(&view_spec.id).await, false)
1485    };
1486    select_query_entities(entities, query, preordered, apply_snapshot_limit)
1487}
1488
1489fn select_query_entities(
1490    mut entities: Vec<(String, Value)>,
1491    query: &SubscriptionQuery,
1492    preordered: bool,
1493    apply_snapshot_limit: bool,
1494) -> Vec<(String, Value)> {
1495    entities.retain(|(key, data)| query_matches_entity(query, key, data));
1496    if !preordered {
1497        entities.sort_by(|left, right| {
1498            let left_seq = left.1.get("_seq").and_then(Value::as_str).unwrap_or("");
1499            let right_seq = right.1.get("_seq").and_then(Value::as_str).unwrap_or("");
1500            let order = if query.after.is_some() {
1501                cmp_seq(left_seq, right_seq)
1502            } else {
1503                cmp_seq(right_seq, left_seq)
1504            };
1505            order.then_with(|| left.0.cmp(&right.0))
1506        });
1507    }
1508
1509    let skip = query.skip.unwrap_or(0);
1510    let take = query.take.unwrap_or(usize::MAX);
1511    let mut selected: Vec<_> = entities.into_iter().skip(skip).take(take).collect();
1512    if apply_snapshot_limit {
1513        if let Some(limit) = query.snapshot_limit {
1514            selected.truncate(limit);
1515        }
1516    }
1517    selected
1518}
1519
1520fn query_matches_entity(query: &SubscriptionQuery, key: &str, data: &Value) -> bool {
1521    if !query.matches_key(key) {
1522        return false;
1523    }
1524    if let Some(partition) = &query.partition {
1525        if value_at_dot_path(data, "_partition") != Some(&Value::String(partition.clone())) {
1526            return false;
1527        }
1528    }
1529    if let Some(after) = &query.after {
1530        let Some(seq) = data.get("_seq").and_then(Value::as_str) else {
1531            return false;
1532        };
1533        if cmp_seq(seq, after) != std::cmp::Ordering::Greater {
1534            return false;
1535        }
1536    }
1537    query
1538        .filters
1539        .iter()
1540        .all(|(path, expected)| value_at_dot_path(data, path) == Some(expected))
1541}
1542
1543fn value_at_dot_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
1544    path.split('.')
1545        .try_fold(value, |current, segment| current.get(segment))
1546}
1547
1548#[cfg(test)]
1549mod tests {
1550    use super::*;
1551    use crate::cache::EntityCacheConfig;
1552    use crate::view::{Delivery, Filters, Projection};
1553    use serde_json::json;
1554    use tokio::sync::oneshot;
1555
1556    fn list_spec() -> ViewSpec {
1557        ViewSpec {
1558            id: "Thing/list".to_string(),
1559            export: "Thing".to_string(),
1560            mode: Mode::List,
1561            wire_format: Default::default(),
1562            projection: Projection::all(),
1563            filters: Filters::all(),
1564            delivery: Delivery::default(),
1565            pipeline: None,
1566            source_view: None,
1567        }
1568    }
1569
1570    #[test]
1571    fn snapshot_batches_share_identity_and_completion() {
1572        let entities = ["one", "two", "three"].map(|key| SnapshotEntity {
1573            key: key.to_string(),
1574            data: json!({"key": key}),
1575        });
1576        let batches = create_snapshot_batches(
1577            &entities,
1578            SnapshotMetadata {
1579                subscription_id: "sub-1",
1580                snapshot_id: "snapshot-1",
1581                authoritative: true,
1582                mode: Mode::List,
1583                view_id: "Thing/list",
1584                key: None,
1585            },
1586            &SnapshotBatchConfig {
1587                initial_batch_size: 2,
1588                subsequent_batch_size: 1,
1589            },
1590        );
1591        assert_eq!(batches.len(), 2);
1592        assert!(batches.iter().all(|batch| batch.subscription_id == "sub-1"));
1593        assert!(batches
1594            .iter()
1595            .all(|batch| batch.snapshot_id == "snapshot-1"));
1596        assert!(!batches[0].complete);
1597        assert!(batches[1].complete);
1598        assert!(batches.iter().all(|batch| batch.authoritative));
1599    }
1600
1601    #[test]
1602    fn empty_incremental_snapshot_is_explicitly_non_authoritative() {
1603        let batches = create_snapshot_batches(
1604            &[],
1605            SnapshotMetadata {
1606                subscription_id: "sub-1",
1607                snapshot_id: "snapshot-1",
1608                authoritative: false,
1609                mode: Mode::State,
1610                view_id: "Thing/state",
1611                key: Some("missing"),
1612            },
1613            &SnapshotBatchConfig {
1614                initial_batch_size: 1,
1615                subsequent_batch_size: 1,
1616            },
1617        );
1618        assert_eq!(batches.len(), 1);
1619        assert!(!batches[0].authoritative);
1620        assert!(batches[0].complete);
1621        assert_eq!(batches[0].key.as_deref(), Some("missing"));
1622    }
1623
1624    #[test]
1625    fn dot_path_filters_are_exact_and_type_sensitive() {
1626        let mut query = SubscriptionQuery {
1627            view: "Thing/list".to_string(),
1628            ..Default::default()
1629        };
1630        query
1631            .filters
1632            .insert("state.status".to_string(), json!("open"));
1633        query.filters.insert("metrics.count".to_string(), json!(2));
1634        assert!(query_matches_entity(
1635            &query,
1636            "one",
1637            &json!({"state": {"status": "open"}, "metrics": {"count": 2}}),
1638        ));
1639        assert!(!query_matches_entity(
1640            &query,
1641            "one",
1642            &json!({"state": {"status": "open"}, "metrics": {"count": "2"}}),
1643        ));
1644    }
1645
1646    #[test]
1647    fn take_and_skip_define_independent_deterministic_windows() {
1648        let entities: Vec<_> = (1..=6)
1649            .map(|id| {
1650                (
1651                    id.to_string(),
1652                    json!({"id": id, "_seq": format!("10:{id:012}")}),
1653                )
1654            })
1655            .collect();
1656        let first = SubscriptionQuery {
1657            view: "Thing/list".to_string(),
1658            take: Some(2),
1659            skip: Some(0),
1660            ..Default::default()
1661        };
1662        let second = SubscriptionQuery {
1663            skip: Some(2),
1664            ..first.clone()
1665        };
1666        let first_keys: Vec<_> = select_query_entities(entities.clone(), &first, false, false)
1667            .into_iter()
1668            .map(|(key, _)| key)
1669            .collect();
1670        let second_keys: Vec<_> = select_query_entities(entities, &second, false, false)
1671            .into_iter()
1672            .map(|(key, _)| key)
1673            .collect();
1674        assert_eq!(first_keys, ["6", "5"]);
1675        assert_eq!(second_keys, ["4", "3"]);
1676    }
1677
1678    #[tokio::test]
1679    async fn state_receiver_is_installed_before_snapshot_awaits() {
1680        let bus = BusManager::new();
1681        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
1682        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
1683        let bus_for_task = bus.clone();
1684        let task = tokio::spawn(async move {
1685            subscribe_state_then_snapshot(&bus_for_task, "Thing/state", "one", || async move {
1686                snapshot_started_tx.send(()).unwrap();
1687                release_snapshot_rx.await.unwrap();
1688            })
1689            .await
1690            .0
1691        });
1692        snapshot_started_rx.await.unwrap();
1693        bus.publish_state(
1694            "Thing/state",
1695            "one",
1696            Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
1697        )
1698        .await;
1699        release_snapshot_tx.send(()).unwrap();
1700        let mut receiver = task.await.unwrap();
1701        receiver.changed().await.unwrap();
1702        assert!(!receiver.borrow().is_empty());
1703    }
1704
1705    async fn assert_list_receiver_precedes_snapshot(view: &'static str) {
1706        let bus = BusManager::new();
1707        let (snapshot_started_tx, snapshot_started_rx) = oneshot::channel();
1708        let (release_snapshot_tx, release_snapshot_rx) = oneshot::channel();
1709        let bus_for_task = bus.clone();
1710        let task = tokio::spawn(async move {
1711            subscribe_list_then_snapshot(&bus_for_task, view, || async move {
1712                snapshot_started_tx.send(()).unwrap();
1713                release_snapshot_rx.await.unwrap();
1714            })
1715            .await
1716            .0
1717        });
1718        snapshot_started_rx.await.unwrap();
1719        bus.publish_list(
1720            view,
1721            Arc::new(BusMessage {
1722                key: "one".to_string(),
1723                entity: view.to_string(),
1724                payload: Arc::new(Bytes::from_static(br#"{"op":"patch"}"#)),
1725            }),
1726        )
1727        .await;
1728        release_snapshot_tx.send(()).unwrap();
1729        let mut receiver = task.await.unwrap();
1730        assert_eq!(receiver.recv().await.unwrap().key, "one");
1731    }
1732
1733    #[tokio::test]
1734    async fn list_receiver_is_installed_before_snapshot() {
1735        assert_list_receiver_precedes_snapshot("Thing/list").await;
1736    }
1737
1738    #[tokio::test]
1739    async fn append_receiver_is_installed_before_snapshot() {
1740        assert_list_receiver_precedes_snapshot("Thing/append").await;
1741    }
1742
1743    #[tokio::test]
1744    async fn derived_source_receiver_is_installed_before_snapshot() {
1745        assert_list_receiver_precedes_snapshot("Thing/list-source").await;
1746    }
1747
1748    #[tokio::test]
1749    async fn snapshot_limit_does_not_change_live_take_skip_membership() {
1750        let cache = EntityCache::with_config(EntityCacheConfig {
1751            max_entities_per_view: 10,
1752            ..Default::default()
1753        });
1754        for id in 1..=4 {
1755            cache
1756                .upsert(
1757                    "Thing/list",
1758                    &id.to_string(),
1759                    json!({"_seq": format!("10:{id:012}")}),
1760                )
1761                .await;
1762        }
1763        let query = SubscriptionQuery {
1764            view: "Thing/list".to_string(),
1765            take: Some(3),
1766            skip: Some(1),
1767            snapshot_limit: Some(1),
1768            ..Default::default()
1769        };
1770        let live = load_query_entities(&cache, None, &list_spec(), &query, false).await;
1771        let snapshot = load_query_entities(&cache, None, &list_spec(), &query, true).await;
1772        assert_eq!(live.len(), 3);
1773        assert_eq!(snapshot.len(), 1);
1774    }
1775
1776    #[test]
1777    fn fixture_manifest_covers_required_conformance_cases() {
1778        let manifest: Value = serde_json::from_str(include_str!(
1779            "../../../../tests/fixtures/websocket-v2/manifest.json"
1780        ))
1781        .unwrap();
1782        let names: HashSet<_> = manifest["fixtures"]
1783            .as_array()
1784            .unwrap()
1785            .iter()
1786            .filter_map(Value::as_str)
1787            .collect();
1788        for required in [
1789            "keyed-state.json",
1790            "list-windows.json",
1791            "filters.json",
1792            "multi-batch-authoritative.json",
1793            "empty-snapshot.json",
1794            "remove.json",
1795            "delete.json",
1796            "incremental-snapshot.json",
1797            "reconnect-replacement.json",
1798            "errors.json",
1799        ] {
1800            assert!(names.contains(required), "missing fixture {required}");
1801        }
1802
1803        for document in [
1804            include_str!("../../../../tests/fixtures/websocket-v2/keyed-state.json"),
1805            include_str!("../../../../tests/fixtures/websocket-v2/list-windows.json"),
1806            include_str!("../../../../tests/fixtures/websocket-v2/filters.json"),
1807            include_str!("../../../../tests/fixtures/websocket-v2/multi-batch-authoritative.json"),
1808            include_str!("../../../../tests/fixtures/websocket-v2/empty-snapshot.json"),
1809            include_str!("../../../../tests/fixtures/websocket-v2/remove.json"),
1810            include_str!("../../../../tests/fixtures/websocket-v2/delete.json"),
1811            include_str!("../../../../tests/fixtures/websocket-v2/incremental-snapshot.json"),
1812            include_str!("../../../../tests/fixtures/websocket-v2/reconnect-replacement.json"),
1813            include_str!("../../../../tests/fixtures/websocket-v2/errors.json"),
1814        ] {
1815            let fixture: Value = serde_json::from_str(document).unwrap();
1816            assert!(fixture["name"].is_string());
1817        }
1818    }
1819}