dakera-inference 0.11.77

//! Cross-encoder reranker for improving recall precision.
//!
//! Uses BAAI/bge-reranker-base (Xenova ONNX INT8 quantized) to score
//! (query, passage) pairs for relevance. More accurate than bi-encoder
//! vector similarity but slower — used as a second-stage reranker after
//! ANN candidate retrieval.
//!
//! # Architecture
//!
//! ```text
//! query + passage → [CLS] query [SEP] passage [SEP]
//!                       ↓ BERT forward pass
//!                   logits [batch, 1]
//!                       ↓ sigmoid
//!                   relevance scores ∈ [0, 1]
//! ```
//!
//! # Session Pool
//!
//! The engine maintains `RERANKER_POOL_SIZE` independent ONNX sessions.
//! Large candidate lists are split into chunks of `RERANKER_CHUNK_SIZE` and
//! dispatched in parallel across the pool, eliminating head-of-line blocking
//! when multiple recall calls arrive concurrently (DAK-5873).
//!
//! # ONNX Mini-Batching
//!
//! Within each dispatched chunk, pairs are further split into mini-batches of
//! `RERANKER_ONNX_BATCH_SIZE` for a single `session.run()` call. Smaller mini-batches
//! reduce sequence-padding overhead: each mini-batch pads to its own maximum
//! sequence length rather than the chunk maximum, cutting wasted compute when
//! passage lengths vary widely (DAK-5883).

use crate::engine::EmbeddingEngine;
use crate::error::{InferenceError, Result};
use ort::execution_providers::CUDAExecutionProvider;
use ort::inputs;
use ort::session::builder::GraphOptimizationLevel;
use ort::session::Session;
use ort::value::Tensor;
use parking_lot::Mutex;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokenizers::{
    EncodeInput, InputSequence, PaddingParams, PaddingStrategy, Tokenizer, TruncationParams,
};
use tracing::{info, instrument, warn};

/// The reranker model Xenova HuggingFace repo ID (ONNX INT8).
const RERANKER_REPO_ID: &str = "Xenova/bge-reranker-base";
/// ONNX quantized model filename within the repo.
const RERANKER_ONNX_FILE: &str = "onnx/model_quantized.onnx";
/// Maximum token length for cross-encoder input (query + passage combined).
const MAX_SEQ_LENGTH: usize = 512;
/// Number of independent ONNX sessions in the reranker pool.
///
/// Two sessions allow concurrent recall requests to rerank in parallel without
/// head-of-line mutex blocking. Each session uses `intra_threads=4`; two sessions
/// occupy all 8 vCPUs on the production CPX32 server (DAK-5873).
const RERANKER_POOL_SIZE: usize = 2;

/// Maximum concurrent `score_pairs` callers allowed before returning `Overloaded`.
///
/// Root cause of DAK-5893: with `RERANKER_POOL_SIZE=2` and 8 concurrent bench recall
/// requests, the 7th/8th caller waited >120s for a Mutex slot → client timeout →
/// 8-attempt retry loop → 19-minute stall. Capping at `POOL_SIZE * 3 = 6` lets
/// 6 requests queue shallowly (each waits at most 2 ahead on its session) while
/// the 7th+ returns immediately so the API falls back to unranked results.
const RERANKER_MAX_CONCURRENT: usize = RERANKER_POOL_SIZE * 3;

/// Maximum candidates per session sub-batch (parallel dispatch unit).
///
/// Large candidate lists (e.g. temporal `fetch_n = top_k × 8 = 160`) are split
/// into chunks of this size and dispatched concurrently across the pool. With
/// `RERANKER_POOL_SIZE=2` and `RERANKER_CHUNK_SIZE=32`, a 160-candidate list
/// produces 5 chunks: sessions[0] handles chunks 0/2/4, sessions[1] handles
/// chunks 1/3 — 3 serial chunks on the busier session vs. 5 serial before.
const RERANKER_CHUNK_SIZE: usize = 32;
/// Maximum candidate pairs per single ONNX `session.run()` call (inner batch).
///
/// Within each dispatched chunk, pairs are further split into mini-batches of
/// this size. Each mini-batch is padded to its own maximum sequence length,
/// reducing wasted computation when passage lengths vary (padding overhead on
/// shorter passages is bounded by the max within the mini-batch, not the chunk).
///
/// Tuning guide:
/// - Smaller (8): less padding waste, more `session.run()` calls per chunk.
/// - Larger (32, equal to CHUNK_SIZE): behaves like the pre-DAK-5883 single-call
///   mode, effectively disabling inner batching.
/// - Default 16: halves padding overhead on mixed-length candidate lists while
///   keeping `session.run()` call count to 2 per full chunk.
const RERANKER_ONNX_BATCH_SIZE: usize = 16;

/// RAII guard that decrements the `active_requests` counter when dropped.
struct ActiveGuard(Arc<AtomicUsize>);

impl Drop for ActiveGuard {
    fn drop(&mut self) {
        self.0.fetch_sub(1, Ordering::SeqCst);
    }
}

/// Cross-encoder reranking engine.
///
/// Thread-safe — shared via `Arc`. Maintains a pool of independent ONNX sessions
/// so concurrent rerank calls never contend on a single mutex.
pub struct CrossEncoderEngine {
    /// Pool of independent ONNX sessions (round-robin dispatch).
    sessions: Vec<Arc<Mutex<Session>>>,
    tokenizer: Arc<Tokenizer>,
    /// Whether the loaded ONNX model expects a `token_type_ids` input tensor.
    /// bge-reranker-base only has `input_ids` + `attention_mask`; some other
    /// cross-encoders include `token_type_ids`. Determined at load time.
    has_token_type_ids: bool,
    /// Round-robin counter for session assignment.
    next_session: AtomicUsize,
    /// Active concurrent callers of `score_pairs`. When this reaches
    /// `RERANKER_MAX_CONCURRENT`, new callers return `Overloaded` immediately
    /// so the API can fall back to unranked results rather than queuing
    /// indefinitely — root cause fix for DAK-5893 SIGTERM stall.
    active_requests: Arc<AtomicUsize>,
}

impl CrossEncoderEngine {
    /// Load or download the reranker model.
    ///
    /// Downloads `Xenova/bge-reranker-base` ONNX INT8 model from HuggingFace Hub
    /// if not already cached. Builds `RERANKER_POOL_SIZE` independent sessions.
    #[instrument(skip_all)]
    pub async fn new(cache_dir: Option<String>) -> Result<Self> {
        info!("Initializing cross-encoder reranker: {}", RERANKER_REPO_ID);

        let (tokenizer_path, onnx_path) =
            tokio::task::spawn_blocking(move || download_reranker_files(cache_dir))
                .await
                .map_err(|e| InferenceError::ModelLoadError(format!("spawn_blocking: {e}")))?
                .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?;

        info!("Loading reranker tokenizer from {:?}", tokenizer_path);
        let mut tokenizer = Tokenizer::from_file(&tokenizer_path)
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))?;

        // Configure padding + truncation for uniform batch shapes
        let padding = PaddingParams {
            strategy: PaddingStrategy::BatchLongest,
            pad_id: tokenizer.get_padding().map_or(0, |p| p.pad_id),
            pad_token: tokenizer
                .get_padding()
                .map_or("[PAD]".to_string(), |p| p.pad_token.clone()),
            ..Default::default()
        };
        tokenizer.with_padding(Some(padding));
        let truncation = TruncationParams {
            max_length: MAX_SEQ_LENGTH,
            ..Default::default()
        };
        let _ = tokenizer.with_truncation(Some(truncation));

        info!(
            "Loading reranker ONNX model from {:?} (pool_size={}, onnx_batch_size={})",
            onnx_path, RERANKER_POOL_SIZE, RERANKER_ONNX_BATCH_SIZE
        );

        let use_gpu = std::env::var("DAKERA_USE_GPU")
            .map(|v| v == "1")
            .unwrap_or(false);
        if use_gpu {
            info!("CUDA execution provider enabled for reranker (DAKERA_USE_GPU=1)");
        }

        // Build pool of independent ONNX sessions — each has its own ORT context
        // so pool members never block each other under concurrent rerank calls.
        let (sessions, has_token_type_ids) =
            tokio::task::spawn_blocking(move || -> Result<(Vec<Arc<Mutex<Session>>>, bool)> {
                let raw: Result<Vec<Session>> = (0..RERANKER_POOL_SIZE)
                    .map(|_| {
                        let builder = Session::builder()
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                            .with_optimization_level(GraphOptimizationLevel::Level3)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                            .with_intra_threads(4)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?;

                        let mut builder = if use_gpu {
                            builder
                                .with_execution_providers(
                                    [CUDAExecutionProvider::default().build()],
                                )
                                .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                        } else {
                            builder
                        };

                        builder
                            .commit_from_file(&onnx_path)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))
                    })
                    .collect();
                let raw = raw?;
                // Inspect first session to detect optional token_type_ids input.
                let has_tti = raw[0].inputs().iter().any(|i| i.name() == "token_type_ids");
                let sessions: Vec<Arc<Mutex<Session>>> =
                    raw.into_iter().map(|s| Arc::new(Mutex::new(s))).collect();
                Ok((sessions, has_tti))
            })
            .await
            .map_err(|e| InferenceError::ModelLoadError(format!("spawn_blocking: {e}")))?
            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?;

        info!(
            has_token_type_ids,
            pool_size = sessions.len(),
            onnx_batch_size = RERANKER_ONNX_BATCH_SIZE,
            "Cross-encoder reranker loaded successfully"
        );

        Ok(Self {
            sessions,
            tokenizer: Arc::new(tokenizer),
            has_token_type_ids,
            next_session: AtomicUsize::new(0),
            active_requests: Arc::new(AtomicUsize::new(0)),
        })
    }

    /// Score a batch of (query, passage) pairs.
    ///
    /// Passages are split into chunks of [`RERANKER_CHUNK_SIZE`] and dispatched
    /// in parallel across the session pool (round-robin). Within each chunk,
    /// pairs are processed in mini-batches of [`RERANKER_ONNX_BATCH_SIZE`] to
    /// reduce sequence-padding overhead. Chunk results are reassembled in input
    /// order.
    ///
    /// Returns `Err(InferenceError::Overloaded)` immediately when more than
    /// `RERANKER_MAX_CONCURRENT` callers are active — the API layer falls back
    /// to unranked results rather than queuing indefinitely (DAK-5893 fix).
    ///
    /// Returns a relevance score in `[0, 1]` for each passage.
    /// Higher scores indicate greater relevance to the query.
    #[instrument(skip(self, passages), fields(n_passages = passages.len()))]
    pub async fn score_pairs(&self, query: &str, passages: &[String]) -> Result<Vec<f32>> {
        if passages.is_empty() {
            return Ok(Vec::new());
        }

        // Concurrency gate: if already at capacity, return Overloaded so the API
        // falls back to unranked results instead of queuing for >120s (DAK-5893).
        let prev = self.active_requests.fetch_add(1, Ordering::SeqCst);
        if prev >= RERANKER_MAX_CONCURRENT {
            self.active_requests.fetch_sub(1, Ordering::SeqCst);
            warn!(
                active = prev,
                max = RERANKER_MAX_CONCURRENT,
                "Cross-encoder at capacity — returning Overloaded (API will use unranked results)"
            );
            return Err(InferenceError::Overloaded {
                active: prev,
                max: RERANKER_MAX_CONCURRENT,
            });
        }
        // RAII decrement: always release the slot on return, including errors.
        let _guard = ActiveGuard(Arc::clone(&self.active_requests));

        let pool_len = self.sessions.len();
        // Round-robin start: each concurrent caller gets a different initial slot
        // so concurrent requests don't all contend on sessions[0].
        let start_idx = self.next_session.fetch_add(1, Ordering::Relaxed);
        let tokenizer = Arc::clone(&self.tokenizer);
        let has_token_type_ids = self.has_token_type_ids;
        let query_str = query.to_string();

        // Split candidates into RERANKER_CHUNK_SIZE sub-batches.
        let chunks: Vec<Vec<String>> = passages
            .chunks(RERANKER_CHUNK_SIZE)
            .map(<[String]>::to_vec)
            .collect();

        // Spawn all chunks concurrently; each acquires its own session slot.
        let mut handles = Vec::with_capacity(chunks.len());
        for (i, chunk) in chunks.into_iter().enumerate() {
            let session = Arc::clone(&self.sessions[(start_idx + i) % pool_len]);
            let tok = Arc::clone(&tokenizer);
            let q = query_str.clone();
            handles.push(tokio::task::spawn_blocking(move || {
                score_pairs_blocking(&session, &tok, &q, &chunk, has_token_type_ids)
            }));
        }

        // Collect results in chunk order to preserve passage ordering.
        let mut scores = Vec::with_capacity(passages.len());
        for handle in handles {
            let chunk_scores = handle
                .await
                .map_err(|e| InferenceError::InferenceError(format!("spawn_blocking: {e}")))??;
            scores.extend(chunk_scores);
        }

        Ok(scores)
    }

    /// Number of parallel ONNX sessions in the pool.
    pub fn pool_size(&self) -> usize {
        self.sessions.len()
    }

    /// Configured ONNX mini-batch size (pairs per `session.run()` call).
    pub fn onnx_batch_size(&self) -> usize {
        RERANKER_ONNX_BATCH_SIZE
    }

    /// Current number of active concurrent `score_pairs` calls.
    /// Used by metrics and health checks (DAK-5893).
    pub fn active_requests_count(&self) -> usize {
        self.active_requests.load(Ordering::Relaxed)
    }

    /// Maximum concurrent `score_pairs` calls before `Overloaded` is returned.
    pub fn max_concurrent(&self) -> usize {
        RERANKER_MAX_CONCURRENT
    }
}

/// Blocking cross-encoder inference for one chunk — runs inside `spawn_blocking`.
///
/// The chunk is processed as a sequence of mini-batches of [`RERANKER_ONNX_BATCH_SIZE`]
/// pairs. Each mini-batch issues one `session.run()` call and pads to its own
/// maximum sequence length, reducing waste compared to padding the full chunk.
/// The session mutex is held for all mini-batches in the chunk to avoid per-mini-batch
/// acquire/release overhead.
fn score_pairs_blocking(
    session: &Arc<Mutex<Session>>,
    tokenizer: &Tokenizer,
    query: &str,
    passages: &[String],
    has_token_type_ids: bool,
) -> Result<Vec<f32>> {
    let total = passages.len();
    if total == 0 {
        return Ok(Vec::new());
    }

    let mut all_scores = Vec::with_capacity(total);
    // Hold the lock for the entire chunk to eliminate per-mini-batch
    // acquire/release cost. Total lock duration is unchanged vs. the
    // pre-DAK-5883 single-call approach since total compute is the same.
    let mut sess = session.lock();

    for mini_batch in passages.chunks(RERANKER_ONNX_BATCH_SIZE) {
        let batch_size = mini_batch.len();

        // Build EncodeInput pairs: [CLS] query [SEP] passage [SEP]
        let inputs: Vec<EncodeInput> = mini_batch
            .iter()
            .map(|p| EncodeInput::Dual(InputSequence::from(query), InputSequence::from(p.as_str())))
            .collect();

        let encodings = tokenizer
            .encode_batch(inputs, true)
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))?;

        let seq_len = encodings.first().map(|e| e.get_ids().len()).unwrap_or(0);
        if seq_len == 0 {
            all_scores.extend(std::iter::repeat_n(0.5f32, batch_size));
            continue;
        }

        // Flatten to i64 arrays (ORT BERT models expect int64)
        let mut input_ids = Vec::with_capacity(batch_size * seq_len);
        let mut attention_mask = Vec::with_capacity(batch_size * seq_len);
        let mut token_type_ids = Vec::with_capacity(batch_size * seq_len);

        for enc in &encodings {
            input_ids.extend(enc.get_ids().iter().map(|&id| id as i64));
            attention_mask.extend(enc.get_attention_mask().iter().map(|&m| m as i64));
            let type_ids = enc.get_type_ids();
            if type_ids.is_empty() {
                token_type_ids.extend(std::iter::repeat_n(0i64, seq_len));
            } else {
                token_type_ids.extend(type_ids.iter().map(|&t| t as i64));
            }
        }

        // Build ORT tensors
        let input_ids_tensor = Tensor::<i64>::from_array(([batch_size, seq_len], input_ids))
            .map_err(|e| InferenceError::InferenceError(e.to_string()))?;
        let attention_mask_tensor =
            Tensor::<i64>::from_array(([batch_size, seq_len], attention_mask))
                .map_err(|e| InferenceError::InferenceError(e.to_string()))?;
        let token_type_ids_tensor =
            Tensor::<i64>::from_array(([batch_size, seq_len], token_type_ids))
                .map_err(|e| InferenceError::InferenceError(e.to_string()))?;

        // Run inference in a scoped block so `outputs` drops before the next
        // mini-batch iteration reuses the session — the borrow on `sess` from
        // `SessionOutputs` must end before the next `sess.run()` call.
        let mini_scores: Vec<f32> = {
            let outputs = if has_token_type_ids {
                sess.run(inputs![
                    "input_ids" => input_ids_tensor,
                    "attention_mask" => attention_mask_tensor,
                    "token_type_ids" => token_type_ids_tensor
                ])
                .map_err(|e: ort::Error| InferenceError::InferenceError(e.to_string()))?
            } else {
                sess.run(inputs![
                    "input_ids" => input_ids_tensor,
                    "attention_mask" => attention_mask_tensor
                ])
                .map_err(|e: ort::Error| InferenceError::InferenceError(e.to_string()))?
            };

            // Extract logits — bge-reranker-base output shape is [batch_size, 1]
            let (out_shape, logits_slice) = outputs[0]
                .try_extract_tensor::<f32>()
                .map_err(|e| InferenceError::InferenceError(e.to_string()))?;

            if out_shape.is_empty() || out_shape[0] as usize != batch_size {
                warn!(
                    "Reranker output shape mismatch: expected [{}, 1], got {:?}",
                    batch_size, out_shape
                );
            }

            // Apply sigmoid → owned Vec<f32>; borrow on outputs/sess ends here.
            logits_slice.iter().map(|&logit| sigmoid(logit)).collect()
            // outputs drops here (in reverse declaration order: logits_slice, outputs)
        };

        let n_scores = mini_scores.len();
        if n_scores != batch_size {
            warn!(
                "Reranker score count mismatch: expected {}, got {}",
                batch_size, n_scores
            );
            let mut padded = mini_scores;
            padded.resize(batch_size, 0.5);
            all_scores.extend(padded);
        } else {
            all_scores.extend(mini_scores);
        }
    }
    // sess drops here, releasing the mutex

    Ok(all_scores)
}

/// Sigmoid activation: 1 / (1 + exp(-x))
#[inline]
fn sigmoid(x: f32) -> f32 {
    1.0 / (1.0 + (-x).exp())
}

/// Download tokenizer and ONNX model files for the reranker.
/// Reuses `EmbeddingEngine::download_hf_file_pub` for redirect-aware caching.
fn download_reranker_files(
    cache_dir: Option<String>,
) -> std::result::Result<(PathBuf, PathBuf), InferenceError> {
    let cache = match cache_dir {
        Some(dir) => {
            let p = PathBuf::from(dir);
            std::fs::create_dir_all(&p)
                .map_err(|e| InferenceError::ModelLoadError(format!("cache_dir create: {e}")))?;
            p
        }
        None => {
            let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
            PathBuf::from(home)
                .join(".cache")
                .join("huggingface")
                .join("dakera")
                .join(RERANKER_REPO_ID.replace('/', "--"))
        }
    };

    std::fs::create_dir_all(&cache)
        .map_err(|e| InferenceError::ModelLoadError(format!("create cache dir: {e}")))?;

    let files = [
        "tokenizer.json",
        "tokenizer_config.json",
        "special_tokens_map.json",
        RERANKER_ONNX_FILE,
    ];

    for filename in &files {
        EmbeddingEngine::download_hf_file_pub(RERANKER_REPO_ID, filename, &cache)
            .map_err(|e| InferenceError::HubError(format!("download {filename}: {e}")))?;
    }

    let tokenizer_path = cache.join("tokenizer.json");
    let onnx_path = cache.join(RERANKER_ONNX_FILE);
    Ok((tokenizer_path, onnx_path))
}

impl std::fmt::Debug for CrossEncoderEngine {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CrossEncoderEngine")
            .field("model", &RERANKER_REPO_ID)
            .field("pool_size", &self.sessions.len())
            .field("onnx_batch_size", &RERANKER_ONNX_BATCH_SIZE)
            .field(
                "active_requests",
                &self.active_requests.load(Ordering::Relaxed),
            )
            .field("max_concurrent", &RERANKER_MAX_CONCURRENT)
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── GPU env-var read path ────────────────────────────────────────────────

    #[test]
    fn test_use_gpu_default_is_false() {
        // Without DAKERA_USE_GPU set, use_gpu must resolve to false (CPU default).
        use std::sync::Mutex;
        static ENV_LOCK: Mutex<()> = Mutex::new(());
        let _guard = ENV_LOCK.lock().unwrap();
        unsafe { std::env::remove_var("DAKERA_USE_GPU") };
        let use_gpu = std::env::var("DAKERA_USE_GPU")
            .map(|v| v == "1")
            .unwrap_or(false);
        assert!(
            !use_gpu,
            "expected CPU default when DAKERA_USE_GPU is unset"
        );
    }

    #[test]
    fn test_use_gpu_enabled_when_env_var_is_1() {
        use std::sync::Mutex;
        static ENV_LOCK: Mutex<()> = Mutex::new(());
        let _guard = ENV_LOCK.lock().unwrap();
        unsafe { std::env::set_var("DAKERA_USE_GPU", "1") };
        let use_gpu = std::env::var("DAKERA_USE_GPU")
            .map(|v| v == "1")
            .unwrap_or(false);
        unsafe { std::env::remove_var("DAKERA_USE_GPU") };
        assert!(use_gpu, "expected GPU mode when DAKERA_USE_GPU=1");
    }

    #[test]
    fn test_use_gpu_not_enabled_for_other_values() {
        use std::sync::Mutex;
        static ENV_LOCK: Mutex<()> = Mutex::new(());
        let _guard = ENV_LOCK.lock().unwrap();
        for val in ["0", "true", "yes", "gpu", ""] {
            unsafe { std::env::set_var("DAKERA_USE_GPU", val) };
            let use_gpu = std::env::var("DAKERA_USE_GPU")
                .map(|v| v == "1")
                .unwrap_or(false);
            unsafe { std::env::remove_var("DAKERA_USE_GPU") };
            assert!(
                !use_gpu,
                "expected CPU when DAKERA_USE_GPU={val:?} (only '1' enables GPU)"
            );
        }
    }

    #[test]
    fn test_sigmoid() {
        assert!((sigmoid(0.0) - 0.5).abs() < 1e-6);
        assert!(sigmoid(10.0) > 0.99);
        assert!(sigmoid(-10.0) < 0.01);
    }

    #[test]
    fn test_chunk_count_exact() {
        // 64 passages / chunk_size=32 → exactly 2 chunks
        let passages: Vec<String> = (0..64).map(|i| format!("passage {i}")).collect();
        let chunks: Vec<Vec<String>> = passages
            .chunks(RERANKER_CHUNK_SIZE)
            .map(<[String]>::to_vec)
            .collect();
        assert_eq!(chunks.len(), 2);
        assert_eq!(chunks[0].len(), 32);
        assert_eq!(chunks[1].len(), 32);
    }

    #[test]
    fn test_chunk_count_remainder() {
        // 50 passages / chunk_size=32 → 2 chunks (32 + 18)
        let passages: Vec<String> = (0..50).map(|i| format!("passage {i}")).collect();
        let chunks: Vec<Vec<String>> = passages
            .chunks(RERANKER_CHUNK_SIZE)
            .map(<[String]>::to_vec)
            .collect();
        assert_eq!(chunks.len(), 2);
        assert_eq!(chunks[0].len(), 32);
        assert_eq!(chunks[1].len(), 18);
    }

    #[test]
    fn test_chunk_count_small_batch() {
        // 10 passages → single chunk, no splitting overhead
        let passages: Vec<String> = (0..10).map(|i| format!("passage {i}")).collect();
        let chunks: Vec<Vec<String>> = passages
            .chunks(RERANKER_CHUNK_SIZE)
            .map(<[String]>::to_vec)
            .collect();
        assert_eq!(chunks.len(), 1);
        assert_eq!(chunks[0].len(), 10);
    }

    #[test]
    fn test_chunk_order_preserved() {
        // Chunk splitting must preserve passage order for score reassembly.
        let passages: Vec<String> = (0..70).map(|i| format!("p{i:03}")).collect();
        let reassembled: Vec<String> = passages
            .chunks(RERANKER_CHUNK_SIZE)
            .flat_map(<[String]>::to_vec)
            .collect();
        assert_eq!(passages, reassembled);
    }

    #[test]
    fn test_pool_size_constant() {
        const { assert!(RERANKER_POOL_SIZE >= 1) };
        const { assert!(RERANKER_CHUNK_SIZE >= 1) };
    }

    #[test]
    fn test_max_concurrent_exceeds_pool_size() {
        // RERANKER_MAX_CONCURRENT must be strictly greater than RERANKER_POOL_SIZE
        // so that the pool can be utilised at all before the gate fires (DAK-5893).
        const { assert!(RERANKER_MAX_CONCURRENT > RERANKER_POOL_SIZE) };
        // Must also be reasonable — less than 20 so a 20-request burst gets shed.
        const { assert!(RERANKER_MAX_CONCURRENT < 20) };
    }

    #[test]
    fn test_active_guard_decrements() {
        let counter = Arc::new(AtomicUsize::new(1));
        {
            let _g = ActiveGuard(Arc::clone(&counter));
            assert_eq!(counter.load(Ordering::SeqCst), 1);
        }
        assert_eq!(counter.load(Ordering::SeqCst), 0);
    }

    #[test]
    fn test_round_robin_wraps() {
        let pool_len = RERANKER_POOL_SIZE;
        // Simulate 10 concurrent callers; each gets a different start_idx.
        // Verify no start_idx exceeds pool_len when used with modulo.
        for start in 0usize..10 {
            let idx = start % pool_len;
            assert!(idx < pool_len);
        }
    }

    // ── ONNX mini-batch tests (DAK-5883) ────────────────────────────────────

    #[test]
    fn test_onnx_batch_size_constant_invariants() {
        // ONNX batch size must be positive and no larger than the chunk size.
        // If ONNX_BATCH_SIZE > CHUNK_SIZE the inner loop always produces one
        // mini-batch (identical to pre-DAK-5883 behaviour), which is allowed
        // but defeats the purpose of the constant.
        const { assert!(RERANKER_ONNX_BATCH_SIZE >= 1) };
        const { assert!(RERANKER_ONNX_BATCH_SIZE <= RERANKER_CHUNK_SIZE) };
    }

    #[test]
    fn test_onnx_mini_batch_count_full_chunk() {
        // A full chunk (32 passages) with ONNX_BATCH_SIZE=16 → exactly 2 mini-batches.
        let passages: Vec<String> = (0..RERANKER_CHUNK_SIZE).map(|i| format!("p{i}")).collect();
        let mini_batches: Vec<&[String]> = passages.chunks(RERANKER_ONNX_BATCH_SIZE).collect();
        let expected = RERANKER_CHUNK_SIZE.div_ceil(RERANKER_ONNX_BATCH_SIZE);
        assert_eq!(mini_batches.len(), expected);
        // Each full mini-batch has exactly ONNX_BATCH_SIZE items.
        for mb in &mini_batches[..mini_batches.len() - 1] {
            assert_eq!(mb.len(), RERANKER_ONNX_BATCH_SIZE);
        }
    }

    #[test]
    fn test_onnx_mini_batch_count_partial_chunk() {
        // ONNX_BATCH_SIZE + 1 passages → 2 mini-batches (full + remainder of 1).
        let n = RERANKER_ONNX_BATCH_SIZE + 1;
        let passages: Vec<String> = (0..n).map(|i| format!("p{i}")).collect();
        let mini_batches: Vec<&[String]> = passages.chunks(RERANKER_ONNX_BATCH_SIZE).collect();
        assert_eq!(mini_batches.len(), 2);
        assert_eq!(mini_batches[0].len(), RERANKER_ONNX_BATCH_SIZE);
        assert_eq!(mini_batches[1].len(), 1);
    }

    #[test]
    fn test_onnx_mini_batch_count_smaller_than_batch_size() {
        // Fewer passages than ONNX_BATCH_SIZE → single mini-batch (no overhead).
        let n = RERANKER_ONNX_BATCH_SIZE / 2;
        let passages: Vec<String> = (0..n).map(|i| format!("p{i}")).collect();
        let mini_batches: Vec<&[String]> = passages.chunks(RERANKER_ONNX_BATCH_SIZE).collect();
        assert_eq!(mini_batches.len(), 1);
        assert_eq!(mini_batches[0].len(), n);
    }

    #[test]
    fn test_onnx_mini_batch_order_preserved() {
        // Mini-batch splitting and reassembly must preserve input order exactly.
        // This guards against score[i] ↔ passage[j] mismatches in reassembly.
        let passages: Vec<String> = (0..70).map(|i| format!("p{i:03}")).collect();
        let reassembled: Vec<String> = passages
            .chunks(RERANKER_ONNX_BATCH_SIZE)
            .flat_map(|mb| mb.to_vec())
            .collect();
        assert_eq!(passages, reassembled);
    }

    #[test]
    fn test_onnx_mini_batch_total_score_count_matches_input() {
        // Whatever the input size, total score count after reassembly must equal
        // the number of input passages (covers exact multiples and remainders).
        for n in [1, 8, 15, 16, 17, 32, 33, 47, 64] {
            let passages: Vec<String> = (0..n).map(|i| format!("p{i}")).collect();
            let total: usize = passages
                .chunks(RERANKER_ONNX_BATCH_SIZE)
                .map(|mb| mb.len())
                .sum();
            assert_eq!(total, n, "score count mismatch for n={n}");
        }
    }

    #[test]
    fn test_onnx_batch_size_accessor() {
        // Verify that onnx_batch_size() returns the compile-time constant.
        // Requires constructing a CrossEncoderEngine — only possible in integration
        // tests with a real model. Instead, verify the constant directly via the
        // public constant's value relationship.
        assert_eq!(RERANKER_ONNX_BATCH_SIZE, 16);
    }
}