klieo-memory-qdrant 3.3.0

//! Qdrant connection configuration + bootstrap helpers.

use crate::error::store_err;
use dashmap::DashMap;
use klieo_core::error::MemoryError;
use qdrant_client::{
    qdrant::{CreateCollectionBuilder, Distance, VectorParamsBuilder},
    Qdrant,
};
use secrecy::{ExposeSecret, SecretString};
use std::sync::Arc;
use tokio::sync::OnceCell;

/// Configuration for [`MemoryQdrant`](crate::MemoryQdrant).
///
/// `api_key` is a [`SecretString`] (W2.A9): `Debug` redacts the value,
/// `Drop` zeroises it. Operators can pass any `Into<SecretString>` (a
/// `String` or another `SecretString`); `expose_secret()` is only
/// invoked at the `Qdrant::from_url` builder boundary.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct QdrantConfig {
    url: String,
    api_key: Option<SecretString>,
    /// Per-deployment collection-name prefix. Default: `"klieo_facts"`.
    collection_prefix: String,
    /// Operator opt-in to allow plaintext `http://` to non-loopback
    /// hosts (W2.A11 / CWE-319). Default false — `QdrantHandle::connect`
    /// refuses such URLs because the gRPC handshake transmits the
    /// `api-key` header before any TLS negotiation.
    allow_plaintext_remote: bool,
    /// Identifier of the embedder this store was provisioned for. M2
    /// `FilterableLongTermMemory::embedder_id()` surfaces this so
    /// `GraphAwareLongTerm` can hard-fail on cross-embedder queries
    /// (silent semantic drift is far worse than a typed error at the
    /// seam). Default: `"default"`.
    embedder_id: String,
}

impl QdrantConfig {
    /// Build a config from a Qdrant gRPC endpoint URL.
    pub fn new(url: impl Into<String>) -> Self {
        Self {
            url: url.into(),
            api_key: None,
            collection_prefix: "klieo_facts".to_string(),
            allow_plaintext_remote: false,
            embedder_id: "default".to_string(),
        }
    }

    /// Set the API key used in `api-key` headers.
    pub fn with_api_key(mut self, key: impl Into<SecretString>) -> Self {
        self.api_key = Some(key.into());
        self
    }

    /// Override the collection-name prefix (default `klieo_facts`).
    pub fn with_collection_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.collection_prefix = prefix.into();
        self
    }

    /// Explicitly allow `http://` to non-loopback hosts. Use only when
    /// Qdrant runs inside a sealed pod network (mesh-mTLS) or behind a
    /// TLS-terminating sidecar. Otherwise the api-key + collection
    /// payloads transit in cleartext (W2.A11).
    pub fn allow_plaintext_remote(mut self) -> Self {
        self.allow_plaintext_remote = true;
        self
    }

    /// Set the embedder identifier this store is provisioned for. M2
    /// `GraphAwareLongTerm` calls `recall_filtered_checked` with the
    /// caller's embedder id and hard-fails on mismatch — silent
    /// semantic drift across embedder versions is far worse than a
    /// typed error at the seam.
    pub fn with_embedder_id(mut self, id: impl Into<String>) -> Self {
        self.embedder_id = id.into();
        self
    }

    /// Embedder identifier configured for this store. Returns
    /// `"default"` unless overridden via [`Self::with_embedder_id`].
    pub fn embedder_id(&self) -> &str {
        &self.embedder_id
    }

    pub(crate) fn url(&self) -> &str {
        &self.url
    }

    pub(crate) fn api_key(&self) -> Option<&str> {
        self.api_key
            .as_ref()
            .map(|s: &SecretString| s.expose_secret())
    }

    pub(crate) fn collection_prefix(&self) -> &str {
        &self.collection_prefix
    }
}

/// Connected Qdrant handle, cheap to clone via `Arc`.
///
/// `ensure_collection` is cached per-process via a `DashMap` of
/// `(collection, dim)` → [`OnceCell<()>`]. The first call for a given
/// key performs the gRPC `collection_exists` + `create_collection`
/// dance; every subsequent call short-circuits on the cached `Ok(())`.
/// Concurrent first-callers share one in-flight `OnceCell::get_or_try_init`
/// future — no duplicate `create_collection` RPCs.
#[derive(Clone)]
pub(crate) struct QdrantHandle {
    pub(crate) client: Arc<Qdrant>,
    pub(crate) collection_prefix: String,
    bootstrapped: Arc<DashMap<String, Arc<OnceCell<()>>>>,
}

impl QdrantHandle {
    pub(crate) fn connect(cfg: &QdrantConfig) -> Result<Self, MemoryError> {
        // W2.A11 / CWE-319: refuse plaintext http:// to non-loopback
        // hosts unless the operator opted in via
        // `QdrantConfig::allow_plaintext_remote()`. Without this guard
        // the gRPC handshake transmits the `api-key` header in
        // cleartext to remote Qdrant clusters.
        check_plaintext_scheme(cfg.url(), cfg.allow_plaintext_remote)?;
        let mut builder = Qdrant::from_url(cfg.url());
        if let Some(key) = cfg.api_key() {
            builder = builder.api_key(key.to_string());
        }
        let client = builder.build().map_err(store_err)?;
        Ok(Self {
            client: Arc::new(client),
            collection_prefix: cfg.collection_prefix().to_string(),
            bootstrapped: Arc::new(DashMap::new()),
        })
    }

    /// Idempotent collection bootstrap, cached per-process.
    ///
    /// On the first call for a given `(name, dim)` pair the helper
    /// issues `collection_exists` + (if absent) `create_collection`.
    /// Subsequent calls short-circuit on the cached `OnceCell` —
    /// no gRPC round-trip. Bootstrap failures propagate; the
    /// `OnceCell` is left empty so the next call retries naturally.
    pub(crate) async fn ensure_collection(
        &self,
        name: &str,
        vector_dim: u64,
    ) -> Result<(), MemoryError> {
        let key = format!("{name}/{vector_dim}");
        let cell = self
            .bootstrapped
            .entry(key)
            .or_insert_with(|| Arc::new(OnceCell::new()))
            .clone();
        cell.get_or_try_init(|| self.bootstrap_collection(name, vector_dim))
            .await?;
        Ok(())
    }

    async fn bootstrap_collection(&self, name: &str, vector_dim: u64) -> Result<(), MemoryError> {
        let exists = self
            .client
            .collection_exists(name)
            .await
            .map_err(store_err)?;
        if exists {
            // An existing collection's vector dimension is fixed at creation. If
            // the configured embedder now produces a different dimension (a model
            // swap), writes would silently mismatch — fail closed instead.
            return self.verify_existing_dim(name, vector_dim).await;
        }
        let req = CreateCollectionBuilder::new(name)
            .vectors_config(VectorParamsBuilder::new(vector_dim, Distance::Cosine));
        self.client
            .create_collection(req)
            .await
            .map_err(store_err)?;
        Ok(())
    }

    /// Fail closed when an existing single-vector collection's declared
    /// dimension differs from `expected_dim`. Named-vector collections (not
    /// created by this crate) and incomplete configs are left unvalidated.
    async fn verify_existing_dim(&self, name: &str, expected_dim: u64) -> Result<(), MemoryError> {
        let info = self.client.collection_info(name).await.map_err(store_err)?;
        let actual = info
            .result
            .and_then(|i| i.config)
            .and_then(|c| c.params)
            .and_then(|p| p.vectors_config)
            .and_then(|vc| vc.config)
            .and_then(single_vector_dim);
        match actual {
            Some(actual) if actual != expected_dim => Err(MemoryError::Embedding(format!(
                "qdrant collection `{name}` has dimension {actual}, but the configured \
                 embedder produces {expected_dim}-dim vectors; a model change needs a new \
                 collection or a re-index"
            ))),
            _ => Ok(()),
        }
    }
}

/// Dimension of a single-vector collection config; `None` for named-vector
/// (`ParamsMap`) collections, which this crate never creates and so leaves
/// unvalidated.
fn single_vector_dim(config: qdrant_client::qdrant::vectors_config::Config) -> Option<u64> {
    use qdrant_client::qdrant::vectors_config::Config;
    match config {
        Config::Params(params) => Some(params.size),
        Config::ParamsMap(_) => None,
    }
}

/// Returns Err when `url` is `http://` to a non-loopback host and the
/// caller did not set `QdrantConfig::allow_plaintext_remote()`. W2.A11.
fn check_plaintext_scheme(url: &str, allow_plaintext_remote: bool) -> Result<(), MemoryError> {
    if !url.starts_with("http://") {
        return Ok(());
    }
    if allow_plaintext_remote {
        return Ok(());
    }
    let after_scheme = url.trim_start_matches("http://");
    let authority = after_scheme.split('/').next().unwrap_or("");
    // IPv6 hosts are bracketed: `[::1]:6333` → host=`[::1]`. IPv4/dns
    // hosts have no brackets: `localhost:6333` → host=`localhost`.
    let host = if let Some(end) = authority.find(']') {
        &authority[..=end]
    } else {
        authority.split(':').next().unwrap_or("")
    };
    if matches!(host, "localhost" | "127.0.0.1" | "[::1]" | "::1") {
        return Ok(());
    }
    Err(MemoryError::Store(format!(
        "QdrantConfig URL `{url}` is plaintext http:// to a non-loopback host; \
         the api-key is sent in cleartext. Use https://, switch to a loopback \
         binding, or set QdrantConfig::allow_plaintext_remote() if a sealed \
         pod network (mTLS/sidecar) handles the encryption."
    )))
}

#[cfg(test)]
mod tests {
    use super::*;
    use qdrant_client::qdrant::vectors_config::Config;
    use qdrant_client::qdrant::{VectorParams, VectorParamsMap};
    use std::sync::atomic::{AtomicUsize, Ordering};

    #[test]
    fn single_vector_dim_reads_unnamed_vector_size() {
        let config = Config::Params(VectorParams {
            size: 768,
            ..Default::default()
        });
        assert_eq!(single_vector_dim(config), Some(768));
    }

    #[test]
    fn single_vector_dim_skips_named_vector_collections() {
        let config = Config::ParamsMap(VectorParamsMap {
            map: Default::default(),
        });
        assert_eq!(single_vector_dim(config), None);
    }

    /// Asserts the `OnceCell` semantics independently of Qdrant:
    /// N concurrent `get_or_try_init` calls run the inner closure
    /// exactly once.
    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn once_cell_runs_init_exactly_once_across_concurrent_callers() {
        let counter = Arc::new(AtomicUsize::new(0));
        let cell: Arc<OnceCell<()>> = Arc::new(OnceCell::new());

        let mut handles = Vec::new();
        for _ in 0..16 {
            let cell = cell.clone();
            let counter = counter.clone();
            handles.push(tokio::spawn(async move {
                cell.get_or_try_init(|| async {
                    counter.fetch_add(1, Ordering::SeqCst);
                    tokio::task::yield_now().await;
                    Ok::<_, MemoryError>(())
                })
                .await
                .unwrap();
            }));
        }
        for handle in handles {
            handle.await.unwrap();
        }
        assert_eq!(counter.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn bootstrapped_keys_are_collection_plus_dim() {
        let key_a = format!("{}/{}", "klieo_facts:dev", 768u64);
        let key_b = format!("{}/{}", "klieo_facts:dev", 1024u64);
        assert_ne!(key_a, key_b);
    }

    /// W2.A11 / CWE-319: refuse plaintext to a non-loopback host
    /// unless the operator opts in.
    #[test]
    fn check_plaintext_scheme_rejects_remote_http() {
        assert!(check_plaintext_scheme("http://qdrant.example.com:6333", false).is_err());
        assert!(check_plaintext_scheme("http://10.0.0.1:6333", false).is_err());
    }

    #[test]
    fn check_plaintext_scheme_allows_loopback_http() {
        assert!(check_plaintext_scheme("http://localhost:6333", false).is_ok());
        assert!(check_plaintext_scheme("http://127.0.0.1:6333", false).is_ok());
        assert!(check_plaintext_scheme("http://[::1]:6333", false).is_ok());
    }

    #[test]
    fn check_plaintext_scheme_allows_https() {
        assert!(check_plaintext_scheme("https://qdrant.example.com:6333", false).is_ok());
    }

    #[test]
    fn check_plaintext_scheme_honours_opt_in() {
        assert!(check_plaintext_scheme("http://qdrant.example.com:6333", true).is_ok());
    }
}