Skip to main content

klieo_memory_qdrant/
client.rs

1//! Qdrant connection configuration + bootstrap helpers.
2
3use crate::error::store_err;
4use dashmap::DashMap;
5use klieo_core::error::MemoryError;
6use qdrant_client::{
7    qdrant::{CreateCollectionBuilder, Distance, VectorParamsBuilder},
8    Qdrant,
9};
10use secrecy::{ExposeSecret, SecretString};
11use std::sync::Arc;
12use tokio::sync::OnceCell;
13
14/// Configuration for [`MemoryQdrant`](crate::MemoryQdrant).
15///
16/// `api_key` is a [`SecretString`] (W2.A9): `Debug` redacts the value,
17/// `Drop` zeroises it. Operators can pass any `Into<SecretString>` (a
18/// `String` or another `SecretString`); `expose_secret()` is only
19/// invoked at the `Qdrant::from_url` builder boundary.
20#[derive(Clone, Debug)]
21#[non_exhaustive]
22pub struct QdrantConfig {
23    url: String,
24    api_key: Option<SecretString>,
25    /// Per-deployment collection-name prefix. Default: `"klieo_facts"`.
26    collection_prefix: String,
27    /// Operator opt-in to allow plaintext `http://` to non-loopback
28    /// hosts (W2.A11 / CWE-319). Default false — `QdrantHandle::connect`
29    /// refuses such URLs because the gRPC handshake transmits the
30    /// `api-key` header before any TLS negotiation.
31    allow_plaintext_remote: bool,
32    /// Identifier of the embedder this store was provisioned for. M2
33    /// `FilterableLongTermMemory::embedder_id()` surfaces this so
34    /// `GraphAwareLongTerm` can hard-fail on cross-embedder queries
35    /// (silent semantic drift is far worse than a typed error at the
36    /// seam). Default: `"default"`.
37    embedder_id: String,
38}
39
40impl QdrantConfig {
41    /// Build a config from a Qdrant gRPC endpoint URL.
42    pub fn new(url: impl Into<String>) -> Self {
43        Self {
44            url: url.into(),
45            api_key: None,
46            collection_prefix: "klieo_facts".to_string(),
47            allow_plaintext_remote: false,
48            embedder_id: "default".to_string(),
49        }
50    }
51
52    /// Set the API key used in `api-key` headers.
53    pub fn with_api_key(mut self, key: impl Into<SecretString>) -> Self {
54        self.api_key = Some(key.into());
55        self
56    }
57
58    /// Override the collection-name prefix (default `klieo_facts`).
59    pub fn with_collection_prefix(mut self, prefix: impl Into<String>) -> Self {
60        self.collection_prefix = prefix.into();
61        self
62    }
63
64    /// Explicitly allow `http://` to non-loopback hosts. Use only when
65    /// Qdrant runs inside a sealed pod network (mesh-mTLS) or behind a
66    /// TLS-terminating sidecar. Otherwise the api-key + collection
67    /// payloads transit in cleartext (W2.A11).
68    pub fn allow_plaintext_remote(mut self) -> Self {
69        self.allow_plaintext_remote = true;
70        self
71    }
72
73    /// Set the embedder identifier this store is provisioned for. M2
74    /// `GraphAwareLongTerm` calls `recall_filtered_checked` with the
75    /// caller's embedder id and hard-fails on mismatch — silent
76    /// semantic drift across embedder versions is far worse than a
77    /// typed error at the seam.
78    pub fn with_embedder_id(mut self, id: impl Into<String>) -> Self {
79        self.embedder_id = id.into();
80        self
81    }
82
83    /// Embedder identifier configured for this store. Returns
84    /// `"default"` unless overridden via [`Self::with_embedder_id`].
85    pub fn embedder_id(&self) -> &str {
86        &self.embedder_id
87    }
88
89    pub(crate) fn url(&self) -> &str {
90        &self.url
91    }
92
93    pub(crate) fn api_key(&self) -> Option<&str> {
94        self.api_key
95            .as_ref()
96            .map(|s: &SecretString| s.expose_secret())
97    }
98
99    pub(crate) fn collection_prefix(&self) -> &str {
100        &self.collection_prefix
101    }
102}
103
104/// Connected Qdrant handle, cheap to clone via `Arc`.
105///
106/// `ensure_collection` is cached per-process via a `DashMap` of
107/// `(collection, dim)` → [`OnceCell<()>`]. The first call for a given
108/// key performs the gRPC `collection_exists` + `create_collection`
109/// dance; every subsequent call short-circuits on the cached `Ok(())`.
110/// Concurrent first-callers share one in-flight `OnceCell::get_or_try_init`
111/// future — no duplicate `create_collection` RPCs.
112#[derive(Clone)]
113pub(crate) struct QdrantHandle {
114    pub(crate) client: Arc<Qdrant>,
115    pub(crate) collection_prefix: String,
116    bootstrapped: Arc<DashMap<String, Arc<OnceCell<()>>>>,
117}
118
119impl QdrantHandle {
120    pub(crate) fn connect(cfg: &QdrantConfig) -> Result<Self, MemoryError> {
121        // W2.A11 / CWE-319: refuse plaintext http:// to non-loopback
122        // hosts unless the operator opted in via
123        // `QdrantConfig::allow_plaintext_remote()`. Without this guard
124        // the gRPC handshake transmits the `api-key` header in
125        // cleartext to remote Qdrant clusters.
126        check_plaintext_scheme(cfg.url(), cfg.allow_plaintext_remote)?;
127        let mut builder = Qdrant::from_url(cfg.url());
128        if let Some(key) = cfg.api_key() {
129            builder = builder.api_key(key.to_string());
130        }
131        let client = builder.build().map_err(store_err)?;
132        Ok(Self {
133            client: Arc::new(client),
134            collection_prefix: cfg.collection_prefix().to_string(),
135            bootstrapped: Arc::new(DashMap::new()),
136        })
137    }
138
139    /// Idempotent collection bootstrap, cached per-process.
140    ///
141    /// On the first call for a given `(name, dim)` pair the helper
142    /// issues `collection_exists` + (if absent) `create_collection`.
143    /// Subsequent calls short-circuit on the cached `OnceCell` —
144    /// no gRPC round-trip. Bootstrap failures propagate; the
145    /// `OnceCell` is left empty so the next call retries naturally.
146    pub(crate) async fn ensure_collection(
147        &self,
148        name: &str,
149        vector_dim: u64,
150    ) -> Result<(), MemoryError> {
151        let key = format!("{name}/{vector_dim}");
152        let cell = self
153            .bootstrapped
154            .entry(key)
155            .or_insert_with(|| Arc::new(OnceCell::new()))
156            .clone();
157        cell.get_or_try_init(|| self.bootstrap_collection(name, vector_dim))
158            .await?;
159        Ok(())
160    }
161
162    async fn bootstrap_collection(&self, name: &str, vector_dim: u64) -> Result<(), MemoryError> {
163        let exists = self
164            .client
165            .collection_exists(name)
166            .await
167            .map_err(store_err)?;
168        if exists {
169            // An existing collection's vector dimension is fixed at creation. If
170            // the configured embedder now produces a different dimension (a model
171            // swap), writes would silently mismatch — fail closed instead.
172            return self.verify_existing_dim(name, vector_dim).await;
173        }
174        let req = CreateCollectionBuilder::new(name)
175            .vectors_config(VectorParamsBuilder::new(vector_dim, Distance::Cosine));
176        self.client
177            .create_collection(req)
178            .await
179            .map_err(store_err)?;
180        Ok(())
181    }
182
183    /// Fail closed when an existing single-vector collection's declared
184    /// dimension differs from `expected_dim`. Named-vector collections (not
185    /// created by this crate) and incomplete configs are left unvalidated.
186    async fn verify_existing_dim(&self, name: &str, expected_dim: u64) -> Result<(), MemoryError> {
187        let info = self.client.collection_info(name).await.map_err(store_err)?;
188        let actual = info
189            .result
190            .and_then(|i| i.config)
191            .and_then(|c| c.params)
192            .and_then(|p| p.vectors_config)
193            .and_then(|vc| vc.config)
194            .and_then(single_vector_dim);
195        match actual {
196            Some(actual) if actual != expected_dim => Err(MemoryError::Embedding(format!(
197                "qdrant collection `{name}` has dimension {actual}, but the configured \
198                 embedder produces {expected_dim}-dim vectors; a model change needs a new \
199                 collection or a re-index"
200            ))),
201            _ => Ok(()),
202        }
203    }
204}
205
206/// Dimension of a single-vector collection config; `None` for named-vector
207/// (`ParamsMap`) collections, which this crate never creates and so leaves
208/// unvalidated.
209fn single_vector_dim(config: qdrant_client::qdrant::vectors_config::Config) -> Option<u64> {
210    use qdrant_client::qdrant::vectors_config::Config;
211    match config {
212        Config::Params(params) => Some(params.size),
213        Config::ParamsMap(_) => None,
214    }
215}
216
217/// Returns Err when `url` is `http://` to a non-loopback host and the
218/// caller did not set `QdrantConfig::allow_plaintext_remote()`. W2.A11.
219fn check_plaintext_scheme(url: &str, allow_plaintext_remote: bool) -> Result<(), MemoryError> {
220    if !url.starts_with("http://") {
221        return Ok(());
222    }
223    if allow_plaintext_remote {
224        return Ok(());
225    }
226    let after_scheme = url.trim_start_matches("http://");
227    let authority = after_scheme.split('/').next().unwrap_or("");
228    // IPv6 hosts are bracketed: `[::1]:6333` → host=`[::1]`. IPv4/dns
229    // hosts have no brackets: `localhost:6333` → host=`localhost`.
230    let host = if let Some(end) = authority.find(']') {
231        &authority[..=end]
232    } else {
233        authority.split(':').next().unwrap_or("")
234    };
235    if matches!(host, "localhost" | "127.0.0.1" | "[::1]" | "::1") {
236        return Ok(());
237    }
238    Err(MemoryError::Store(format!(
239        "QdrantConfig URL `{url}` is plaintext http:// to a non-loopback host; \
240         the api-key is sent in cleartext. Use https://, switch to a loopback \
241         binding, or set QdrantConfig::allow_plaintext_remote() if a sealed \
242         pod network (mTLS/sidecar) handles the encryption."
243    )))
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249    use qdrant_client::qdrant::vectors_config::Config;
250    use qdrant_client::qdrant::{VectorParams, VectorParamsMap};
251    use std::sync::atomic::{AtomicUsize, Ordering};
252
253    #[test]
254    fn single_vector_dim_reads_unnamed_vector_size() {
255        let config = Config::Params(VectorParams {
256            size: 768,
257            ..Default::default()
258        });
259        assert_eq!(single_vector_dim(config), Some(768));
260    }
261
262    #[test]
263    fn single_vector_dim_skips_named_vector_collections() {
264        let config = Config::ParamsMap(VectorParamsMap {
265            map: Default::default(),
266        });
267        assert_eq!(single_vector_dim(config), None);
268    }
269
270    /// Asserts the `OnceCell` semantics independently of Qdrant:
271    /// N concurrent `get_or_try_init` calls run the inner closure
272    /// exactly once.
273    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
274    async fn once_cell_runs_init_exactly_once_across_concurrent_callers() {
275        let counter = Arc::new(AtomicUsize::new(0));
276        let cell: Arc<OnceCell<()>> = Arc::new(OnceCell::new());
277
278        let mut handles = Vec::new();
279        for _ in 0..16 {
280            let cell = cell.clone();
281            let counter = counter.clone();
282            handles.push(tokio::spawn(async move {
283                cell.get_or_try_init(|| async {
284                    counter.fetch_add(1, Ordering::SeqCst);
285                    tokio::task::yield_now().await;
286                    Ok::<_, MemoryError>(())
287                })
288                .await
289                .unwrap();
290            }));
291        }
292        for handle in handles {
293            handle.await.unwrap();
294        }
295        assert_eq!(counter.load(Ordering::SeqCst), 1);
296    }
297
298    #[tokio::test]
299    async fn bootstrapped_keys_are_collection_plus_dim() {
300        let key_a = format!("{}/{}", "klieo_facts:dev", 768u64);
301        let key_b = format!("{}/{}", "klieo_facts:dev", 1024u64);
302        assert_ne!(key_a, key_b);
303    }
304
305    /// W2.A11 / CWE-319: refuse plaintext to a non-loopback host
306    /// unless the operator opts in.
307    #[test]
308    fn check_plaintext_scheme_rejects_remote_http() {
309        assert!(check_plaintext_scheme("http://qdrant.example.com:6333", false).is_err());
310        assert!(check_plaintext_scheme("http://10.0.0.1:6333", false).is_err());
311    }
312
313    #[test]
314    fn check_plaintext_scheme_allows_loopback_http() {
315        assert!(check_plaintext_scheme("http://localhost:6333", false).is_ok());
316        assert!(check_plaintext_scheme("http://127.0.0.1:6333", false).is_ok());
317        assert!(check_plaintext_scheme("http://[::1]:6333", false).is_ok());
318    }
319
320    #[test]
321    fn check_plaintext_scheme_allows_https() {
322        assert!(check_plaintext_scheme("https://qdrant.example.com:6333", false).is_ok());
323    }
324
325    #[test]
326    fn check_plaintext_scheme_honours_opt_in() {
327        assert!(check_plaintext_scheme("http://qdrant.example.com:6333", true).is_ok());
328    }
329}