arcature 2026.2.1

//! Typed SSE event-stream wrapper over `axum::response::sse` (PROGRAM.md
//! §AP2.1-8).
//!
//! A thin typed layer that adds the realtime safety core the raw SSE
//! response does not provide: explicit origin policy, connection-limit
//! enforcement, backpressure, graceful shutdown, and stable tracing spans.
//! Raw `axum::response::sse` remains a first-class escape hatch
//! (AGENTS.md §16): an application can return a raw `Sse` directly and
//! skip this wrapper.
//!
//! # No proprietary protocol
//!
//! The wrapper emits standard SSE frames (`event:`, `data:`, `id:`,
//! `retry:`); it owns no `X-Arcature-*` header and no page protocol. A
//! stock browser `EventSource` works without knowing Arcature exists.
//!
//! # Backpressure
//!
//! `axum::response::sse` drives its source `Stream` on demand (it only
//! polls `poll_next` when the HTTP response body is being flushed). When
//! the client network is slow, the body flush blocks, so the source is
//! not polled — natural backpressure. The wrapper additionally maps
//! broadcast lag/closed to typed outcomes (a comment frame on lag, stream
//! end on closed) so the client is not silently starved.
//!
//! # Lifecycle / reconnect
//!
//! The server does not own client reconnect. `EventSource` reconnects
//! automatically; the wrapper emits a `retry:` interval the application
//! configures. The server owns admission, one live subscription per
//! stream, and a graceful stream-end (a final comment + EOF) on drain.

use std::pin::Pin;
use std::time::Duration;

use futures::stream::{self, Stream};

use crate::axum::http::HeaderMap;
use crate::axum::response::sse::{Event, KeepAlive, Sse};
use crate::axum::response::{IntoResponse, Response};
use crate::realtime::channel::{Broadcast, ChannelPayload, Subscription};
use crate::realtime::error::{ChannelError, ProtocolHint, RealtimeError, admission_status};
use crate::realtime::origin::{OriginDecision, OriginPolicy};
use crate::realtime::registry::{ConnectionGuard, Registry};
use crate::realtime::shutdown::ShutdownConfig;

/// SSE timing and retry limits (attacker-facing bounds, AGENTS.md §29).
#[derive(Clone, Copy, Debug)]
pub struct SseLimits {
    /// The `retry:` value sent to the client (ms between reconnects). The
    /// server advertises it; `EventSource` honors it.
    pub retry_ms: u64,
    /// Keep-alive comment interval (prevents idle proxies from closing the
    /// connection). Zero disables keep-alive.
    pub keep_alive_interval: Duration,
}

impl SseLimits {
    /// Conservative defaults: 3s retry, 15s keep-alive.
    #[must_use]
    pub fn conservative() -> Self {
        Self {
            retry_ms: 3000,
            keep_alive_interval: Duration::from_secs(15),
        }
    }
}

/// The configuration for an SSE endpoint. Like the WebSocket endpoint, the
/// application constructs one and clones it into every handler.
#[derive(Clone)]
pub struct SseEndpoint {
    broadcast: Broadcast,
    origin: OriginPolicy,
    registry: Registry,
    limits: SseLimits,
    shutdown: ShutdownConfig,
}

impl SseEndpoint {
    /// Construct an SSE endpoint. The `broadcast` is the channel the stream
    /// subscribes to; `origin`/`registry`/`shutdown` are shared with the
    /// WebSocket endpoint (the application passes the same values).
    #[must_use]
    pub fn new(
        broadcast: Broadcast,
        origin: OriginPolicy,
        registry: Registry,
        limits: SseLimits,
        shutdown: ShutdownConfig,
    ) -> Self {
        Self {
            broadcast,
            origin,
            registry,
            limits,
            shutdown,
        }
    }

    /// The Axum handler. Accept the request headers and an application-
    /// defined `channel_id` (the handler extracts it). Performs admission
    /// (origin, connection limit) and returns the SSE response.
    ///
    /// Authorization for SSE: because SSE is a plain GET (no upgrade), the
    /// application is expected to wire per-channel authorization via an
    /// Axum layer / route guard before this handler runs (e.g. a
    /// `Policy<M>` guard that rejects unauthorized channels with 403). The
    /// `broadcast` here is therefore the one the application already
    /// resolved after its own authorization — channel names never
    /// implicitly authorize (PROGRAM.md §AP2.1-8). The wrapper enforces
    /// origin and the connection cap; the application enforces channel
    /// authorization upstream.
    pub async fn handle(self, headers: HeaderMap, channel_id: String) -> Response {
        let _ = channel_id;

        // ── Admission: origin ──────────────────────────────────────────
        if self.origin.authorize(headers.get("origin")) == OriginDecision::Denied {
            tracing::debug!(
                target: "arcature::realtime::sse::admit",
                realtime_transport = "sse",
                error_category = "origin",
                "realtime sse rejected: origin policy"
            );
            return admission_status(&RealtimeError::Origin).into_response();
        }

        // ── Admission: connection limit ────────────────────────────────
        let guard = match self.registry.acquire(self.shutdown.max_connections()) {
            Ok(g) => g,
            Err(_) => {
                tracing::debug!(
                    target: "arcature::realtime::sse::admit",
                    realtime_transport = "sse",
                    error_category = "limit",
                    "realtime sse rejected: connection limit"
                );
                return admission_status(&RealtimeError::ConnectionLimit).into_response();
            }
        };

        // ── Build the typed event stream ────────────────────────────────
        // The stream owns the subscription and the guard; dropping the
        // stream (which happens when axum drops the response body — client
        // disconnect or stream end) drops both, so the live counts
        // decrement cleanly (no orphans).
        let sub = self.broadcast.subscribe();
        let shutdown = self.shutdown;
        let retry_ms = self.limits.retry_ms;

        let event_stream = sse_event_stream(sub, guard, shutdown, retry_ms);

        let keep_alive = if self.limits.keep_alive_interval.is_zero() {
            None
        } else {
            Some(
                KeepAlive::new()
                    .interval(self.limits.keep_alive_interval)
                    .text("arcature-keepalive"),
            )
        };
        let event_stream: Pin<
            Box<dyn Stream<Item = Result<Event, std::convert::Infallible>> + Send>,
        > = Box::pin(event_stream);
        // The keep-alive and no-keep-alive arms produce different concrete
        // `Sse<...>` types (keep-alive wraps in `KeepAliveStream`), so each
        // arm converts to `Response` directly rather than unifying the
        // `Sse` type. Boxing the event stream keeps the inner stream type
        // identical; the `match` is purely over which `into_response` path
        // runs.
        match keep_alive {
            Some(ka) => Sse::new(event_stream).keep_alive(ka).into_response(),
            None => Sse::new(event_stream).into_response(),
        }
    }
}

/// The per-stream state owned by the `unfold` closure. Moving the state
/// into the async closure avoids the self-referential borrow a boxed
/// `recv` future would create on a manual `Stream` impl.
struct SseState {
    sub: Subscription,
    _guard: ConnectionGuard,
    shutdown: ShutdownConfig,
    retry_ms: u64,
    // True until the retry event has been yielded.
    yielded_retry: bool,
    // True once the draining comment has been emitted; the next step
    // returns `None` to end the stream.
    draining: bool,
    // True once the channel closed comment has been emitted.
    channel_closed: bool,
}

/// Build the typed SSE event `Stream` from a broadcast subscription and a
/// connection guard. The first yielded event advertises the `retry:`
/// interval; subsequent events are broadcast payloads, with comment
/// frames for lag/closed, and a final comment + `None` on server drain.
fn sse_event_stream(
    sub: Subscription,
    guard: ConnectionGuard,
    shutdown: ShutdownConfig,
    retry_ms: u64,
) -> impl Stream<Item = Result<Event, std::convert::Infallible>> + Send {
    let init = SseState {
        sub,
        _guard: guard,
        shutdown,
        retry_ms,
        yielded_retry: false,
        draining: false,
        channel_closed: false,
    };
    stream::unfold(init, |mut state| async move {
        // First event: advertise the retry interval.
        if !state.yielded_retry {
            state.yielded_retry = true;
            return Some((
                Ok(Event::default().retry(Duration::from_millis(state.retry_ms))),
                state,
            ));
        }

        // Graceful drain: emit one final comment, then end on the next poll.
        if state.draining {
            return None;
        }
        if state.shutdown.is_draining() {
            state.draining = true;
            return Some((Ok(Event::default().comment("server-draining")), state));
        }

        // If the broadcast channel is closed, end the stream. A closed
        // `tokio::sync::broadcast` channel returns `RecvError::Closed`
        // immediately on every subsequent `recv()`, so without this guard
        // the stream would emit `channel-closed` comments forever in a
        // hot loop. The single `channel-closed` comment was already emitted
        // on the poll that first observed the close; now we terminate.
        if state.channel_closed {
            return None;
        }

        // Pull the next broadcast payload, or exit promptly on drain. The
        // `select!` against `drain_notified()` ensures a drain exits the
        // stream even when `sub.recv()` is blocked waiting for a broadcast
        // that never comes.
        tokio::select! {
            _ = state.shutdown.drain_notified() => {
                state.draining = true;
                Some((Ok(Event::default().comment("server-draining")), state))
            }
            recv = state.sub.recv() => {
                match recv {
                    Ok(payload) => match event_from_payload(payload) {
                        Ok(ev) => Some((Ok(ev), state)),
                        Err(hint) => {
                            Some((
                                Ok(Event::default().comment(format!("decode-error:{hint}"))),
                                state,
                            ))
                        }
                    },
                    Err(ChannelError::Lagged) => {
                        Some((Ok(Event::default().comment("channel-lagged")), state))
                    }
                    Err(ChannelError::Closed) => {
                        state.channel_closed = true;
                        Some((Ok(Event::default().comment("channel-closed")), state))
                    }
                    Err(ChannelError::Full) => None,
                }
            }
        }
    })
}

/// Build an SSE [`Event`] from a broadcast payload. The payload is the
/// application's serialized bytes (typically JSON). A non-UTF-8 payload
/// yields a typed protocol hint rather than a panic (AGENTS.md §17).
fn event_from_payload(payload: ChannelPayload) -> Result<Event, ProtocolHint> {
    let s = std::str::from_utf8(payload.as_bytes()).map_err(|_| ProtocolHint::Utf8)?;
    Ok(Event::default().data(s))
}