Skip to main content

reddb_server/server/
output_stream.rs

1//! Output streaming primitives for issue #760 (PRD #759 / ADR 0029) — the
2//! "shim slice" of bidirectional streaming.
3//!
4//! Scope of this module:
5//!   - [`StreamConfig`]: capture the `stream.*` namespace from `red_config`
6//!     at lease open and freeze it for the lease's lifetime (acceptance
7//!     criterion: KV mutations mid-stream do not retroactively change
8//!     behaviour).
9//!   - [`StreamLease`]: an internal, unforwarded handle bound to a
10//!     snapshot LSN and a frozen config. No external surface yet (S2 will
11//!     add quotas + per-principal accounting).
12//!   - [`open_stream`]: refuses with `stream_in_transaction_unsupported`
13//!     when the caller already has an active `BEGIN` on the session.
14//!   - [`ChunkProducer`]: page-aligned (N × 16 KiB) production buffer
15//!     that flushes on the first of byte / row / latency cap.
16//!   - [`Clock`]: trait-injected time source so TTL expiry is testable.
17//!
18//! The HTTP NDJSON wire framing built on top of these primitives lives in
19//! `handlers_query::handle_query_ndjson_stream` and is dispatched from
20//! `routing::try_route_streaming`.
21
22use std::collections::HashMap;
23use std::sync::atomic::{AtomicU64, Ordering};
24use std::sync::{Arc, Mutex};
25use std::time::{SystemTime, UNIX_EPOCH};
26
27use crate::runtime::RedDBRuntime;
28use crate::storage::query::engine::cancel::CancelToken;
29use reddb_types::types::Value;
30
31const RED_CONFIG_COLLECTION: &str = "red_config";
32
33/// Engine page size — the production buffer is always a multiple of this.
34/// Matches `storage::engine::PAGE_SIZE`.
35pub const PAGE_SIZE: usize = 16 * 1024;
36
37/// Injectable time source. Production code uses [`SystemClock`]; tests
38/// drive TTL expiry with [`FakeClock`] so they don't depend on wall time.
39pub trait Clock: Send + Sync {
40    fn now_ms(&self) -> u64;
41}
42
43#[derive(Debug, Default)]
44pub struct SystemClock;
45
46impl Clock for SystemClock {
47    fn now_ms(&self) -> u64 {
48        SystemTime::now()
49            .duration_since(UNIX_EPOCH)
50            .map(|d| d.as_millis() as u64)
51            .unwrap_or(0)
52    }
53}
54
55#[derive(Debug)]
56pub struct FakeClock {
57    now_ms: AtomicU64,
58}
59
60impl FakeClock {
61    pub fn new(now_ms: u64) -> Self {
62        Self {
63            now_ms: AtomicU64::new(now_ms),
64        }
65    }
66
67    pub fn advance(&self, ms: u64) {
68        self.now_ms.fetch_add(ms, Ordering::SeqCst);
69    }
70}
71
72impl Clock for FakeClock {
73    fn now_ms(&self) -> u64 {
74        self.now_ms.load(Ordering::SeqCst)
75    }
76}
77
78/// Frozen snapshot of the `stream.*` namespace at lease open. Acceptance
79/// criterion: a `red_config` mutation while a stream is running does not
80/// retroactively change the running stream's behaviour — that is, the
81/// per-lease config is value-typed and not a back-reference.
82#[derive(Debug, Clone, Copy)]
83pub struct StreamConfig {
84    pub snapshot_ttl_ms: u64,
85    pub chunk_default_pages: usize,
86    pub chunk_min_pages: usize,
87    pub chunk_max_pages: usize,
88    pub chunk_max_rows: usize,
89    pub chunk_max_latency_ms: u64,
90    /// Process-wide concurrent stream cap (issue #761 / S2).
91    /// Acquiring the (N+1)th slot is refused with
92    /// `server_stream_capacity_exhausted`.
93    pub max_global_streams: usize,
94    /// Per-principal concurrent stream cap (issue #761 / S2).
95    /// Acquiring the (M+1)th slot for a single principal is refused
96    /// with `principal_stream_quota_exhausted` even when the global
97    /// counter still has room.
98    pub max_per_principal_streams: usize,
99    /// Issue #765 / S6 — default end-to-end verification mode for input
100    /// streams when the open frame does not request one
101    /// (`stream.integrity.default_verify`, `"none"` per ADR 0029).
102    pub default_verify: crate::runtime::integrity_tombstone::VerifyMode,
103}
104
105impl Default for StreamConfig {
106    fn default() -> Self {
107        Self::DEFAULT
108    }
109}
110
111impl StreamConfig {
112    /// Defaults match ADR 0029. `snapshot_ttl_ms` is the only key
113    /// observable mid-stream; the chunk caps only affect framing.
114    pub const DEFAULT: StreamConfig = StreamConfig {
115        snapshot_ttl_ms: 60_000,
116        chunk_default_pages: 4,
117        chunk_min_pages: 1,
118        chunk_max_pages: 64,
119        chunk_max_rows: 1000,
120        chunk_max_latency_ms: 50,
121        max_global_streams: 256,
122        max_per_principal_streams: 32,
123        default_verify: crate::runtime::integrity_tombstone::VerifyMode::None,
124    };
125
126    /// Read every `stream.*` key from `red_config` over the runtime's KV
127    /// store. Missing keys fall back to [`Self::DEFAULT`]. Unparseable or
128    /// out-of-range values fall back silently — bad config never gets to
129    /// terminate a request that would otherwise succeed.
130    pub fn load(runtime: &RedDBRuntime) -> Self {
131        let db = runtime.db();
132        let read_u64 = |key: &str| -> Option<u64> {
133            match db.get_kv(RED_CONFIG_COLLECTION, key) {
134                Some((Value::Integer(v), _)) if v >= 0 => Some(v as u64),
135                Some((Value::UnsignedInteger(v), _)) => Some(v),
136                Some((Value::Text(text), _)) => text.parse().ok(),
137                _ => None,
138            }
139        };
140
141        let mut cfg = Self::DEFAULT;
142        if let Some(v) = read_u64("stream.snapshot.ttl_ms") {
143            cfg.snapshot_ttl_ms = v;
144        }
145        if let Some(v) = read_u64("stream.chunk.default_pages") {
146            cfg.chunk_default_pages = v as usize;
147        }
148        if let Some(v) = read_u64("stream.chunk.min_pages") {
149            cfg.chunk_min_pages = v as usize;
150        }
151        if let Some(v) = read_u64("stream.chunk.max_pages") {
152            cfg.chunk_max_pages = v as usize;
153        }
154        if let Some(v) = read_u64("stream.chunk.max_rows") {
155            cfg.chunk_max_rows = v as usize;
156        }
157        if let Some(v) = read_u64("stream.chunk.max_latency_ms") {
158            cfg.chunk_max_latency_ms = v;
159        }
160        if let Some(v) = read_u64("stream.max_global") {
161            cfg.max_global_streams = v as usize;
162        }
163        if let Some(v) = read_u64("stream.max_per_principal") {
164            cfg.max_per_principal_streams = v as usize;
165        }
166        // Issue #765 / S6 — opt-in integrity default. Stored as a Text key
167        // (`"none"` / `"sha256"`); anything else falls back to `None`.
168        if let Some((Value::Text(text), _)) =
169            db.get_kv(RED_CONFIG_COLLECTION, "stream.integrity.default_verify")
170        {
171            cfg.default_verify = crate::runtime::integrity_tombstone::VerifyMode::parse(&text);
172        }
173        cfg.normalize();
174        cfg
175    }
176
177    /// Clamp interrelated fields to a self-consistent state. Floors come
178    /// from ADR 0029 ("hard floor 1 page"); ceilings prevent zero-row /
179    /// zero-byte caps that would prevent the producer from ever flushing.
180    fn normalize(&mut self) {
181        if self.chunk_min_pages == 0 {
182            self.chunk_min_pages = 1;
183        }
184        if self.chunk_max_pages < self.chunk_min_pages {
185            self.chunk_max_pages = self.chunk_min_pages;
186        }
187        if self.chunk_default_pages < self.chunk_min_pages {
188            self.chunk_default_pages = self.chunk_min_pages;
189        }
190        if self.chunk_default_pages > self.chunk_max_pages {
191            self.chunk_default_pages = self.chunk_max_pages;
192        }
193        if self.chunk_max_rows == 0 {
194            self.chunk_max_rows = 1;
195        }
196        if self.max_global_streams == 0 {
197            self.max_global_streams = 1;
198        }
199        if self.max_per_principal_streams == 0 {
200            self.max_per_principal_streams = 1;
201        }
202    }
203
204    /// Page-aligned production buffer size in bytes. Acceptance criterion:
205    /// "production buffer is always N × 16 KiB".
206    pub fn production_buffer_bytes(&self) -> usize {
207        self.chunk_default_pages.saturating_mul(PAGE_SIZE)
208    }
209}
210
211/// Monotonic, process-local lease id. The id is internal — the bearer
212/// token still authenticates the open, the lease only identifies the
213/// stream for audit and termination (ADR 0029 "Authorization").
214///
215/// Issue #767 / S8 — the internal id is **never** forwarded on the
216/// wire. The client receives [`StreamLease::lease_handle`], a 128-bit
217/// opaque random value, in `OpenAck` / the `{"end": …}` envelope.
218static NEXT_LEASE_ID: AtomicU64 = AtomicU64::new(1);
219
220/// Issue #767 / S8 — wire-visible lease handle, 128 bits drawn from
221/// the OS CSPRNG. Hex-encoded so it survives JSON / header transport.
222/// Opacity property: a client cannot derive the internal lease id
223/// from the handle (the two are independent), and the handle is not
224/// sequential across opens.
225pub const LEASE_HANDLE_BYTES: usize = 16;
226
227/// Generate an opaque 128-bit lease handle (32-char hex). Per ADR 0029,
228/// the handle is the only identifier the wire sees; the internal
229/// `StreamLease::id` stays server-side.
230pub fn generate_lease_handle() -> String {
231    let mut bytes = [0u8; LEASE_HANDLE_BYTES];
232    if crate::crypto::os_random::fill_bytes(&mut bytes).is_err() {
233        // CSPRNG failure is exceedingly rare on supported platforms;
234        // fall back to a high-entropy mix so the open path never
235        // terminates a stream that would otherwise succeed. The
236        // fallback is *not* secure unguessable — production paths
237        // will hit /dev/urandom successfully.
238        let lo = NEXT_LEASE_ID.fetch_add(0, Ordering::SeqCst).to_le_bytes();
239        let now = crate::utils::now_unix_nanos().to_le_bytes();
240        bytes[..8].copy_from_slice(&lo);
241        bytes[8..].copy_from_slice(&now);
242    }
243    crate::utils::to_hex(&bytes)
244}
245
246/// Issue #767 / S8 — non-validating JWT `exp` claim extractor. Returns
247/// the expiry time in milliseconds since the epoch when `token` is a
248/// JWT carrying an integer `exp` claim, otherwise `None`.
249///
250/// The extractor does **not** validate the JWT signature, issuer, or
251/// audience — those gates run during `is_authorized` at OpenStream.
252/// This helper is consulted *after* the bearer has been accepted, to
253/// decide whether the lease will outlive the bearer credential and an
254/// audit event should be emitted accordingly.
255pub fn parse_jwt_exp_ms(token: &str) -> Option<u64> {
256    let parts: Vec<&str> = token.split('.').collect();
257    if parts.len() != 3 {
258        return None;
259    }
260    let payload = base64url_decode_padded(parts[1])?;
261    let v: crate::json::Value = crate::json::from_slice(&payload).ok()?;
262    let exp_secs = v.get("exp").and_then(|n| n.as_f64())? as i64;
263    if exp_secs <= 0 {
264        return None;
265    }
266    Some((exp_secs as u64).saturating_mul(1000))
267}
268
269fn base64url_decode_padded(input: &str) -> Option<Vec<u8>> {
270    let mut s = String::with_capacity(input.len() + 4);
271    for ch in input.chars() {
272        match ch {
273            '-' => s.push('+'),
274            '_' => s.push('/'),
275            _ => s.push(ch),
276        }
277    }
278    while !s.len().is_multiple_of(4) {
279        s.push('=');
280    }
281    reddb_wire::redwire::handshake::base64_std_decode(&s)
282}
283
284/// Issue #767 / S8 — wire-visible reason for `stream.closed`. The
285/// audit emit site decides the variant once the stream's exit shape
286/// has been observed.
287#[derive(Debug, Clone, Copy, PartialEq, Eq)]
288pub enum CloseReason {
289    Ok,
290    Cancelled,
291    Error,
292    SnapshotExpired,
293    CapacityRefused,
294    IntegrityFailed,
295}
296
297impl CloseReason {
298    pub fn as_str(&self) -> &'static str {
299        match self {
300            CloseReason::Ok => "ok",
301            CloseReason::Cancelled => "cancelled",
302            CloseReason::Error => "error",
303            CloseReason::SnapshotExpired => "snapshot_expired",
304            CloseReason::CapacityRefused => "capacity_refused",
305            CloseReason::IntegrityFailed => "integrity_failed",
306        }
307    }
308}
309
310#[derive(Debug)]
311pub struct StreamLease {
312    pub id: u64,
313    /// Issue #767 / S8 — opaque handle handed to the client. 32-char
314    /// hex (128 bits of CSPRNG output). The internal `id` is not
315    /// derivable from this value.
316    pub lease_handle: String,
317    pub snapshot_lsn: u64,
318    pub opened_at_ms: u64,
319    pub config: StreamConfig,
320}
321
322impl StreamLease {
323    /// `true` once `snapshot_ttl_ms` has elapsed since the lease was opened.
324    /// The shim slice materialises the result first, so expiry in practice
325    /// only fires for streams that take longer than `ttl_ms` to drain to
326    /// the client. The check is wired in so the wire envelope can carry
327    /// `snapshot_expired` exactly the way later slices (pull-based
328    /// executors) will need it.
329    pub fn snapshot_expired(&self, now_ms: u64) -> bool {
330        now_ms.saturating_sub(self.opened_at_ms) >= self.config.snapshot_ttl_ms
331    }
332}
333
334#[derive(Debug, PartialEq, Eq)]
335pub enum OpenStreamError {
336    /// Acceptance criterion #4: `OpenStream` against a session with an
337    /// active `BEGIN` is refused with this code. The wire shape carries
338    /// it back as `{"error": {"code": "stream_in_transaction_unsupported"}}`.
339    TransactionActive,
340}
341
342impl OpenStreamError {
343    pub fn code(&self) -> &'static str {
344        match self {
345            OpenStreamError::TransactionActive => "stream_in_transaction_unsupported",
346        }
347    }
348
349    pub fn message(&self) -> &'static str {
350        match self {
351            OpenStreamError::TransactionActive => {
352                "cannot open output stream while a transaction is active on this session"
353            }
354        }
355    }
356}
357
358/// Issue a lease. The caller is responsible for binding `snapshot_lsn`
359/// to the same MVCC view the underlying executor will read from; in the
360/// shim slice this is `runtime.cdc_current_lsn()` captured before
361/// `execute_query`.
362pub fn open_stream(
363    config: StreamConfig,
364    snapshot_lsn: u64,
365    in_transaction: bool,
366    clock: &dyn Clock,
367) -> Result<StreamLease, OpenStreamError> {
368    if in_transaction {
369        return Err(OpenStreamError::TransactionActive);
370    }
371    Ok(StreamLease {
372        id: NEXT_LEASE_ID.fetch_add(1, Ordering::SeqCst),
373        lease_handle: generate_lease_handle(),
374        snapshot_lsn,
375        opened_at_ms: clock.now_ms(),
376        config,
377    })
378}
379
380/// Page-aligned chunk producer. The producer accumulates byte-encoded
381/// rows in an N × 16 KiB buffer; on the first of byte / row / latency
382/// cap it forwards the buffer to the supplied flush closure, which the
383/// transport layer turns into a chunked-encoding frame.
384///
385/// The struct does not know about HTTP, NDJSON, or chunked transfer —
386/// it is wire-agnostic so the gRPC and RedWire paths can reuse it.
387pub struct ChunkProducer<'a> {
388    buf: Vec<u8>,
389    rows: usize,
390    window_started_ms: u64,
391    cap_bytes: usize,
392    cap_rows: usize,
393    cap_latency_ms: u64,
394    clock: &'a dyn Clock,
395    total_flushes: u64,
396    total_bytes: u64,
397    total_rows: u64,
398    last_flush_reason: Option<FlushReason>,
399}
400
401#[derive(Debug, Clone, Copy, PartialEq, Eq)]
402pub enum FlushReason {
403    Byte,
404    Row,
405    Latency,
406    Terminal,
407}
408
409impl<'a> ChunkProducer<'a> {
410    pub fn new(config: &StreamConfig, clock: &'a dyn Clock) -> Self {
411        let cap_bytes = config.production_buffer_bytes();
412        Self {
413            buf: Vec::with_capacity(cap_bytes),
414            rows: 0,
415            window_started_ms: clock.now_ms(),
416            cap_bytes,
417            cap_rows: config.chunk_max_rows,
418            cap_latency_ms: config.chunk_max_latency_ms,
419            clock,
420            total_flushes: 0,
421            total_bytes: 0,
422            total_rows: 0,
423            last_flush_reason: None,
424        }
425    }
426
427    /// Append one already-encoded line (NDJSON: bytes + `\n`). Returns
428    /// `true` if the append triggered a flush.
429    pub fn push_line<F>(&mut self, line: &[u8], flush: &mut F) -> std::io::Result<bool>
430    where
431        F: FnMut(&[u8]) -> std::io::Result<()>,
432    {
433        crate::runtime::query_budget::charge_output(line.len()).map_err(std::io::Error::other)?;
434        self.buf.extend_from_slice(line);
435        self.rows += 1;
436        self.total_rows += 1;
437
438        if self.buf.len() >= self.cap_bytes {
439            self.flush(flush, FlushReason::Byte)?;
440            return Ok(true);
441        }
442        if self.rows >= self.cap_rows {
443            self.flush(flush, FlushReason::Row)?;
444            return Ok(true);
445        }
446        let elapsed = self.clock.now_ms().saturating_sub(self.window_started_ms);
447        if elapsed >= self.cap_latency_ms {
448            self.flush(flush, FlushReason::Latency)?;
449            return Ok(true);
450        }
451        Ok(false)
452    }
453
454    /// Issue #768 / S9 — drive a pull-based line source into the
455    /// production buffer. The producer *pulls* one encoded line at a
456    /// time from `source` and routes it through [`push_line`], so the
457    /// resident working set is the page-aligned buffer plus the single
458    /// line currently in hand — never the full result set. Pair this
459    /// with the pull-based scan iterators
460    /// (`parallel_scan::parallel_scan_rows`,
461    /// `bitmap_scan::execute_bitmap_scan_stream`) whose records are
462    /// encoded lazily by `encode`.
463    ///
464    /// Returns the number of lines consumed. Flush caps (byte / row /
465    /// latency) fire mid-drain exactly as they would for hand-driven
466    /// [`push_line`] calls, so first-line latency stays bounded by
467    /// `chunk.max_latency_ms` regardless of how many lines the source
468    /// will ultimately yield.
469    ///
470    /// [`push_line`]: ChunkProducer::push_line
471    pub fn drive_lines<S, R, Enc, F>(
472        &mut self,
473        source: S,
474        mut encode: Enc,
475        flush: &mut F,
476    ) -> std::io::Result<u64>
477    where
478        S: IntoIterator<Item = R>,
479        Enc: FnMut(&R) -> Vec<u8>,
480        F: FnMut(&[u8]) -> std::io::Result<()>,
481    {
482        let mut count = 0u64;
483        for record in source {
484            let line = encode(&record);
485            self.push_line(&line, flush)?;
486            count += 1;
487        }
488        Ok(count)
489    }
490
491    /// Force-flush any buffered bytes — used after the final NDJSON line
492    /// (`{"end": …}`) to push the tail of the buffer before closing the
493    /// connection.
494    pub fn finish<F>(&mut self, flush: &mut F) -> std::io::Result<()>
495    where
496        F: FnMut(&[u8]) -> std::io::Result<()>,
497    {
498        if !self.buf.is_empty() {
499            self.flush(flush, FlushReason::Terminal)?;
500        }
501        Ok(())
502    }
503
504    fn flush<F>(&mut self, flush: &mut F, reason: FlushReason) -> std::io::Result<()>
505    where
506        F: FnMut(&[u8]) -> std::io::Result<()>,
507    {
508        flush(&self.buf)?;
509        self.total_bytes += self.buf.len() as u64;
510        self.total_flushes += 1;
511        self.last_flush_reason = Some(reason);
512        self.buf.clear();
513        self.rows = 0;
514        self.window_started_ms = self.clock.now_ms();
515        Ok(())
516    }
517
518    pub fn total_flushes(&self) -> u64 {
519        self.total_flushes
520    }
521    pub fn total_bytes(&self) -> u64 {
522        self.total_bytes
523    }
524    pub fn total_rows(&self) -> u64 {
525        self.total_rows
526    }
527    pub fn last_flush_reason(&self) -> Option<FlushReason> {
528        self.last_flush_reason
529    }
530}
531
532/// HTTP chunked transfer encoding helpers. We do not pull in `hyper` for
533/// the streaming path; the existing SSE handler also hand-rolls its own
534/// HTTP framing, and matching that style keeps the diff narrow.
535pub fn write_chunked_response_header<W: std::io::Write>(
536    writer: &mut W,
537    status: u16,
538    content_type: &str,
539) -> std::io::Result<()> {
540    // Streaming responses write their own head here instead of going
541    // through `HttpResponse::to_http_bytes`, so they must repeat the
542    // same permissive CORS posture — otherwise a browser that can call
543    // the JSON endpoints cross-origin would still be blocked on the
544    // NDJSON/SSE streaming routes (graph/vector results). API auth is
545    // by header/API key, never cookies, so wildcard origin is safe and
546    // Allow-Credentials is deliberately never sent.
547    let mut header = format!(
548        "HTTP/1.1 {} {}\r\nContent-Type: {}\r\nTransfer-Encoding: chunked\r\nCache-Control: no-cache\r\nConnection: close\r\n",
549        status,
550        crate::server::transport::status_text(status),
551        content_type,
552    );
553    // Same permissive CORS posture as the buffered path — sourced from
554    // the single `cors_header_pairs` choke point so streaming routes
555    // (graph/vector NDJSON results) are reachable cross-origin too.
556    for (name, value) in crate::server::transport::cors_header_pairs() {
557        header.push_str(name);
558        header.push_str(": ");
559        header.push_str(&value);
560        header.push_str("\r\n");
561    }
562    header.push_str("\r\n");
563    writer.write_all(header.as_bytes())?;
564    writer.flush()
565}
566
567/// Emit one HTTP chunk (`<hex-size>\r\n<bytes>\r\n`). A zero-length
568/// payload is silently dropped — the terminator chunk lives in
569/// [`write_chunked_terminator`].
570pub fn write_chunk<W: std::io::Write>(writer: &mut W, bytes: &[u8]) -> std::io::Result<()> {
571    if bytes.is_empty() {
572        return Ok(());
573    }
574    let size = format!("{:x}\r\n", bytes.len());
575    writer.write_all(size.as_bytes())?;
576    writer.write_all(bytes)?;
577    writer.write_all(b"\r\n")?;
578    writer.flush()
579}
580
581/// Final `0\r\n\r\n` chunk that terminates a chunked body.
582pub fn write_chunked_terminator<W: std::io::Write>(writer: &mut W) -> std::io::Result<()> {
583    writer.write_all(b"0\r\n\r\n")?;
584    writer.flush()
585}
586
587/// Issue #761 / S2 — process-wide stream capacity registry. Holds two
588/// counters: a global concurrent-stream count and a per-principal map.
589/// Both are decremented when the [`StreamCapacityGuard`] handed back
590/// from a successful `try_acquire` is dropped, so the release path
591/// covers every normal exit (success, mid-stream error, snapshot
592/// expiry, client disconnect that drops the writer chain, panic
593/// unwind through the stack frame holding the guard).
594#[derive(Debug, Default)]
595pub struct StreamCapacityRegistry {
596    inner: Mutex<CapacityInner>,
597}
598
599#[derive(Debug, Default)]
600struct CapacityInner {
601    global_count: usize,
602    per_principal: HashMap<String, usize>,
603}
604
605/// Failure modes of [`StreamCapacityRegistry::try_acquire`]. Each
606/// variant carries the cap that fired and the live counter value at
607/// refusal time so clients can back off intelligently (the HTTP layer
608/// surfaces these inside the structured 429 body).
609#[derive(Debug, Clone, PartialEq, Eq)]
610pub enum AcquireError {
611    /// `stream.max_global` exceeded. Per acceptance criterion #1.
612    GlobalExhausted { limit: usize, current: usize },
613    /// `stream.max_per_principal` exceeded for `principal`. Per
614    /// acceptance criterion #2. The principal is surfaced verbatim;
615    /// callers must escape it on the wire.
616    PrincipalExhausted {
617        principal: String,
618        limit: usize,
619        current: usize,
620    },
621}
622
623impl AcquireError {
624    pub fn code(&self) -> &'static str {
625        match self {
626            AcquireError::GlobalExhausted { .. } => "server_stream_capacity_exhausted",
627            AcquireError::PrincipalExhausted { .. } => "principal_stream_quota_exhausted",
628        }
629    }
630
631    pub fn message(&self) -> String {
632        match self {
633            AcquireError::GlobalExhausted { limit, current } => {
634                format!("server stream capacity exhausted (limit {limit}, current {current})")
635            }
636            AcquireError::PrincipalExhausted {
637                principal,
638                limit,
639                current,
640            } => format!(
641                "principal {principal} stream quota exhausted (limit {limit}, current {current})"
642            ),
643        }
644    }
645}
646
647impl StreamCapacityRegistry {
648    pub fn new() -> Arc<Self> {
649        Arc::new(Self::default())
650    }
651
652    /// Attempt to acquire one slot. Both caps are checked under the
653    /// same lock, so a concurrent acquire+release pair cannot over-
654    /// issue beyond either ceiling (acceptance criterion #5).
655    pub fn try_acquire(
656        self: &Arc<Self>,
657        principal: &str,
658        max_global: usize,
659        max_per_principal: usize,
660    ) -> Result<StreamCapacityGuard, AcquireError> {
661        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
662        if inner.global_count >= max_global {
663            return Err(AcquireError::GlobalExhausted {
664                limit: max_global,
665                current: inner.global_count,
666            });
667        }
668        let current = inner.per_principal.get(principal).copied().unwrap_or(0);
669        if current >= max_per_principal {
670            return Err(AcquireError::PrincipalExhausted {
671                principal: principal.to_string(),
672                limit: max_per_principal,
673                current,
674            });
675        }
676        inner.global_count += 1;
677        inner
678            .per_principal
679            .insert(principal.to_string(), current + 1);
680        Ok(StreamCapacityGuard {
681            registry: Arc::clone(self),
682            principal: principal.to_string(),
683            released: false,
684        })
685    }
686
687    fn release(&self, principal: &str) {
688        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
689        if inner.global_count > 0 {
690            inner.global_count -= 1;
691        }
692        if let Some(count) = inner.per_principal.get_mut(principal) {
693            if *count > 0 {
694                *count -= 1;
695            }
696            if *count == 0 {
697                inner.per_principal.remove(principal);
698            }
699        }
700    }
701
702    /// Visible for tests and audit handlers.
703    pub fn snapshot(&self) -> (usize, HashMap<String, usize>) {
704        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
705        (inner.global_count, inner.per_principal.clone())
706    }
707}
708
709/// RAII slot returned by [`StreamCapacityRegistry::try_acquire`].
710/// Decrements both counters on drop — acceptance criterion #4
711/// ("releasing a slot on stream end (any reason) decrements both
712/// counters atomically").
713#[must_use = "dropping the guard immediately releases the stream slot"]
714#[derive(Debug)]
715pub struct StreamCapacityGuard {
716    registry: Arc<StreamCapacityRegistry>,
717    principal: String,
718    released: bool,
719}
720
721impl StreamCapacityGuard {
722    pub fn principal(&self) -> &str {
723        &self.principal
724    }
725}
726
727impl Drop for StreamCapacityGuard {
728    fn drop(&mut self) {
729        if !self.released {
730            self.registry.release(&self.principal);
731            self.released = true;
732        }
733    }
734}
735
736// ──────── Issue #766 / S7 — resume coordinator ────────
737
738/// Resumability assessment of a query plan. The shim slice runs a
739/// textual classifier over the SQL string: a query is resumable iff
740/// it has a stable total order over a unique key. By default we
741/// promise RID ASC; an explicit `ORDER BY rid` (or `ORDER BY rid ASC`)
742/// is also resumable. Anything that aggregates / groups / windows or
743/// orders on a non-unique column is not.
744pub fn assess_resumability(query: &str) -> bool {
745    let upper = query.to_uppercase();
746    let trimmed = upper.trim_start();
747    if !trimmed.starts_with("SELECT ") && !trimmed.starts_with("SELECT\n") {
748        return false;
749    }
750    const FORBIDDEN: &[&str] = &[
751        " GROUP BY ",
752        " HAVING ",
753        " DISTINCT ",
754        "DISTINCT ",
755        "COUNT(",
756        "SUM(",
757        "AVG(",
758        "MIN(",
759        "MAX(",
760        "ARRAY_AGG(",
761        "JSON_AGG(",
762        "OVER(",
763        " OVER (",
764        " JOIN ",
765    ];
766    for kw in FORBIDDEN {
767        if upper.contains(kw) {
768            return false;
769        }
770    }
771    if let Some(idx) = upper.find("ORDER BY") {
772        let tail = &upper[idx + "ORDER BY".len()..];
773        // Strip trailing LIMIT and statement terminator.
774        let mut clause = tail.to_string();
775        if let Some(lim) = clause.find(" LIMIT ") {
776            clause.truncate(lim);
777        }
778        if let Some(semi) = clause.find(';') {
779            clause.truncate(semi);
780        }
781        let clause = clause.trim();
782        if !matches!(clause, "RID" | "RID ASC") {
783            return false;
784        }
785    }
786    true
787}
788
789/// Resume-eligibility ledger. Holds `(snapshot_lsn → opened_at_ms,
790/// ttl_ms)` so a resume request can be checked against TTL without
791/// trusting the wall clock on the client. The shim slice does not
792/// implement true MVCC pinning — the registry's role is to make
793/// `snapshot_expired` deterministic and testable.
794#[derive(Debug, Default)]
795pub struct LeaseRegistry {
796    inner: Mutex<HashMap<u64, LeaseEntry>>,
797}
798
799#[derive(Debug, Clone, Copy)]
800struct LeaseEntry {
801    opened_at_ms: u64,
802    ttl_ms: u64,
803}
804
805#[derive(Debug, Clone, Copy, PartialEq, Eq)]
806pub enum LeaseLookup {
807    /// No lease ever recorded for this snapshot LSN.
808    Unknown,
809    /// Lease exists but its TTL has elapsed.
810    Expired,
811    /// Lease exists and is still within TTL.
812    Live,
813}
814
815impl LeaseRegistry {
816    pub fn new() -> Arc<Self> {
817        Arc::new(Self::default())
818    }
819
820    /// Record a freshly-opened lease. Idempotent — re-inserting the
821    /// same snapshot_lsn refreshes the timestamp (the client cannot
822    /// observe lease identity through the snapshot LSN alone, so this
823    /// matches "the latest open wins" semantics).
824    pub fn record(&self, snapshot_lsn: u64, opened_at_ms: u64, ttl_ms: u64) {
825        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
826        inner.insert(
827            snapshot_lsn,
828            LeaseEntry {
829                opened_at_ms,
830                ttl_ms,
831            },
832        );
833    }
834
835    /// Resume-time lookup. Returns whether the lease is unknown,
836    /// expired, or still live as of `now_ms`.
837    pub fn lookup(&self, snapshot_lsn: u64, now_ms: u64) -> LeaseLookup {
838        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
839        match inner.get(&snapshot_lsn) {
840            None => LeaseLookup::Unknown,
841            Some(entry) => {
842                if now_ms.saturating_sub(entry.opened_at_ms) >= entry.ttl_ms {
843                    LeaseLookup::Expired
844                } else {
845                    LeaseLookup::Live
846                }
847            }
848        }
849    }
850
851    /// Visible for tests / audit.
852    #[doc(hidden)]
853    pub fn len(&self) -> usize {
854        self.inner.lock().unwrap_or_else(|e| e.into_inner()).len()
855    }
856}
857
858/// Issue #807 / PRD #750 — opaque cursor token. 192 bits drawn from the
859/// OS CSPRNG, hex-encoded so it survives JSON transport. The token is the
860/// only cursor identifier the wire sees; the pinned snapshot, scope, and
861/// query stay server-side in the [`CursorRegistry`]. Opacity property: a
862/// client cannot derive the pinned snapshot LSN or another principal's
863/// token from the value, and tokens are not sequential across opens.
864pub const CURSOR_TOKEN_BYTES: usize = 24;
865
866static NEXT_CURSOR_SALT: AtomicU64 = AtomicU64::new(1);
867
868/// Generate an opaque 192-bit cursor token (48-char hex). Falls back to a
869/// high-entropy time/counter mix if the CSPRNG is unavailable so a fresh
870/// open never fails to mint a token (the fallback is not unguessable, but
871/// production paths reach `/dev/urandom` successfully).
872pub fn generate_cursor_token() -> String {
873    let mut bytes = [0u8; CURSOR_TOKEN_BYTES];
874    if crate::crypto::os_random::fill_bytes(&mut bytes).is_err() {
875        let salt = NEXT_CURSOR_SALT
876            .fetch_add(1, Ordering::SeqCst)
877            .to_le_bytes();
878        let now = crate::utils::now_unix_nanos().to_le_bytes();
879        bytes[..8].copy_from_slice(&salt);
880        bytes[8..16].copy_from_slice(&now);
881    }
882    crate::utils::to_hex(&bytes)
883}
884
885/// Server-side cursor record. Pins the read snapshot and the query so a
886/// resume replays the same point-in-time view, and binds the entry to the
887/// `(tenant, principal)` that opened it for authorization scoping.
888#[derive(Debug, Clone)]
889struct CursorEntry {
890    snapshot_lsn: u64,
891    tenant: String,
892    principal: String,
893    query: String,
894    entity_types: Option<Vec<String>>,
895    capabilities: Option<Vec<String>>,
896    opened_at_ms: u64,
897    ttl_ms: u64,
898    /// Issue #808 / 750d — shared cancel token observed by the executor for
899    /// this stream. The cancel endpoint and client-disconnect detection
900    /// raise it; the engine iterators poll it. Cloned out to the handler at
901    /// open so the live stream and a later cancel request share one flag.
902    cancel_token: CancelToken,
903    /// Issue #808 / 750d — tombstone marker. Set once the stream is
904    /// cancelled (explicitly or by disconnect); a tombstoned cursor can
905    /// never be resumed, only reported as cleanly cancelled to its owner.
906    cancelled: bool,
907}
908
909/// Pinned read state returned to the handler when a resume token resolves
910/// cleanly. The handler re-executes `query` against `snapshot_lsn` to
911/// re-stream the same view the original cursor referenced.
912#[derive(Debug, Clone, PartialEq, Eq)]
913pub struct CursorResume {
914    pub snapshot_lsn: u64,
915    pub query: String,
916    pub entity_types: Option<Vec<String>>,
917    pub capabilities: Option<Vec<String>>,
918    pub expires_at_ms: u64,
919}
920
921/// Why a resume token was refused. The handler maps every variant to a
922/// structured wire error. [`CursorReject::NotFound`] deliberately covers
923/// both "never issued" and "owned by a different tenant/principal" so an
924/// unauthorized caller cannot tell a foreign cursor from a nonexistent one.
925#[derive(Debug, Clone, Copy, PartialEq, Eq)]
926pub enum CursorReject {
927    /// Token unknown, or known to a different tenant/principal (masked so
928    /// existence never leaks to an unauthorized caller).
929    NotFound,
930    /// Token owned by this caller but its TTL has elapsed.
931    Expired,
932    /// Issue #808 / 750d — token owned by this caller but the stream was
933    /// cancelled (explicitly or by client disconnect) and tombstoned. A
934    /// resume is refused with a clean, structured error rather than
935    /// replaying a stream the client already abandoned.
936    Cancelled,
937}
938
939/// Issue #807 / PRD #750 — process-wide cursor registry for `/query/stream`.
940///
941/// On a fresh stream the server mints an opaque token, pins the read
942/// snapshot LSN, and records it here scoped to the requesting
943/// `(tenant, principal)` with a TTL. A later resume request presents the
944/// token; [`resolve`](CursorRegistry::resolve) returns the pinned query +
945/// snapshot only when the token is live AND the caller's scope matches.
946///
947/// In-memory is acceptable for this first cut (per the issue): the registry
948/// is shared via `Arc` so cloned server handles enforce against one map.
949#[derive(Debug, Default)]
950pub struct CursorRegistry {
951    inner: Mutex<HashMap<String, CursorEntry>>,
952}
953
954impl CursorRegistry {
955    pub fn new() -> Arc<Self> {
956        Arc::new(Self::default())
957    }
958
959    /// Mint + record a cursor for a freshly-opened stream, returning the
960    /// opaque token handed to the client. The entry expires at
961    /// `opened_at_ms + ttl_ms` (saturating).
962    #[allow(clippy::too_many_arguments)]
963    pub fn register(
964        &self,
965        snapshot_lsn: u64,
966        tenant: &str,
967        principal: &str,
968        query: &str,
969        entity_types: Option<Vec<String>>,
970        capabilities: Option<Vec<String>>,
971        opened_at_ms: u64,
972        ttl_ms: u64,
973    ) -> String {
974        let token = generate_cursor_token();
975        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
976        inner.insert(
977            token.clone(),
978            CursorEntry {
979                snapshot_lsn,
980                tenant: tenant.to_string(),
981                principal: principal.to_string(),
982                query: query.to_string(),
983                entity_types,
984                capabilities,
985                opened_at_ms,
986                ttl_ms,
987                cancel_token: CancelToken::new(),
988                cancelled: false,
989            },
990        );
991        token
992    }
993
994    /// Issue #808 / 750d — clone out the shared cancel token for a cursor
995    /// the caller just minted, so the live streaming handler can poll the
996    /// same flag a later cancel request raises. No scope check: the only
997    /// caller is the handler that owns the freshly-registered token.
998    pub fn cancel_token_for(&self, token: &str) -> Option<CancelToken> {
999        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
1000        inner.get(token).map(|e| e.cancel_token.clone())
1001    }
1002
1003    /// Issue #808 / 750d — cancel + tombstone a cursor on behalf of
1004    /// `(tenant, principal)`. Raises the shared cancel token (so an
1005    /// in-flight stream observing it stops) and marks the entry cancelled
1006    /// so a later resume is refused with [`CursorReject::Cancelled`].
1007    ///
1008    /// Scope is checked exactly as in [`resolve`](CursorRegistry::resolve):
1009    /// a token owned by a different tenant/principal — or one that never
1010    /// existed — returns [`CursorReject::NotFound`] so an unauthorized
1011    /// caller cannot cancel, or even confirm the existence of, a foreign
1012    /// stream. Cancelling an already-cancelled cursor is idempotent.
1013    pub fn cancel(
1014        &self,
1015        token: &str,
1016        tenant: &str,
1017        principal: &str,
1018    ) -> Result<CancelToken, CursorReject> {
1019        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
1020        let entry = inner.get_mut(token).ok_or(CursorReject::NotFound)?;
1021        if entry.tenant != tenant || entry.principal != principal {
1022            return Err(CursorReject::NotFound);
1023        }
1024        entry.cancelled = true;
1025        entry.cancel_token.cancel();
1026        Ok(entry.cancel_token.clone())
1027    }
1028
1029    /// Resolve a resume token for `(tenant, principal)` as of `now_ms`.
1030    ///
1031    /// Scope is checked **before** expiry: a token owned by a different
1032    /// tenant or principal resolves to [`CursorReject::NotFound`] —
1033    /// identical to a token that was never issued — so an unauthorized
1034    /// caller learns nothing about its existence. Only the rightful owner
1035    /// can observe [`CursorReject::Expired`].
1036    pub fn resolve(
1037        &self,
1038        token: &str,
1039        tenant: &str,
1040        principal: &str,
1041        now_ms: u64,
1042    ) -> Result<CursorResume, CursorReject> {
1043        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
1044        let entry = inner.get(token).ok_or(CursorReject::NotFound)?;
1045        // Scope gate first — a mismatch must be indistinguishable from
1046        // absence so a cross-tenant/cross-principal probe cannot confirm
1047        // the token exists.
1048        if entry.tenant != tenant || entry.principal != principal {
1049            return Err(CursorReject::NotFound);
1050        }
1051        // Issue #808 / 750d — a tombstoned (cancelled) cursor is refused
1052        // with a dedicated reason before the TTL check, so the owner learns
1053        // the stream was cancelled rather than seeing a generic expiry.
1054        if entry.cancelled {
1055            return Err(CursorReject::Cancelled);
1056        }
1057        if now_ms.saturating_sub(entry.opened_at_ms) >= entry.ttl_ms {
1058            return Err(CursorReject::Expired);
1059        }
1060        Ok(CursorResume {
1061            snapshot_lsn: entry.snapshot_lsn,
1062            query: entry.query.clone(),
1063            entity_types: entry.entity_types.clone(),
1064            capabilities: entry.capabilities.clone(),
1065            expires_at_ms: entry.opened_at_ms.saturating_add(entry.ttl_ms),
1066        })
1067    }
1068
1069    /// Visible for tests / audit.
1070    #[doc(hidden)]
1071    pub fn len(&self) -> usize {
1072        self.inner.lock().unwrap_or_else(|e| e.into_inner()).len()
1073    }
1074
1075    #[doc(hidden)]
1076    pub fn is_empty(&self) -> bool {
1077        self.len() == 0
1078    }
1079}
1080
1081/// Incremental SHA-256 hasher over emitted row lines, hex-encoded on
1082/// finalize. The wire contract is that the server hashes the exact
1083/// byte sequence of each row line (without trailing newline) in the
1084/// order the rows are emitted; the client stores the resulting digest
1085/// from the `end` envelope and replays it on resume.
1086#[derive(Debug, Default)]
1087pub struct PrefixHasher {
1088    inner: Option<sha2::Sha256>,
1089    rows: u64,
1090}
1091
1092impl PrefixHasher {
1093    pub fn new() -> Self {
1094        use sha2::Digest;
1095        Self {
1096            inner: Some(sha2::Sha256::new()),
1097            rows: 0,
1098        }
1099    }
1100
1101    pub fn update(&mut self, line: &[u8]) {
1102        use sha2::Digest;
1103        if let Some(h) = self.inner.as_mut() {
1104            h.update(line);
1105        }
1106        self.rows += 1;
1107    }
1108
1109    pub fn rows(&self) -> u64 {
1110        self.rows
1111    }
1112
1113    /// Hex-encoded digest of everything fed so far. Consumes the
1114    /// hasher (a `PrefixHasher` is single-use).
1115    pub fn finalize_hex(mut self) -> String {
1116        use sha2::Digest;
1117        let hasher = self
1118            .inner
1119            .take()
1120            .expect("PrefixHasher::finalize_hex called twice");
1121        let digest = hasher.finalize();
1122        let mut out = String::with_capacity(64);
1123        for b in digest.iter() {
1124            out.push_str(&format!("{b:02x}"));
1125        }
1126        out
1127    }
1128}
1129
1130/// Issue #767 / S8 — audit emission helpers. Each helper builds an
1131/// `AuditEvent` shaped to the brief and forwards it to the runtime's
1132/// audit log. The helpers are intentionally side-effect-only (no
1133/// return); audit emission must never terminate a stream that would
1134/// otherwise succeed.
1135pub fn audit_stream_opened(
1136    runtime: &RedDBRuntime,
1137    lease_handle: &str,
1138    principal: &str,
1139    snapshot_lsn: u64,
1140    query_hash: &str,
1141) {
1142    use crate::json::{Map, Value as JsonValue};
1143    let mut detail = Map::new();
1144    detail.insert(
1145        "lease_handle".to_string(),
1146        JsonValue::String(lease_handle.to_string()),
1147    );
1148    detail.insert(
1149        "snapshot_lsn".to_string(),
1150        JsonValue::Number(snapshot_lsn as f64),
1151    );
1152    detail.insert(
1153        "query_hash".to_string(),
1154        JsonValue::String(query_hash.to_string()),
1155    );
1156    let event = crate::runtime::audit_log::AuditEvent::builder("stream.opened")
1157        .principal(principal)
1158        .resource(lease_handle.to_string())
1159        .outcome(crate::runtime::audit_log::Outcome::Success)
1160        .detail(JsonValue::Object(detail))
1161        .build();
1162    runtime.audit_log().record_event(event);
1163}
1164
1165pub fn audit_stream_closed(
1166    runtime: &RedDBRuntime,
1167    lease_handle: &str,
1168    principal: &str,
1169    reason: CloseReason,
1170    row_count: u64,
1171    bytes_written: u64,
1172) {
1173    use crate::json::{Map, Value as JsonValue};
1174    let mut stats = Map::new();
1175    stats.insert("row_count".to_string(), JsonValue::Number(row_count as f64));
1176    stats.insert(
1177        "bytes_written".to_string(),
1178        JsonValue::Number(bytes_written as f64),
1179    );
1180    let mut detail = Map::new();
1181    detail.insert(
1182        "lease_handle".to_string(),
1183        JsonValue::String(lease_handle.to_string()),
1184    );
1185    detail.insert(
1186        "reason".to_string(),
1187        JsonValue::String(reason.as_str().to_string()),
1188    );
1189    detail.insert("stats".to_string(), JsonValue::Object(stats));
1190    let outcome = match reason {
1191        CloseReason::Ok => crate::runtime::audit_log::Outcome::Success,
1192        CloseReason::CapacityRefused => crate::runtime::audit_log::Outcome::Denied,
1193        _ => crate::runtime::audit_log::Outcome::Error,
1194    };
1195    let event = crate::runtime::audit_log::AuditEvent::builder("stream.closed")
1196        .principal(principal)
1197        .resource(lease_handle.to_string())
1198        .outcome(outcome)
1199        .detail(JsonValue::Object(detail))
1200        .build();
1201    runtime.audit_log().record_event(event);
1202}
1203
1204pub fn audit_token_expired_during_lease(
1205    runtime: &RedDBRuntime,
1206    lease_handle: &str,
1207    principal: &str,
1208    token_expiry_ms: u64,
1209) {
1210    use crate::json::{Map, Value as JsonValue};
1211    let mut detail = Map::new();
1212    detail.insert(
1213        "lease_handle".to_string(),
1214        JsonValue::String(lease_handle.to_string()),
1215    );
1216    detail.insert(
1217        "token_expiry".to_string(),
1218        JsonValue::Number(token_expiry_ms as f64),
1219    );
1220    detail.insert("lease_continued".to_string(), JsonValue::Bool(true));
1221    let event = crate::runtime::audit_log::AuditEvent::builder("stream.token_expired_during_lease")
1222        .principal(principal)
1223        .resource(lease_handle.to_string())
1224        .outcome(crate::runtime::audit_log::Outcome::Success)
1225        .detail(JsonValue::Object(detail))
1226        .build();
1227    runtime.audit_log().record_event(event);
1228}
1229
1230/// Capacity refusal emits a `stream.closed` event with no lease handle
1231/// (the open never produced one) and `reason: capacity_refused`. The
1232/// brief lists capacity_refused alongside the other close reasons; we
1233/// keep the wire shape identical so downstream audit-log consumers
1234/// don't have to special-case it.
1235pub fn audit_stream_capacity_refused(
1236    runtime: &RedDBRuntime,
1237    principal: &str,
1238    code: &str,
1239    limit: usize,
1240    current: usize,
1241) {
1242    use crate::json::{Map, Value as JsonValue};
1243    let mut detail = Map::new();
1244    detail.insert(
1245        "reason".to_string(),
1246        JsonValue::String(CloseReason::CapacityRefused.as_str().to_string()),
1247    );
1248    detail.insert("code".to_string(), JsonValue::String(code.to_string()));
1249    detail.insert("limit".to_string(), JsonValue::Number(limit as f64));
1250    detail.insert("current".to_string(), JsonValue::Number(current as f64));
1251    let event = crate::runtime::audit_log::AuditEvent::builder("stream.closed")
1252        .principal(principal)
1253        .outcome(crate::runtime::audit_log::Outcome::Denied)
1254        .detail(JsonValue::Object(detail))
1255        .build();
1256    runtime.audit_log().record_event(event);
1257}
1258
1259/// Borrow the lease registry's clock by default — the routing handler
1260/// uses this, tests inject their own.
1261pub fn system_clock() -> Arc<dyn Clock> {
1262    static INSTANCE: std::sync::OnceLock<Arc<dyn Clock>> = std::sync::OnceLock::new();
1263    Arc::clone(INSTANCE.get_or_init(|| Arc::new(SystemClock)))
1264}
1265
1266#[cfg(test)]
1267mod tests {
1268    use super::*;
1269
1270    #[test]
1271    fn open_stream_refuses_when_session_has_active_transaction() {
1272        let clock = FakeClock::new(0);
1273        let err = open_stream(StreamConfig::DEFAULT, 42, true, &clock).unwrap_err();
1274        assert_eq!(err, OpenStreamError::TransactionActive);
1275        assert_eq!(err.code(), "stream_in_transaction_unsupported");
1276    }
1277
1278    #[test]
1279    fn open_stream_succeeds_when_session_is_autocommit() {
1280        let clock = FakeClock::new(1_700_000_000_000);
1281        let lease = open_stream(StreamConfig::DEFAULT, 123, false, &clock).unwrap();
1282        assert_eq!(lease.snapshot_lsn, 123);
1283        assert_eq!(lease.opened_at_ms, 1_700_000_000_000);
1284        assert!(lease.id >= 1);
1285    }
1286
1287    #[test]
1288    fn lease_ids_are_unique_and_monotonic() {
1289        let clock = FakeClock::new(0);
1290        let a = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1291        let b = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1292        assert!(b.id > a.id);
1293    }
1294
1295    #[test]
1296    fn snapshot_expired_uses_injected_clock_and_ttl() {
1297        // TTL fake-clock test: advance past ttl_ms and snapshot_expired
1298        // flips. Acceptance criterion #5.
1299        let clock = FakeClock::new(0);
1300        let mut config = StreamConfig::DEFAULT;
1301        config.snapshot_ttl_ms = 5_000;
1302        let lease = open_stream(config, 0, false, &clock).unwrap();
1303
1304        assert!(!lease.snapshot_expired(clock.now_ms()));
1305        clock.advance(4_999);
1306        assert!(!lease.snapshot_expired(clock.now_ms()));
1307        clock.advance(1);
1308        assert!(lease.snapshot_expired(clock.now_ms()));
1309    }
1310
1311    #[test]
1312    fn stream_config_loads_defaults_when_kv_is_empty() {
1313        // Without runtime, just sanity-check the defaults match ADR 0029.
1314        let cfg = StreamConfig::DEFAULT;
1315        assert_eq!(cfg.snapshot_ttl_ms, 60_000);
1316        assert_eq!(cfg.chunk_default_pages, 4);
1317        assert_eq!(cfg.chunk_min_pages, 1);
1318        assert_eq!(cfg.chunk_max_pages, 64);
1319        assert_eq!(cfg.chunk_max_rows, 1000);
1320        assert_eq!(cfg.chunk_max_latency_ms, 50);
1321        assert_eq!(cfg.production_buffer_bytes(), 64 * 1024);
1322    }
1323
1324    #[test]
1325    fn stream_config_normalize_clamps_inconsistent_inputs() {
1326        let mut cfg = StreamConfig {
1327            snapshot_ttl_ms: 1,
1328            chunk_default_pages: 100,
1329            chunk_min_pages: 0,
1330            chunk_max_pages: 8,
1331            chunk_max_rows: 0,
1332            chunk_max_latency_ms: 1,
1333            max_global_streams: 0,
1334            max_per_principal_streams: 0,
1335            default_verify: crate::runtime::integrity_tombstone::VerifyMode::None,
1336        };
1337        cfg.normalize();
1338        assert_eq!(cfg.chunk_min_pages, 1);
1339        assert_eq!(cfg.chunk_max_pages, 8);
1340        assert_eq!(cfg.chunk_default_pages, 8); // clamped down to max
1341        assert!(cfg.chunk_max_rows >= 1);
1342        assert!(cfg.max_global_streams >= 1);
1343        assert!(cfg.max_per_principal_streams >= 1);
1344    }
1345
1346    /// Test sink that accumulates flushed chunks into an interior-mutable
1347    /// `Vec`. Avoids the closure-capture-mutable-then-borrow-immutable
1348    /// dance that the borrow checker rejects when assertions are
1349    /// interleaved with `push_line` calls.
1350    struct CapturingSink {
1351        chunks: std::cell::RefCell<Vec<Vec<u8>>>,
1352    }
1353    impl CapturingSink {
1354        fn new() -> Self {
1355            Self {
1356                chunks: std::cell::RefCell::new(Vec::new()),
1357            }
1358        }
1359        fn len(&self) -> usize {
1360            self.chunks.borrow().len()
1361        }
1362        fn last_len(&self) -> Option<usize> {
1363            self.chunks.borrow().last().map(|c| c.len())
1364        }
1365    }
1366
1367    fn capture<'a>(sink: &'a CapturingSink) -> impl FnMut(&[u8]) -> std::io::Result<()> + 'a {
1368        move |bytes: &[u8]| {
1369            sink.chunks.borrow_mut().push(bytes.to_vec());
1370            Ok(())
1371        }
1372    }
1373
1374    #[test]
1375    fn chunk_producer_flushes_on_byte_cap() {
1376        let clock = FakeClock::new(0);
1377        let cfg = StreamConfig {
1378            chunk_default_pages: 1, // 16 KiB
1379            chunk_min_pages: 1,
1380            chunk_max_pages: 1,
1381            chunk_max_rows: 1_000_000,
1382            chunk_max_latency_ms: 1_000_000,
1383            ..StreamConfig::DEFAULT
1384        };
1385        let sink = CapturingSink::new();
1386        let mut producer = ChunkProducer::new(&cfg, &clock);
1387        let mut flush = capture(&sink);
1388
1389        producer
1390            .push_line(&vec![b'x'; 8 * 1024], &mut flush)
1391            .unwrap();
1392        assert_eq!(sink.len(), 0);
1393
1394        let triggered = producer
1395            .push_line(&vec![b'y'; 8 * 1024], &mut flush)
1396            .unwrap();
1397        assert!(triggered);
1398        assert_eq!(sink.len(), 1);
1399        assert_eq!(sink.last_len(), Some(16 * 1024));
1400        assert_eq!(producer.last_flush_reason(), Some(FlushReason::Byte));
1401    }
1402
1403    #[test]
1404    fn chunk_producer_flushes_on_row_cap() {
1405        let clock = FakeClock::new(0);
1406        let cfg = StreamConfig {
1407            chunk_default_pages: 4, // 64 KiB — well above any test row size
1408            chunk_min_pages: 1,
1409            chunk_max_pages: 64,
1410            chunk_max_rows: 3,
1411            chunk_max_latency_ms: 1_000_000,
1412            ..StreamConfig::DEFAULT
1413        };
1414        let sink = CapturingSink::new();
1415        let mut producer = ChunkProducer::new(&cfg, &clock);
1416        let mut flush = capture(&sink);
1417        let row = b"{\"row\":{}}\n";
1418        producer.push_line(row, &mut flush).unwrap();
1419        producer.push_line(row, &mut flush).unwrap();
1420        assert_eq!(sink.len(), 0);
1421        let triggered = producer.push_line(row, &mut flush).unwrap();
1422        assert!(triggered);
1423        assert_eq!(sink.len(), 1);
1424        assert_eq!(producer.last_flush_reason(), Some(FlushReason::Row));
1425    }
1426
1427    #[test]
1428    fn chunk_producer_flushes_on_latency_cap() {
1429        let clock = FakeClock::new(0);
1430        let cfg = StreamConfig {
1431            chunk_default_pages: 4,
1432            chunk_min_pages: 1,
1433            chunk_max_pages: 64,
1434            chunk_max_rows: 1_000_000,
1435            chunk_max_latency_ms: 50,
1436            ..StreamConfig::DEFAULT
1437        };
1438        let sink = CapturingSink::new();
1439        let mut producer = ChunkProducer::new(&cfg, &clock);
1440        let mut flush = capture(&sink);
1441        producer.push_line(b"{\"row\":{}}\n", &mut flush).unwrap();
1442        assert_eq!(sink.len(), 0);
1443        clock.advance(60);
1444        let triggered = producer.push_line(b"{\"row\":{}}\n", &mut flush).unwrap();
1445        assert!(triggered);
1446        assert_eq!(producer.last_flush_reason(), Some(FlushReason::Latency));
1447    }
1448
1449    #[test]
1450    fn chunk_producer_finish_emits_terminal_flush() {
1451        let clock = FakeClock::new(0);
1452        let cfg = StreamConfig::DEFAULT;
1453        let sink = CapturingSink::new();
1454        let mut producer = ChunkProducer::new(&cfg, &clock);
1455        let mut flush = capture(&sink);
1456        producer.push_line(b"{\"row\":{}}\n", &mut flush).unwrap();
1457        producer.finish(&mut flush).unwrap();
1458        assert_eq!(sink.len(), 1);
1459        assert_eq!(producer.last_flush_reason(), Some(FlushReason::Terminal));
1460    }
1461
1462    #[test]
1463    fn write_chunked_helpers_produce_well_formed_chunks() {
1464        let mut buf: Vec<u8> = Vec::new();
1465        write_chunked_response_header(&mut buf, 200, "application/x-ndjson").unwrap();
1466        write_chunk(&mut buf, b"{\"row\":{}}\n").unwrap();
1467        write_chunked_terminator(&mut buf).unwrap();
1468        let text = String::from_utf8(buf).unwrap();
1469        assert!(text.starts_with("HTTP/1.1 200 OK\r\n"));
1470        assert!(text.contains("Transfer-Encoding: chunked\r\n"));
1471        // 11 bytes in `{"row":{}}\n` → hex 'b'.
1472        assert!(text.contains("\r\nb\r\n{\"row\":{}}\n\r\n"));
1473        assert!(text.ends_with("0\r\n\r\n"));
1474    }
1475
1476    // ──────── Issue #761 / S2 — capacity registry ────────
1477
1478    #[test]
1479    fn capacity_registry_global_exhausted_returns_structured_error() {
1480        let reg = StreamCapacityRegistry::new();
1481        let _g1 = reg.try_acquire("alice", 2, 32).unwrap();
1482        let _g2 = reg.try_acquire("alice", 2, 32).unwrap();
1483        let err = reg.try_acquire("alice", 2, 32).unwrap_err();
1484        assert_eq!(
1485            err,
1486            AcquireError::GlobalExhausted {
1487                limit: 2,
1488                current: 2,
1489            }
1490        );
1491        assert_eq!(err.code(), "server_stream_capacity_exhausted");
1492    }
1493
1494    #[test]
1495    fn capacity_registry_per_principal_exhausted_independent_of_global() {
1496        // Acceptance criterion #2: per-principal cap fires even when
1497        // global has room. Acceptance criterion #3: counters are
1498        // independent across principals.
1499        let reg = StreamCapacityRegistry::new();
1500        let _a1 = reg.try_acquire("alice", 100, 2).unwrap();
1501        let _a2 = reg.try_acquire("alice", 100, 2).unwrap();
1502        let err = reg.try_acquire("alice", 100, 2).unwrap_err();
1503        assert_eq!(
1504            err,
1505            AcquireError::PrincipalExhausted {
1506                principal: "alice".to_string(),
1507                limit: 2,
1508                current: 2,
1509            }
1510        );
1511        assert_eq!(err.code(), "principal_stream_quota_exhausted");
1512
1513        // Bob is unaffected by Alice's quota.
1514        let _b1 = reg.try_acquire("bob", 100, 2).unwrap();
1515        let _b2 = reg.try_acquire("bob", 100, 2).unwrap();
1516    }
1517
1518    #[test]
1519    fn capacity_registry_release_frees_both_counters() {
1520        // Acceptance criterion #4: drop releases both counters.
1521        let reg = StreamCapacityRegistry::new();
1522        let g1 = reg.try_acquire("alice", 1, 1).unwrap();
1523        assert!(reg.try_acquire("alice", 1, 1).is_err());
1524        drop(g1);
1525        let (global, per_principal) = reg.snapshot();
1526        assert_eq!(global, 0);
1527        assert!(per_principal.is_empty());
1528        // Slot is now reclaimable.
1529        let _g2 = reg.try_acquire("alice", 1, 1).unwrap();
1530    }
1531
1532    #[test]
1533    fn capacity_registry_concurrent_acquire_release_does_not_over_issue() {
1534        // Acceptance criterion #5: stress coverage. Spawn `THREADS`
1535        // threads each running `ITERS` acquire+release cycles against
1536        // a registry sized to fit only `CAP` slots; the live count
1537        // must never exceed `CAP`, and the registry must return to
1538        // zero once every thread has joined.
1539        use std::sync::atomic::{AtomicUsize, Ordering};
1540
1541        const THREADS: usize = 16;
1542        const ITERS: usize = 200;
1543        const CAP_GLOBAL: usize = 4;
1544        const CAP_PER_PRINCIPAL: usize = 4;
1545
1546        let reg = StreamCapacityRegistry::new();
1547        let observed_max = Arc::new(AtomicUsize::new(0));
1548        let mut handles = Vec::new();
1549        for tid in 0..THREADS {
1550            let reg = Arc::clone(&reg);
1551            let observed_max = Arc::clone(&observed_max);
1552            // Two principals share the global cap, each capped at
1553            // `CAP_PER_PRINCIPAL` themselves.
1554            let principal = format!("p{}", tid % 2);
1555            handles.push(std::thread::spawn(move || {
1556                for _ in 0..ITERS {
1557                    if let Ok(guard) = reg.try_acquire(&principal, CAP_GLOBAL, CAP_PER_PRINCIPAL) {
1558                        let (live, _) = reg.snapshot();
1559                        observed_max.fetch_max(live, Ordering::SeqCst);
1560                        // Hold the slot just long enough to let other
1561                        // threads race against the cap.
1562                        std::thread::yield_now();
1563                        drop(guard);
1564                    }
1565                }
1566            }));
1567        }
1568        for h in handles {
1569            h.join().unwrap();
1570        }
1571        let (global_after, per_principal_after) = reg.snapshot();
1572        assert_eq!(global_after, 0, "global counter leaked");
1573        assert!(
1574            per_principal_after.is_empty(),
1575            "per-principal map leaked: {per_principal_after:?}"
1576        );
1577        assert!(
1578            observed_max.load(Ordering::SeqCst) <= CAP_GLOBAL,
1579            "global cap was breached: observed {} > {}",
1580            observed_max.load(Ordering::SeqCst),
1581            CAP_GLOBAL
1582        );
1583    }
1584
1585    // ──────── Issue #766 / S7 — resume coordinator ────────
1586
1587    #[test]
1588    fn assess_resumability_accepts_plain_select() {
1589        assert!(assess_resumability("SELECT id, name FROM t"));
1590        assert!(assess_resumability("select * from t where id > 5"));
1591        assert!(assess_resumability("SELECT a, b FROM t ORDER BY rid"));
1592        assert!(assess_resumability("SELECT a, b FROM t ORDER BY rid ASC"));
1593        assert!(assess_resumability(
1594            "SELECT a FROM t ORDER BY rid ASC LIMIT 10"
1595        ));
1596    }
1597
1598    #[test]
1599    fn assess_resumability_rejects_aggregates_and_unordered() {
1600        assert!(!assess_resumability("SELECT COUNT(*) FROM t"));
1601        assert!(!assess_resumability("SELECT SUM(x) FROM t"));
1602        assert!(!assess_resumability("SELECT a, COUNT(b) FROM t GROUP BY a"));
1603        assert!(!assess_resumability("SELECT DISTINCT a FROM t"));
1604        assert!(!assess_resumability("SELECT a FROM t ORDER BY name"));
1605        assert!(!assess_resumability("SELECT a FROM t ORDER BY rid DESC"));
1606        assert!(!assess_resumability("SELECT a FROM t ORDER BY a, b"));
1607        assert!(!assess_resumability("INSERT INTO t (a) VALUES (1)"));
1608        assert!(!assess_resumability(
1609            "SELECT a FROM t JOIN u ON t.id = u.id"
1610        ));
1611    }
1612
1613    #[test]
1614    fn lease_registry_records_and_expires_against_ttl() {
1615        let reg = LeaseRegistry::new();
1616        reg.record(42, 1_000, 5_000);
1617        assert_eq!(reg.lookup(42, 1_000), LeaseLookup::Live);
1618        assert_eq!(reg.lookup(42, 5_999), LeaseLookup::Live);
1619        assert_eq!(reg.lookup(42, 6_000), LeaseLookup::Expired);
1620        assert_eq!(reg.lookup(99, 1_000), LeaseLookup::Unknown);
1621    }
1622
1623    #[test]
1624    fn prefix_hasher_is_order_sensitive_and_deterministic() {
1625        let mut a = PrefixHasher::new();
1626        a.update(b"{\"row\":{\"id\":1}}");
1627        a.update(b"{\"row\":{\"id\":2}}");
1628        let hash_a = a.finalize_hex();
1629
1630        let mut b = PrefixHasher::new();
1631        b.update(b"{\"row\":{\"id\":1}}");
1632        b.update(b"{\"row\":{\"id\":2}}");
1633        let hash_b = b.finalize_hex();
1634        assert_eq!(hash_a, hash_b);
1635
1636        let mut c = PrefixHasher::new();
1637        c.update(b"{\"row\":{\"id\":2}}");
1638        c.update(b"{\"row\":{\"id\":1}}");
1639        assert_ne!(hash_a, c.finalize_hex());
1640        assert_eq!(hash_a.len(), 64);
1641    }
1642
1643    // ──────── Issue #767 / S8 — opaque lease handle + audit ────────
1644
1645    #[test]
1646    fn lease_handle_is_128_bit_hex_and_unique_per_open() {
1647        // Acceptance criterion #3 — wire-visible handle is opaque,
1648        // 128-bit, non-sequential. The internal monotonic `id` is
1649        // independent of the handle.
1650        let clock = FakeClock::new(0);
1651        let a = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1652        let b = open_stream(StreamConfig::DEFAULT, 1, false, &clock).unwrap();
1653        assert_eq!(
1654            a.lease_handle.len(),
1655            LEASE_HANDLE_BYTES * 2,
1656            "handle must be 128 bits hex-encoded: {}",
1657            a.lease_handle
1658        );
1659        assert!(
1660            a.lease_handle.chars().all(|c| c.is_ascii_hexdigit()),
1661            "handle must be hex: {}",
1662            a.lease_handle
1663        );
1664        assert_ne!(a.lease_handle, b.lease_handle, "handles must differ");
1665        // Sanity — internal id remains monotonic, handle is not
1666        // derivable from it (handle bytes don't encode the id).
1667        assert!(b.id > a.id);
1668    }
1669
1670    #[test]
1671    fn generate_lease_handle_produces_high_entropy_distinct_values() {
1672        // Defensive — large sample to surface a hypothetical CSPRNG
1673        // mis-wiring. 1024 draws with zero collisions on 128 bits is
1674        // statistically certain when the source is sound.
1675        let mut seen = std::collections::HashSet::new();
1676        for _ in 0..1024 {
1677            assert!(
1678                seen.insert(generate_lease_handle()),
1679                "duplicate handle in CSPRNG sequence"
1680            );
1681        }
1682    }
1683
1684    #[test]
1685    fn parse_jwt_exp_ms_extracts_seconds_to_ms() {
1686        // Synthetic JWT — header.payload.signature. We only consult
1687        // the payload (no signature verification).
1688        // Payload: {"exp":1700000000}
1689        // Base64url of that JSON: eyJleHAiOjE3MDAwMDAwMDB9
1690        let token = "eyJhbGciOiJIUzI1NiJ9.eyJleHAiOjE3MDAwMDAwMDB9.sig";
1691        assert_eq!(parse_jwt_exp_ms(token), Some(1_700_000_000_000));
1692    }
1693
1694    #[test]
1695    fn parse_jwt_exp_ms_returns_none_for_opaque_tokens() {
1696        assert_eq!(parse_jwt_exp_ms("not-a-jwt"), None);
1697        assert_eq!(parse_jwt_exp_ms("only.two"), None);
1698        assert_eq!(parse_jwt_exp_ms("a.b.c"), None);
1699    }
1700
1701    // ──────── Issue #807 / 750c — cursor registry ────────
1702
1703    fn register_default_cursor(reg: &CursorRegistry, now_ms: u64) -> String {
1704        reg.register(
1705            42,
1706            "acme",
1707            "bearer:abc",
1708            "SELECT id FROM users ORDER BY rid",
1709            None,
1710            None,
1711            now_ms,
1712            1_000,
1713        )
1714    }
1715
1716    #[test]
1717    fn cursor_token_is_192_bit_hex_and_unique_per_register() {
1718        let reg = CursorRegistry::default();
1719        let a = register_default_cursor(&reg, 0);
1720        let b = register_default_cursor(&reg, 0);
1721        assert_eq!(
1722            a.len(),
1723            CURSOR_TOKEN_BYTES * 2,
1724            "token must be 192 bits hex-encoded: {a}"
1725        );
1726        assert!(
1727            a.chars().all(|c| c.is_ascii_hexdigit()),
1728            "token must be hex: {a}"
1729        );
1730        assert_ne!(a, b, "tokens must differ across registrations");
1731        assert_eq!(reg.len(), 2);
1732    }
1733
1734    #[test]
1735    fn cursor_resolves_for_owner_within_ttl() {
1736        // Happy-path resume — same tenant + principal, before TTL.
1737        let reg = CursorRegistry::default();
1738        let token = register_default_cursor(&reg, 1_000);
1739        let resume = reg
1740            .resolve(&token, "acme", "bearer:abc", 1_500)
1741            .expect("live cursor resolves for its owner");
1742        assert_eq!(resume.snapshot_lsn, 42, "resume re-pins the same snapshot");
1743        assert_eq!(resume.query, "SELECT id FROM users ORDER BY rid");
1744        assert_eq!(resume.expires_at_ms, 2_000);
1745    }
1746
1747    #[test]
1748    fn cursor_rejects_after_ttl_for_owner() {
1749        // TTL expiry — owner sees a distinct `Expired` so they can
1750        // re-issue, but the cursor no longer resolves.
1751        let reg = CursorRegistry::default();
1752        let token = register_default_cursor(&reg, 0);
1753        assert_eq!(
1754            reg.resolve(&token, "acme", "bearer:abc", 1_000),
1755            Err(CursorReject::Expired),
1756            "TTL boundary is inclusive — cursor is dead at opened_at + ttl"
1757        );
1758        assert_eq!(
1759            reg.resolve(&token, "acme", "bearer:abc", 5_000),
1760            Err(CursorReject::Expired)
1761        );
1762    }
1763
1764    #[test]
1765    fn cursor_cross_tenant_is_masked_as_not_found() {
1766        // Tenant isolation — a different tenant cannot tell the token
1767        // from one that never existed (no existence leak), even while the
1768        // cursor is still live for its real owner.
1769        let reg = CursorRegistry::default();
1770        let token = register_default_cursor(&reg, 0);
1771        assert_eq!(
1772            reg.resolve(&token, "evil-corp", "bearer:abc", 100),
1773            Err(CursorReject::NotFound),
1774            "cross-tenant resume must mask existence as NotFound"
1775        );
1776        // Still live for the rightful tenant.
1777        assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1778    }
1779
1780    #[test]
1781    fn cursor_cross_principal_is_masked_as_not_found() {
1782        // Principal isolation — same tenant, different principal is also
1783        // masked as NotFound rather than leaking a permission error.
1784        let reg = CursorRegistry::default();
1785        let token = register_default_cursor(&reg, 0);
1786        assert_eq!(
1787            reg.resolve(&token, "acme", "bearer:other", 100),
1788            Err(CursorReject::NotFound),
1789            "cross-principal resume must mask existence as NotFound"
1790        );
1791        assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1792    }
1793
1794    #[test]
1795    fn cursor_unknown_token_is_not_found() {
1796        let reg = CursorRegistry::default();
1797        assert!(reg.is_empty());
1798        assert_eq!(
1799            reg.resolve("deadbeef", "acme", "bearer:abc", 0),
1800            Err(CursorReject::NotFound)
1801        );
1802    }
1803
1804    #[test]
1805    fn cursor_scope_is_checked_before_expiry() {
1806        // An unauthorized caller hitting an *expired* token must still see
1807        // NotFound, never Expired — Expired would confirm the token once
1808        // existed. Scope precedes expiry in `resolve`.
1809        let reg = CursorRegistry::default();
1810        let token = register_default_cursor(&reg, 0);
1811        assert_eq!(
1812            reg.resolve(&token, "evil-corp", "bearer:abc", 10_000),
1813            Err(CursorReject::NotFound),
1814            "expired + wrong scope must mask as NotFound, not Expired"
1815        );
1816    }
1817
1818    // ──────── Issue #808 / 750d — cancel + tombstone ────────
1819
1820    #[test]
1821    fn cancel_tombstones_cursor_and_raises_token() {
1822        // Cancelling a live cursor raises its shared token (so an in-flight
1823        // executor observing it stops) and tombstones the entry.
1824        let reg = CursorRegistry::default();
1825        let token = register_default_cursor(&reg, 0);
1826        let live = reg
1827            .cancel_token_for(&token)
1828            .expect("freshly-minted cursor exposes its token");
1829        assert!(!live.is_cancelled(), "token starts un-cancelled");
1830
1831        let returned = reg
1832            .cancel(&token, "acme", "bearer:abc")
1833            .expect("owner cancels its own cursor");
1834        assert!(returned.is_cancelled(), "cancel raises the returned token");
1835        assert!(
1836            live.is_cancelled(),
1837            "the handler-held clone observes the cancel"
1838        );
1839    }
1840
1841    #[test]
1842    fn cancelled_cursor_rejects_resume_with_cancelled_reason() {
1843        // A tombstoned cursor refuses resume with a dedicated `Cancelled`
1844        // reason for its owner — distinct from Expired so the client learns
1845        // the stream was cancelled, not aged out.
1846        let reg = CursorRegistry::default();
1847        let token = register_default_cursor(&reg, 0);
1848        reg.cancel(&token, "acme", "bearer:abc")
1849            .expect("owner cancels");
1850        assert_eq!(
1851            reg.resolve(&token, "acme", "bearer:abc", 100),
1852            Err(CursorReject::Cancelled),
1853            "owner resuming a cancelled cursor sees Cancelled"
1854        );
1855    }
1856
1857    #[test]
1858    fn cancel_is_idempotent() {
1859        let reg = CursorRegistry::default();
1860        let token = register_default_cursor(&reg, 0);
1861        assert!(reg.cancel(&token, "acme", "bearer:abc").is_ok());
1862        // Second cancel still succeeds (already tombstoned) and stays cancelled.
1863        let second = reg
1864            .cancel(&token, "acme", "bearer:abc")
1865            .expect("re-cancel is a no-op success");
1866        assert!(second.is_cancelled());
1867    }
1868
1869    #[test]
1870    fn cancel_cross_tenant_is_masked_as_not_found() {
1871        // An unauthorized caller cannot cancel — or confirm the existence
1872        // of — a foreign cursor.
1873        let reg = CursorRegistry::default();
1874        let token = register_default_cursor(&reg, 0);
1875        assert!(
1876            matches!(
1877                reg.cancel(&token, "evil-corp", "bearer:abc"),
1878                Err(CursorReject::NotFound)
1879            ),
1880            "cross-tenant cancel must mask existence as NotFound"
1881        );
1882        // The rightful owner's cursor is untouched — still resumable.
1883        assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1884    }
1885
1886    #[test]
1887    fn cancel_cross_principal_is_masked_as_not_found() {
1888        let reg = CursorRegistry::default();
1889        let token = register_default_cursor(&reg, 0);
1890        assert!(
1891            matches!(
1892                reg.cancel(&token, "acme", "bearer:other"),
1893                Err(CursorReject::NotFound)
1894            ),
1895            "cross-principal cancel must mask existence as NotFound"
1896        );
1897        assert!(reg.resolve(&token, "acme", "bearer:abc", 100).is_ok());
1898    }
1899
1900    #[test]
1901    fn cancel_unknown_token_is_not_found() {
1902        let reg = CursorRegistry::default();
1903        assert!(matches!(
1904            reg.cancel("deadbeef", "acme", "bearer:abc"),
1905            Err(CursorReject::NotFound)
1906        ));
1907    }
1908
1909    #[test]
1910    fn generate_cursor_token_produces_high_entropy_distinct_values() {
1911        let mut seen = std::collections::HashSet::new();
1912        for _ in 0..1024 {
1913            assert!(
1914                seen.insert(generate_cursor_token()),
1915                "duplicate cursor token in CSPRNG sequence"
1916            );
1917        }
1918    }
1919
1920    #[test]
1921    fn close_reason_as_str_covers_every_state_transition() {
1922        // Acceptance criterion #5 — every state transition in the
1923        // brief is representable in the audit close-reason taxonomy.
1924        for (variant, expected) in [
1925            (CloseReason::Ok, "ok"),
1926            (CloseReason::Cancelled, "cancelled"),
1927            (CloseReason::Error, "error"),
1928            (CloseReason::SnapshotExpired, "snapshot_expired"),
1929            (CloseReason::CapacityRefused, "capacity_refused"),
1930            (CloseReason::IntegrityFailed, "integrity_failed"),
1931        ] {
1932            assert_eq!(variant.as_str(), expected);
1933        }
1934    }
1935
1936    #[test]
1937    fn stream_config_defaults_carry_s2_caps() {
1938        assert_eq!(StreamConfig::DEFAULT.max_global_streams, 256);
1939        assert_eq!(StreamConfig::DEFAULT.max_per_principal_streams, 32);
1940    }
1941
1942    // ──────── Issue #768 / S9 — pull-based driver ────────
1943
1944    /// Sink that records every flushed chunk's length so a test can
1945    /// assert the resident working set stayed bounded.
1946    struct SizeSink {
1947        sizes: std::cell::RefCell<Vec<usize>>,
1948    }
1949    impl SizeSink {
1950        fn new() -> Self {
1951            Self {
1952                sizes: std::cell::RefCell::new(Vec::new()),
1953            }
1954        }
1955        fn flushes(&self) -> usize {
1956            self.sizes.borrow().len()
1957        }
1958        fn max_chunk(&self) -> usize {
1959            self.sizes.borrow().iter().copied().max().unwrap_or(0)
1960        }
1961    }
1962    fn size_capture<'a>(sink: &'a SizeSink) -> impl FnMut(&[u8]) -> std::io::Result<()> + 'a {
1963        move |bytes: &[u8]| {
1964            sink.sizes.borrow_mut().push(bytes.len());
1965            Ok(())
1966        }
1967    }
1968
1969    #[test]
1970    fn drive_lines_streams_large_source_with_bounded_working_set() {
1971        // Acceptance #1: a huge scan flows through the chunk buffer
1972        // without ever materialising the full result set. The source
1973        // is a *lazy* range mapped to records — collecting it would
1974        // allocate N rows, but `drive_lines` pulls one at a time. We
1975        // assert every flushed chunk stays within a small multiple of
1976        // the page buffer, i.e. memory tracks the buffer, not N.
1977        let clock = FakeClock::new(0);
1978        let cfg = StreamConfig {
1979            chunk_default_pages: 1, // 16 KiB buffer
1980            chunk_min_pages: 1,
1981            chunk_max_pages: 1,
1982            chunk_max_rows: 1_000_000, // don't let the row cap dominate
1983            chunk_max_latency_ms: 1_000_000,
1984            ..StreamConfig::DEFAULT
1985        };
1986        let sink = SizeSink::new();
1987        let mut producer = ChunkProducer::new(&cfg, &clock);
1988        let mut flush = size_capture(&sink);
1989
1990        const N: u64 = 1_000_000;
1991        // Lazy source: never collected into a Vec.
1992        let source = 0..N;
1993        let consumed = producer
1994            .drive_lines(
1995                source,
1996                |i: &u64| format!("{{\"row\":{{\"id\":{i}}}}}\n").into_bytes(),
1997                &mut flush,
1998            )
1999            .unwrap();
2000        producer.finish(&mut flush).unwrap();
2001
2002        assert_eq!(consumed, N);
2003        assert_eq!(producer.total_rows(), N);
2004        // Many flushes occurred (streamed), not one giant buffer.
2005        assert!(
2006            sink.flushes() > 1000,
2007            "expected the source to stream across many chunks, saw {}",
2008            sink.flushes()
2009        );
2010        // No chunk exceeded the byte cap plus one trailing line — the
2011        // resident buffer is bounded independent of N.
2012        let max_line = format!("{{\"row\":{{\"id\":{}}}}}\n", N - 1).len();
2013        assert!(
2014            sink.max_chunk() <= cfg.production_buffer_bytes() + max_line,
2015            "chunk {} exceeded bounded working set {}",
2016            sink.max_chunk(),
2017            cfg.production_buffer_bytes() + max_line
2018        );
2019    }
2020
2021    #[test]
2022    fn drive_lines_first_chunk_flushes_on_latency_before_source_drains() {
2023        // Acceptance #2: first-row latency is bounded by the latency
2024        // cap, not by full materialisation. The source yields rows
2025        // whose pull advances the fake clock; the first chunk must
2026        // flush as soon as the latency window elapses, long before the
2027        // (large) source is exhausted.
2028        let clock = FakeClock::new(0);
2029        let cfg = StreamConfig {
2030            chunk_default_pages: 64, // large byte cap — won't trip first
2031            chunk_min_pages: 1,
2032            chunk_max_pages: 64,
2033            chunk_max_rows: 1_000_000, // large row cap — won't trip first
2034            chunk_max_latency_ms: 50,
2035            ..StreamConfig::DEFAULT
2036        };
2037        let sink = SizeSink::new();
2038        let mut producer = ChunkProducer::new(&cfg, &clock);
2039
2040        // Drive manually so we can advance the clock between pulls and
2041        // observe the first flush. Each pull advances 20 ms; the 50 ms
2042        // latency cap trips on the third row.
2043        let mut first_flush_after: Option<u64> = None;
2044        let mut row = 0u64;
2045        while row < 1_000_000 {
2046            let line = format!("{{\"row\":{{\"id\":{row}}}}}\n");
2047            clock.advance(20);
2048            let mut flush = size_capture(&sink);
2049            let flushed = producer.push_line(line.as_bytes(), &mut flush).unwrap();
2050            row += 1;
2051            if flushed {
2052                first_flush_after = Some(row);
2053                break;
2054            }
2055        }
2056
2057        assert_eq!(producer.last_flush_reason(), Some(FlushReason::Latency));
2058        let rows_before_flush = first_flush_after.expect("a latency flush must occur");
2059        assert!(
2060            rows_before_flush <= 4,
2061            "first chunk flushed only after {rows_before_flush} rows; latency bound not honoured"
2062        );
2063        // Crucially, the source was nowhere near drained (1e6 rows).
2064        assert!(rows_before_flush < 1_000_000);
2065    }
2066
2067    #[test]
2068    fn drive_lines_parity_with_manual_push_line() {
2069        // The driver must produce byte-identical output to hand-rolled
2070        // push_line calls — the chunk producer's framing is unchanged.
2071        let clock = FakeClock::new(0);
2072        let cfg = StreamConfig::DEFAULT;
2073
2074        let lines: Vec<Vec<u8>> = (0..50)
2075            .map(|i| format!("{{\"row\":{{\"id\":{i}}}}}\n").into_bytes())
2076            .collect();
2077
2078        let driven = CapturingSink::new();
2079        {
2080            let mut p = ChunkProducer::new(&cfg, &clock);
2081            let mut flush = capture(&driven);
2082            p.drive_lines(lines.iter().cloned(), |l: &Vec<u8>| l.clone(), &mut flush)
2083                .unwrap();
2084            p.finish(&mut flush).unwrap();
2085        }
2086
2087        let manual = CapturingSink::new();
2088        {
2089            let mut p = ChunkProducer::new(&cfg, &clock);
2090            let mut flush = capture(&manual);
2091            for l in &lines {
2092                p.push_line(l, &mut flush).unwrap();
2093            }
2094            p.finish(&mut flush).unwrap();
2095        }
2096
2097        assert_eq!(*driven.chunks.borrow(), *manual.chunks.borrow());
2098    }
2099}