Files
Kigi-CLI/crates/codegen/kigi-memory/src/backend.rs
T
ZacharyZhang-NY ed8049cf77 fix(memory): scope embedding credentials to the endpoint that may receive them
`MemoryBackendParams` carried the primary session `AuthManager` and the
session api-key provider unconditionally, while `embed_base_url` is the
CURRENT MODEL's endpoint and `AuthRetryMiddleware` stamps `Authorization`
on every request it wraps. A user who enables `[memory.embedding] model`
while running a BYOK or subscription-OAuth model therefore sent the Kimi
session bearer to that third party.

`EndpointScopedCredentials` binds the credential to the one endpoint it
may reach: `for_endpoint` drops the handle unless the caller vouches for
the URL, and `approved_for` re-checks at provider-build time in release
too, because `MemoryBackendParams` is `Clone` and callers rewrite fields
on the copy.

The shell decides through `CredentialAuthority::manager_for` rather than
a second URL predicate — it answers both whether a credential may ride
and which manager governs it, so a subscription-OAuth platform gets its
own pooled manager. The session's `SharedApiKeyProvider` is not
forwarded at all: it is hard-wired to the primary manager, so at a
pooled platform's host it would resolve the wrong bearer. A platform's
own `embed_api_key` is untouched and keeps serving its own endpoint.

The background reindex built a second provider straight from
`ApiEmbeddingProvider::from_session`, outside the chokepoint and without
401 refresh; it now embeds through the session's own params.

Test strength verified by mutation: with the guard reverted, exactly
`session_credentials_are_withheld_from_a_foreign_endpoint` and
`a_cloned_param_set_cannot_redirect_scoped_credentials` fail.
2026-07-26 23:41:45 -04:00

1480 lines
57 KiB
Rust

//! Concrete `MemoryBackend` implementation using hybrid search.
//!
//! `MemoryBackendImpl` combines FTS5 keyword search with optional vector
//! KNN similarity via `hybrid_search()`. When embeddings are available
//! (embedding config + API key), the query is vectorized and both signals
//! are merged with recency and source weights. When embeddings are
//! unavailable, gracefully degrades to FTS-only.
//!
//! `rusqlite::Connection` is `!Send + !Sync`, so we open a fresh `MemoryIndex`
//! per query. WAL mode ensures concurrent readers don't block.
use std::path::{Path, PathBuf};
use std::sync::Arc;
use kigi_tools::types::memory_backend::{MemoryBackend, MemorySearchResult};
use super::embedding::EmbeddingProvider as _;
use super::storage::MemoryStorage;
use super::watcher::MemoryFileWatcher;
/// The session's embedding credentials, bound to the one endpoint they may
/// reach. Only [`Self::for_endpoint`] retains a live handle; the default fails
/// closed.
///
/// `embed_base_url` is the CURRENT MODEL's endpoint, so a session on a BYOK or
/// subscription-OAuth model aims memory embeddings at that provider's host —
/// and [`kigi_auth::AuthRetryMiddleware`] stamps `Authorization` on every
/// request it wraps, with no idea where the request is going. The caller that
/// owns the credential rule decides `trusted` once, here; a platform's own
/// `embed_api_key` is unaffected and keeps serving its own endpoint.
#[derive(Clone, Default)]
pub struct EndpointScopedCredentials {
endpoint: Option<reqwest::Url>,
auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
}
// Redacts the credential handles; only their presence is printable.
impl std::fmt::Debug for EndpointScopedCredentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EndpointScopedCredentials")
.field("endpoint", &self.endpoint)
.field("has_auth_credentials", &self.auth_credentials.is_some())
.finish()
}
}
impl EndpointScopedCredentials {
/// No session credential may ride — the state every non-session caller wants.
pub fn none() -> Self {
Self::default()
}
/// Retains the handles only for a `trusted`, parsable `endpoint`.
pub fn for_endpoint(
endpoint: &str,
trusted: bool,
auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
) -> Self {
if trusted && let Ok(url) = reqwest::Url::parse(endpoint) {
return Self {
endpoint: Some(url),
auth_credentials,
};
}
if auth_credentials.is_some() {
tracing::info!(
target: kigi_log::memory_log::TARGET,
endpoint,
"memory embeddings: session credentials withheld from this endpoint; \
its own key, if any, still applies"
);
}
Self::none()
}
pub fn is_empty(&self) -> bool {
self.auth_credentials.is_none()
}
/// Enforced at request-build time, in release too: [`MemoryBackendParams`]
/// is `Clone` and callers rewrite fields on the copy, so construction-time
/// scoping alone would not survive a rewritten `embed_base_url`.
fn approved_for(&self, base_url: &str) -> bool {
match &self.endpoint {
None => self.is_empty(),
Some(endpoint) => reqwest::Url::parse(base_url).is_ok_and(|url| &url == endpoint),
}
}
}
/// All configuration needed to build a fully-wired [`MemoryBackendImpl`] for a live session.
///
/// Grouping these in one struct ensures every call site — ToolBridge, first-turn
/// injection, and post-compaction recovery — shares identical config. Without it,
/// different paths silently fell back to FTS-only search and ignored
/// `[memory.search]` config because no single place applied all builder methods.
#[derive(Clone)]
pub struct MemoryBackendParams {
/// Session ID for telemetry events.
pub session_id: String,
/// Embedding provider config — `None` forces FTS-only fallback everywhere.
pub embed_config: Option<kigi_config_types::MemoryEmbeddingConfig>,
/// Base URL for embedding API calls (CLI proxy). Must match the endpoint
/// `embedding_credentials` was scoped to; a mismatch fails closed.
pub embed_base_url: String,
/// API key for embedding API calls.
pub embed_api_key: Option<String>,
/// Hybrid search scoring config (weights, thresholds, decay, MMR).
pub search_config: kigi_config_types::MemorySearchConfig,
/// File watcher for sync-on-search — `None` disables external-edit detection.
pub watcher: Option<Arc<MemoryFileWatcher>>,
/// Seconds before a stale reindex claim is forcibly released.
pub stale_claim_secs: i64,
/// Telemetry label emitted with every search event from this backend.
///
/// Differentiates the three runtime search paths in dashboards and logs:
/// - `"tool"` — model-initiated `memory_search` tool call (ToolBridge)
/// - `"injection"` — first-turn memory context injection
/// - `"compaction_recovery"` — post-compaction context re-injection
pub search_source: &'static str,
/// The session credentials, and the single endpoint they may reach.
pub embedding_credentials: EndpointScopedCredentials,
}
impl MemoryBackendParams {
/// Async so `current_api_key_async` can drive the AuthManager
/// refresh chain; reindex loops outlive the OIDC TTL.
pub async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
build_embedding_provider(
self.embed_config.as_ref(),
&self.embedding_credentials,
self.embed_api_key.as_deref(),
&self.embed_base_url,
)
.await
}
}
async fn build_embedding_provider(
config: Option<&kigi_config_types::MemoryEmbeddingConfig>,
credentials: &EndpointScopedCredentials,
static_api_key: Option<&str>,
base_url: &str,
) -> Option<super::embedding::ApiEmbeddingProvider> {
let config = config?;
if config.model.as_ref().is_none_or(|m| m.is_empty()) {
return None;
}
let approved = credentials.approved_for(base_url);
if !approved {
tracing::error!(
target: kigi_log::memory_log::TARGET,
base_url,
approved_endpoint = ?credentials.endpoint,
"memory embeddings: scoped credentials do not match the request URL; dropping them"
);
}
// Prefer the refresh-capable credential provider — the middleware gives
// 401 retry for free without any per-call key resolution.
if approved && let Some(creds) = credentials.auth_credentials.as_ref() {
let client = super::embedding::build_middleware_client(creds.clone());
return super::embedding::ApiEmbeddingProvider::from_config(
config,
base_url.to_owned(),
client,
);
}
// The platform's own configured key, wrapped in a static middleware client
// (no 401 refresh, but the auth header is still stamped by middleware).
super::embedding::ApiEmbeddingProvider::from_session(
config,
base_url.to_owned(),
static_api_key?.to_owned(),
)
}
/// `MemoryBackend` implementation backed by hybrid search (FTS5 + vector KNN).
///
/// Stores only `Send + Sync` config data. The `MemoryIndex` and
/// `EmbeddingProvider` are constructed on demand per query.
pub struct MemoryBackendImpl {
db_path: PathBuf,
storage: MemoryStorage,
/// Embedding config — `None` disables vector search (FTS-only fallback).
embed_config: Option<kigi_config_types::MemoryEmbeddingConfig>,
/// API base URL for embedding requests (cli-chat-proxy).
embed_base_url: String,
/// API key for embedding requests.
embed_api_key: Option<String>,
/// Search scoring config (weights, min_score, max_results).
search_config: kigi_config_types::MemorySearchConfig,
/// File watcher for detecting external memory edits.
watcher: Option<Arc<MemoryFileWatcher>>,
/// Stale claim threshold for reindex coordination.
stale_claim_secs: i64,
/// Session ID for telemetry events.
session_id: String,
/// Telemetry label for search events — mirrors [`MemoryBackendParams::search_source`].
search_source: &'static str,
/// Shared search counter — read by session summary telemetry.
///
/// Only the ToolBridge backend's counter is shared back to the session actor;
/// injection and compaction-recovery backends use their own local counters.
pub search_counter: std::sync::Arc<std::sync::atomic::AtomicU64>,
/// The session credentials, and the single endpoint they may reach.
embedding_credentials: EndpointScopedCredentials,
}
impl MemoryBackendImpl {
/// Create a new backend. `db_path` must point to an existing SQLite
/// database created by `MemoryIndex::open_or_create()`.
pub fn new(db_path: PathBuf, storage: MemoryStorage) -> Self {
Self {
db_path,
storage,
embed_config: None,
embed_base_url: String::new(),
embed_api_key: None,
search_config: kigi_config_types::MemorySearchConfig::default(),
watcher: None,
stale_claim_secs: 60,
session_id: String::new(),
search_source: "tool",
embedding_credentials: EndpointScopedCredentials::none(),
search_counter: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)),
}
}
/// Set the session ID for telemetry.
pub fn with_session_id(mut self, session_id: String) -> Self {
self.session_id = session_id;
self
}
/// Configure the embedding provider for hybrid search.
///
/// Without this, `search()` falls back to FTS-only.
pub fn with_embedding(
mut self,
config: kigi_config_types::MemoryEmbeddingConfig,
base_url: String,
api_key: Option<String>,
) -> Self {
self.embed_config = Some(config);
self.embed_base_url = base_url;
self.embed_api_key = api_key;
self
}
/// Override the search scoring config (weights, limits, etc.).
pub fn with_search_config(mut self, config: kigi_config_types::MemorySearchConfig) -> Self {
self.search_config = config;
self
}
/// Attach a file watcher for sync-on-search (reindex dirty files before querying).
pub fn with_watcher(mut self, watcher: Arc<MemoryFileWatcher>, stale_claim_secs: i64) -> Self {
self.watcher = Some(watcher);
self.stale_claim_secs = stale_claim_secs;
self
}
/// Open a read-only connection for simple queries (`total_chunks`, `get`).
fn open_readonly(&self) -> Result<rusqlite::Connection, rusqlite::Error> {
// Journal-mode-aware open (busy_timeout included): never mmap a legacy
// WAL -shm on network mounts (SIGBUS); see JournalMode::open_readonly.
kigi_sqlite_journal::JournalMode::for_db_path(&self.db_path).open_readonly(&self.db_path)
}
async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
build_embedding_provider(
self.embed_config.as_ref(),
&self.embedding_credentials,
self.embed_api_key.as_deref(),
&self.embed_base_url,
)
.await
}
/// Build a fully configured backend for a live session.
///
/// Prefer this over calling `new()` + individual builder methods: it ensures
/// session_id, embeddings, search config, and the file watcher are applied
/// consistently at every call site (ToolBridge, first-turn injection,
/// post-compaction recovery). Using the factory eliminates the silent
/// per-site drift where some paths got hybrid search while others fell back
/// to FTS-only, and where `[memory.search]` config was effectively ignored.
pub fn from_session_params(storage: MemoryStorage, params: &MemoryBackendParams) -> Self {
let db_path = storage.workspace_dir().join("index.sqlite");
let mut backend = Self::new(db_path, storage)
.with_session_id(params.session_id.clone())
.with_search_config(params.search_config.clone());
backend.search_source = params.search_source;
if let Some(ec) = &params.embed_config {
backend = backend.with_embedding(
ec.clone(),
params.embed_base_url.clone(),
params.embed_api_key.clone(),
);
}
if let Some(w) = &params.watcher {
backend = backend.with_watcher(w.clone(), params.stale_claim_secs);
}
backend.embedding_credentials = params.embedding_credentials.clone();
backend
}
}
/// Test-only field accessors.
///
/// These expose private fields so tests can assert that `from_session_params`
/// actually stored the values it was given, without routing through a full
/// runtime search call whose semantics override some config fields.
#[cfg(test)]
impl MemoryBackendImpl {
/// Returns the session ID stored in this backend.
pub fn session_id_for_test(&self) -> &str {
&self.session_id
}
/// Returns the search config stored in this backend.
pub fn search_config_for_test(&self) -> &kigi_config_types::MemorySearchConfig {
&self.search_config
}
}
#[async_trait::async_trait]
impl MemoryBackend for MemoryBackendImpl {
#[tracing::instrument(name = "memory.search", skip_all, fields(
session_id = %self.session_id, max_results, min_score,
))]
async fn search(
&self,
query: &str,
max_results: usize,
min_score: f64,
) -> Result<Vec<MemorySearchResult>, Box<dyn std::error::Error + Send + Sync>> {
// Open a MemoryIndex for this query (open-per-query, ~1ms).
//
// IMPORTANT: `MemoryIndex` is `Send` but `!Sync`, so `&MemoryIndex`
// is `!Send`. To keep this future `Send`, we must never hold a
// `&index` borrow across an `.await` point. The code below is
// structured into sync phases (borrow &index) and async phases
// (no &index borrow) to satisfy this constraint.
let embed_dims = self.embed_config.as_ref().map_or(1024, |ec| ec.dimensions);
let mut index = super::index::MemoryIndex::open_or_create(
&self.db_path,
self.storage.clone(),
kigi_config_types::MemoryIndexConfig::default(),
embed_dims,
)
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> {
Box::new(std::io::Error::other(e.to_string()))
})?;
// Sync phase 1: reindex dirty files, collect chunks needing embeddings
let mut reindex_chunks: Vec<(String, String)> = Vec::new();
let mut needs_release = false;
if let Some(ref watcher) = self.watcher
&& watcher.is_dirty()
&& index.try_claim_reindex(self.stale_claim_secs)
{
needs_release = true;
let dirty_files = watcher.take_dirty();
let dirty_count = dirty_files.len();
for file in &dirty_files {
if file.exists() {
// File was created or modified — reindex it.
let source = self.storage.classify_source(file);
let _ = index.reindex_file(file, source);
} else {
// File was deleted — remove its stale chunks from the index so
// they are no longer searchable. Without this call, reindex_file
// returns early when the file is unreadable and leaves orphaned
// chunks behind indefinitely.
let _ = index.delete_path(file);
}
}
if dirty_count > 0 {
reindex_chunks = index.chunks_without_embeddings().unwrap_or_default();
}
}
// Async phase: embed missing chunks (no &index borrow)
let provider = self.make_embedding_provider().await;
if !reindex_chunks.is_empty()
&& let Some(ref provider) = provider
{
let mut upserts: Vec<(String, Vec<f32>)> = Vec::new();
for batch in reindex_chunks.chunks(32) {
let texts: Vec<&str> = batch.iter().map(|(_, t)| t.as_str()).collect();
match provider.embed_batch(&texts).await {
Ok(embeddings) => {
for ((chunk_id, _), emb) in batch.iter().zip(embeddings.into_iter()) {
upserts.push((chunk_id.clone(), emb));
}
}
Err(e) => {
tracing::warn!(
target: kigi_log::memory_log::TARGET,
error = %e,
"embedding batch failed during sync-on-search, skipping"
);
}
}
}
// Sync: upsert embeddings back (borrows &index, no await)
for (chunk_id, emb) in &upserts {
let _ = index.upsert_embedding(chunk_id, emb);
}
}
if needs_release {
index.release_claim();
}
// Sync phase 2: FTS search
let mut search_config = self.search_config.clone();
search_config.max_results = max_results;
search_config.min_score = min_score as f32;
let candidate_limit = search_config.max_results * 3;
let mut fts_results = index.search_fts(query, candidate_limit).unwrap_or_default();
// Supplemental evergreen query: ensure global/workspace MEMORY.md
// chunks appear in candidates even when session volume crowds them
// out of the base FTS results. Mirrors hybrid_search() in search.rs.
let evergreen = index
.search_fts_by_sources(query, candidate_limit, &["global", "workspace"])
.unwrap_or_default();
let existing: std::collections::HashSet<String> =
fts_results.iter().map(|r| r.chunk_id.clone()).collect();
for r in evergreen {
if !existing.contains(&r.chunk_id) {
fts_results.push(r);
}
}
let vec_available = index.vec_available() && provider.is_some();
// Async phase: embed query for vector search (no &index borrow)
let query_embedding = if vec_available {
if let Some(ref provider) = provider {
match provider.embed_batch(&[query]).await {
Ok(embeddings) if !embeddings.is_empty() => {
Some(embeddings.into_iter().next().unwrap())
}
Ok(_) => None,
Err(e) => {
tracing::warn!(error = %e, "embedding query failed, falling back to FTS-only");
None
}
}
} else {
None
}
} else {
None
};
// Sync phase 3: vector search + scoring + merge (borrows &index)
let results = super::search::hybrid_search_merge(
&index,
fts_results,
query_embedding.as_deref(),
&search_config,
)
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> {
Box::new(std::io::Error::other(e.to_string()))
})?;
// Record accesses for the returned chunks so access_count and
// last_accessed stay current. Non-fatal: a failed write is a no-op
// for the caller and does not affect the search response.
for result in &results {
let _ = index.record_access(&result.chunk_id);
}
self.search_counter
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
Ok(results
.into_iter()
.map(|r| MemorySearchResult {
chunk_id: r.chunk_id,
path: r.path,
start_line: r.start_line,
end_line: r.end_line,
score: r.score,
snippet: r.snippet,
source: r.source,
created_at: Some(r.created_at),
})
.collect())
}
fn get(
&self,
path: &str,
from: Option<usize>,
lines: Option<usize>,
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
Ok(self.storage.read_file(Path::new(path), from, lines)?)
}
fn total_chunks(&self) -> Result<usize, Box<dyn std::error::Error + Send + Sync>> {
let conn = self.open_readonly()?;
let count: i64 = conn.query_row("SELECT COUNT(*) FROM chunks", [], |r| r.get(0))?;
Ok(count as usize)
}
/// Return the configured `max_results` from the stored search config.
///
/// Overrides the trait default so the `memory_search` tool honours
/// `[memory.search].max_results` from config when the model does not
/// supply an explicit value.
fn default_search_max_results(&self) -> usize {
self.search_config.max_results
}
/// Return the configured `min_score` from the stored search config.
fn default_search_min_score(&self) -> f64 {
self.search_config.min_score as f64
}
}
#[cfg(test)]
mod factory_tests {
use super::*;
use crate::index::{MemoryIndex, init_sqlite_vec};
use crate::storage::MemoryStorage;
use kigi_config_types::{MemoryEmbeddingConfig, MemorySearchConfig};
use tempfile::TempDir;
fn make_storage(tmp: &TempDir) -> MemoryStorage {
let global = tmp.path().join("memory");
let workspace = global.join("test_ws");
MemoryStorage::with_paths(global, workspace)
}
fn make_params_fts_only(session_id: &str) -> MemoryBackendParams {
MemoryBackendParams {
session_id: session_id.to_string(),
embed_config: None,
embed_base_url: String::new(),
embed_api_key: None,
search_config: MemorySearchConfig::default(),
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
embedding_credentials: EndpointScopedCredentials::none(),
}
}
/// from_session_params stores the session_id it was given.
///
/// Direct assertion via the `#[cfg(test)]` accessor proves the factory
/// actually stored the value rather than discarding it. The counter
/// increment check additionally confirms the backend is functional.
#[tokio::test]
async fn test_factory_sets_session_id() {
let tmp = TempDir::new().unwrap();
init_sqlite_vec();
let storage = make_storage(&tmp);
let db_path = storage.workspace_dir().join("index.sqlite");
let mut idx = MemoryIndex::open_or_create(
&db_path,
storage.clone(),
kigi_config_types::MemoryIndexConfig::default(),
4,
)
.unwrap();
let file = tmp.path().join("note.md");
std::fs::write(&file, "# Facts\n\nRust is fast.").unwrap();
idx.reindex_file(&file, "workspace").unwrap();
drop(idx);
let params = make_params_fts_only("test-session-abc");
let backend = MemoryBackendImpl::from_session_params(storage, &params);
// Direct assertion: the stored session_id matches what the factory was given.
assert_eq!(
backend.session_id_for_test(),
"test-session-abc",
"session_id must be stored exactly as supplied"
);
// Functional check: the backend actually runs a search.
let before = backend
.search_counter
.load(std::sync::atomic::Ordering::Relaxed);
let _ = backend.search("rust", 5, 0.0).await;
let after = backend
.search_counter
.load(std::sync::atomic::Ordering::Relaxed);
assert_eq!(
after,
before + 1,
"search counter must increment per search"
);
}
/// from_session_params stores the search_config it was given.
///
/// Direct assertion via the `#[cfg(test)]` accessor proves the factory
/// propagated the config into the backend rather than discarding it.
/// `max_results` is verified because the `search()` method overrides it
/// with the caller's argument — so checking the *stored* value is the only
/// way to confirm the factory wired it correctly.
#[tokio::test]
async fn test_factory_wires_search_config() {
let tmp = TempDir::new().unwrap();
init_sqlite_vec();
let storage = make_storage(&tmp);
let db_path = storage.workspace_dir().join("index.sqlite");
let mut idx = MemoryIndex::open_or_create(
&db_path,
storage.clone(),
kigi_config_types::MemoryIndexConfig::default(),
4,
)
.unwrap();
for i in 0..10 {
let f = tmp.path().join(format!("note{i}.md"));
std::fs::write(&f, format!("# Entry {i}\n\nRust tip number {i}.")).unwrap();
idx.reindex_file(&f, "workspace").unwrap();
}
drop(idx);
let params = MemoryBackendParams {
search_config: MemorySearchConfig {
max_results: 3,
..Default::default()
},
..make_params_fts_only("test-search-config")
};
let backend = MemoryBackendImpl::from_session_params(storage, &params);
// Direct: the stored config has exactly the value the factory was given.
assert_eq!(
backend.search_config_for_test().max_results,
3,
"stored max_results must equal what was supplied to the factory"
);
}
/// from_session_params wires non-overridable config fields (MMR, temporal decay)
/// that `search()` never replaces with caller arguments.
///
/// This is the clearest proof that `[memory.search]` config is actually wired
/// rather than silently ignored: fields the caller cannot override must arrive
/// in the stored search_config exactly as given.
#[test]
fn test_factory_wires_non_overridable_search_config_fields() {
let tmp = TempDir::new().unwrap();
let storage = make_storage(&tmp);
let custom_search = MemorySearchConfig {
max_results: 7,
mmr: kigi_config_types::MmrConfig {
enabled: true,
lambda: 0.42,
},
temporal_decay: kigi_config_types::TemporalDecayConfig {
enabled: true,
half_life_days: 14.0,
},
..Default::default()
};
let params = MemoryBackendParams {
search_config: custom_search,
..make_params_fts_only("test-full-config")
};
let backend = MemoryBackendImpl::from_session_params(storage, &params);
let stored = backend.search_config_for_test();
// None of these are overridden by the caller in search() — they must
// survive the factory path `unchanged`.
assert_eq!(stored.max_results, 7);
assert!(stored.mmr.enabled, "MMR enabled must be stored");
assert!(
(stored.mmr.lambda - 0.42).abs() < f64::EPSILON,
"MMR lambda must be stored exactly"
);
assert!(
stored.temporal_decay.enabled,
"temporal_decay enabled must be stored"
);
assert!(
(stored.temporal_decay.half_life_days - 14.0).abs() < f64::EPSILON,
"temporal_decay half_life_days must be stored exactly"
);
}
/// from_session_params propagates search_source into the backend.
///
/// Correctness test: every caller (tool, injection,
/// compaction_recovery) must be able to set a distinct source label so
/// dashboards can separate the three search paths.
#[test]
fn test_factory_propagates_search_source() {
let tmp = TempDir::new().unwrap();
let storage = make_storage(&tmp);
for source in ["tool", "injection", "compaction_recovery"] {
let params = MemoryBackendParams {
search_source: source,
..make_params_fts_only("test-source")
};
let backend = MemoryBackendImpl::from_session_params(storage.clone(), &params);
assert_eq!(
backend.search_source, source,
"search_source must be propagated for source='{source}'"
);
}
}
/// The default search_source is "tool" when constructing via new().
#[test]
fn test_default_search_source_is_tool() {
let tmp = TempDir::new().unwrap();
let db_path = tmp.path().join("test.sqlite");
let storage = make_storage(&tmp);
let backend = MemoryBackendImpl::new(db_path, storage);
assert_eq!(backend.search_source, "tool");
}
/// MemoryBackendParams with different search_source values is Clone.
#[test]
fn test_params_clone_preserves_search_source() {
let params = MemoryBackendParams {
search_source: "injection",
..make_params_fts_only("test-clone-source")
};
let cloned = params.clone();
assert_eq!(cloned.search_source, "injection");
}
/// Watcher startup telemetry reflects actual runtime state.
///
/// `watcher.is_some()` is `true` only when the watcher started successfully.
/// With a valid directory the watcher should start; without one it should return None.
/// This guards the contract that `watcher_started` in telemetry must reflect
/// runtime outcome, not configuration intent.
#[test]
fn test_params_watcher_started_reflects_runtime() {
let tmp = TempDir::new().unwrap();
// Success path: directory exists → watcher starts.
let watch_dir = tmp.path().join("memory");
std::fs::create_dir_all(&watch_dir).unwrap();
let watcher = crate::watcher::MemoryFileWatcher::start(&watch_dir);
let params_with_watcher = MemoryBackendParams {
watcher: watcher.map(std::sync::Arc::new),
..make_params_fts_only("test-watcher-runtime")
};
// watcher.is_some() reflects whether startup succeeded.
// (On environments without inotify/FSEvents this may be None; skip rather than fail.)
// just verify it compiles
let _ = params_with_watcher.watcher.is_some();
// Failure path: non-existent directory → watcher must return None.
let missing = tmp.path().join("does_not_exist");
let no_watcher = crate::watcher::MemoryFileWatcher::start(&missing);
assert!(
no_watcher.is_none(),
"watcher must return None for a non-existent directory"
);
let params_no_watcher = MemoryBackendParams {
watcher: None,
..make_params_fts_only("test-no-watcher")
};
assert!(
params_no_watcher.watcher.is_none(),
"params.watcher.is_none() means telemetry reports watcher_started=false"
);
}
/// default_search_max_results returns the configured value from search_config.
///
/// Verifies that the MemoryBackend trait override in MemoryBackendImpl
/// exposes search_config.max_results rather than the hardcoded default (6).
#[test]
fn test_default_search_max_results_from_config() {
let tmp = TempDir::new().unwrap();
let storage = make_storage(&tmp);
let params = MemoryBackendParams {
search_config: MemorySearchConfig {
max_results: 12,
..Default::default()
},
..make_params_fts_only("test-defaults")
};
let backend = MemoryBackendImpl::from_session_params(storage, &params);
assert_eq!(
backend.default_search_max_results(),
12,
"default_search_max_results must return search_config.max_results"
);
}
/// default_search_min_score returns the configured value from search_config.
#[test]
fn test_default_search_min_score_from_config() {
let tmp = TempDir::new().unwrap();
let storage = make_storage(&tmp);
let params = MemoryBackendParams {
search_config: MemorySearchConfig {
min_score: 0.42,
..Default::default()
},
..make_params_fts_only("test-defaults")
};
let backend = MemoryBackendImpl::from_session_params(storage, &params);
assert!(
(backend.default_search_min_score() - 0.42_f64).abs() < 1e-6,
"default_search_min_score must return search_config.min_score"
);
}
/// from_session_params without embed_config produces a backend that does not panic
/// and returns results using FTS-only path.
#[tokio::test]
async fn test_factory_fts_only_without_embed() {
let tmp = TempDir::new().unwrap();
init_sqlite_vec();
let storage = make_storage(&tmp);
let db_path = storage.workspace_dir().join("index.sqlite");
let mut idx = MemoryIndex::open_or_create(
&db_path,
storage.clone(),
kigi_config_types::MemoryIndexConfig::default(),
4,
)
.unwrap();
let f = tmp.path().join("note.md");
std::fs::write(&f, "# Guide\n\nRust ownership rules.").unwrap();
idx.reindex_file(&f, "workspace").unwrap();
drop(idx);
let params = make_params_fts_only("test-fts-only");
let backend = MemoryBackendImpl::from_session_params(storage, &params);
let results = backend.search("rust ownership", 5, 0.0).await.unwrap();
assert!(
!results.is_empty(),
"FTS-only backend should return results"
);
let ts = results[0].created_at;
assert!(
ts.is_some() && ts.unwrap() > 0,
"created_at must be Some(positive) after backend search (got {ts:?})"
);
}
/// from_session_params with embed_config but no api_key gracefully falls back
/// to FTS-only (the embedding provider requires a key).
#[tokio::test]
async fn test_factory_embed_config_without_key_falls_back_to_fts() {
let tmp = TempDir::new().unwrap();
init_sqlite_vec();
let storage = make_storage(&tmp);
let db_path = storage.workspace_dir().join("index.sqlite");
let mut idx = MemoryIndex::open_or_create(
&db_path,
storage.clone(),
kigi_config_types::MemoryIndexConfig::default(),
4,
)
.unwrap();
let f = tmp.path().join("note.md");
std::fs::write(&f, "# Guide\n\nRust borrow checker.").unwrap();
idx.reindex_file(&f, "workspace").unwrap();
drop(idx);
let params = MemoryBackendParams {
embed_config: Some(MemoryEmbeddingConfig::default()),
embed_base_url: "http://localhost".to_string(),
// no key → provider cannot be created
embed_api_key: None,
..make_params_fts_only("test-embed-no-key")
};
let backend = MemoryBackendImpl::from_session_params(storage, &params);
// Must not panic; FTS results should still come back.
let results = backend.search("rust borrow", 5, 0.0).await.unwrap();
assert!(
!results.is_empty(),
"should fall back to FTS when api_key is None"
);
}
/// MemoryBackendParams is Clone.
#[test]
fn test_params_is_clone() {
let params = make_params_fts_only("clone-test");
let _cloned = params.clone();
}
/// from_session_params without watcher produces a backend that searches correctly.
#[tokio::test]
async fn test_factory_no_watcher() {
let tmp = TempDir::new().unwrap();
init_sqlite_vec();
let storage = make_storage(&tmp);
let db_path = storage.workspace_dir().join("index.sqlite");
let mut idx = MemoryIndex::open_or_create(
&db_path,
storage.clone(),
kigi_config_types::MemoryIndexConfig::default(),
4,
)
.unwrap();
let f = tmp.path().join("note.md");
std::fs::write(&f, "# Tip\n\nAlways write tests.").unwrap();
idx.reindex_file(&f, "workspace").unwrap();
drop(idx);
let params = MemoryBackendParams {
watcher: None,
..make_params_fts_only("test-no-watcher")
};
let backend = MemoryBackendImpl::from_session_params(storage, &params);
let results = backend.search("tests", 5, 0.0).await.unwrap();
assert!(
!results.is_empty(),
"no-watcher backend should still return results"
);
}
/// `ensure_initialized` must be called before watcher startup.
///
/// Regression test for the ordering fix: on a first-use machine the
/// memory directories do not exist yet. If the watcher tries to watch a
/// non-existent directory it returns `None` (silently dropping the feature).
/// After `ensure_initialized()` the directories exist and the watcher can
/// start successfully.
///
/// This mirrors the ordering enforced in `spawn_session_actor`:
/// 1. `storage.ensure_initialized()`
/// 2. `MemoryFileWatcher::start(storage.global_dir())`
#[test]
fn test_ensure_initialized_before_watcher_ordering() {
let tmp = TempDir::new().unwrap();
let global = tmp.path().join("memory");
let workspace = global.join("test_ws");
let storage = MemoryStorage::with_paths(global.clone(), workspace.clone());
// Precondition: neither directory exists yet (fresh machine simulation).
assert!(
!global.exists(),
"global memory dir must not exist before initialization"
);
// Wrong ordering (watcher before init)
// The watcher returns None because the directory does not exist.
let watcher_before_init = crate::watcher::MemoryFileWatcher::start(&global);
assert!(
watcher_before_init.is_none(),
"watcher must fail (None) when directory does not exist yet"
);
// Correct ordering (init, then watcher)
// After ensure_initialized the directories and MEMORY.md templates exist.
storage.ensure_initialized().unwrap();
assert!(
global.exists(),
"global dir must exist after ensure_initialized"
);
assert!(
workspace.exists(),
"workspace dir must exist after ensure_initialized"
);
assert!(
global.join("MEMORY.md").exists(),
"global MEMORY.md template must exist"
);
assert!(
workspace.join("MEMORY.md").exists(),
"workspace MEMORY.md template must exist"
);
// Watcher now succeeds because the directory exists.
// (Allowed to return None in environments without inotify/kqueue
// support — e.g. some CI containers — but must not error-panic.)
let watcher_after_init = crate::watcher::MemoryFileWatcher::start(&global);
// If a watcher was returned we can confirm it is usable (not dirty yet).
if let Some(w) = watcher_after_init {
assert!(
!w.is_dirty(),
"freshly started watcher must report no dirty files"
);
}
// If None, the test environment does not support file-watching —
// that is acceptable; the directories themselves are what matter here.
}
/// End-to-end regression test for the watcher-driven delete path.
///
/// Tests the full chain:
/// 1. file is indexed
/// 2. watcher is started
/// 3. first `backend.search()` confirms content is found
/// 4. file is deleted (OS fires a Remove event to the watcher)
/// 5. second `backend.search()` triggers sync-on-search, which calls
/// `delete_path()` because the file no longer exists
/// 6. content is no longer returned
///
/// This test guards against regressions in the `file.exists() → else
/// delete_path()` branch that would be invisible to the `delete_path`
/// unit tests alone.
#[tokio::test]
async fn test_watcher_delete_clears_stale_chunks() {
let tmp = TempDir::new().unwrap();
init_sqlite_vec();
let global = tmp.path().join("memory");
let workspace = global.join("test_ws");
std::fs::create_dir_all(&global).unwrap();
std::fs::create_dir_all(&workspace).unwrap();
let storage = MemoryStorage::with_paths(global.clone(), workspace);
let db_path = storage.workspace_dir().join("index.sqlite");
// Step 1: Write + canonicalize the file path BEFORE indexing.
//
// On macOS, TempDir paths may live under /private/tmp (via a symlink
// from /tmp). FSEvents returns canonicalized paths, so the path stored
// in the index must match what the watcher event delivers.
let file_raw = global.join("note.md");
std::fs::write(&file_raw, "# Unique\n\nXyzzy-watcher-delete-token.").unwrap();
let file = dunce::canonicalize(&file_raw).unwrap_or(file_raw);
{
let mut idx = MemoryIndex::open_or_create(
&db_path,
storage.clone(),
kigi_config_types::MemoryIndexConfig::default(),
4,
)
.unwrap();
// Index with the canonical path so DB key matches watcher event paths.
idx.reindex_file(&file, "workspace").unwrap();
}
// Step 2: Start watcher AFTER indexing so the Remove event for the
// upcoming deletion is the first event the watcher ever sees.
let watch_dir = dunce::canonicalize(&global).unwrap_or(global.clone());
let watcher = match crate::watcher::MemoryFileWatcher::start(&watch_dir) {
Some(w) => w,
None => {
// File-watching not supported in this environment (e.g., some CI
// containers without inotify/FSEvents). Skip rather than fail.
return;
}
};
let watcher_arc = std::sync::Arc::new(watcher);
let params = MemoryBackendParams {
watcher: Some(watcher_arc.clone()),
..make_params_fts_only("test-watcher-delete")
};
let backend = MemoryBackendImpl::from_session_params(storage, &params);
// Step 3: Confirm content is found before deletion.
let before = backend
.search("Xyzzy-watcher-delete-token", 5, 0.0)
.await
.unwrap();
assert!(
!before.is_empty(),
"content must be found before file is deleted"
);
// Step 4: Delete the file — the OS will fire a Remove event.
std::fs::remove_file(&file).unwrap();
// Poll until the watcher detects the event (more reliable than a fixed
// sleep on macOS where FSEvents delivery time varies considerably).
// Give up after 2 s and skip the timing-sensitive assertion rather than
// flake — delete_path unit tests cover the underlying logic.
let mut event_delivered = false;
for _ in 0..20 {
std::thread::sleep(std::time::Duration::from_millis(100));
if watcher_arc.is_dirty() {
event_delivered = true;
break;
}
}
if !event_delivered {
// FSEvents not delivered within 2 s — environment is too slow.
// Skip silently; the logic is covered by delete_path unit tests.
return;
}
// Step 5+6: search triggers sync-on-search, which detects file.exists()
// == false and calls delete_path(), clearing all stale chunks.
let after = backend
.search("Xyzzy-watcher-delete-token", 5, 0.0)
.await
.unwrap();
assert!(
after.is_empty(),
"deleted file's content must not appear after watcher-driven delete sync"
);
}
struct ProbeCredentials;
impl kigi_auth::HttpAuth for ProbeCredentials {
fn apply(
&self,
builder: reqwest::RequestBuilder,
_base_url: &str,
) -> reqwest::RequestBuilder {
builder.bearer_auth("session-bearer")
}
}
#[async_trait::async_trait]
impl kigi_auth::AuthCredentialProvider for ProbeCredentials {
fn snapshot(&self) -> kigi_auth::CredentialSnapshot {
kigi_auth::CredentialSnapshot {
token: Some("session-bearer".into()),
..Default::default()
}
}
async fn refresh_after_unauthorized(&self) -> bool {
false
}
}
const SESSION_ENDPOINT: &str = "https://api.kimi.com/coding/v1";
const FOREIGN_ENDPOINT: &str = "https://api.anthropic.com/v1";
fn params_at(base_url: &str, credentials: EndpointScopedCredentials) -> MemoryBackendParams {
MemoryBackendParams {
session_id: "s1".into(),
embed_config: Some(MemoryEmbeddingConfig {
model: Some("test-embed-model".into()),
..Default::default()
}),
embed_base_url: base_url.into(),
embed_api_key: None,
search_config: MemorySearchConfig::default(),
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
embedding_credentials: credentials,
}
}
/// The session bearer must never ride to a third-party embedding host.
///
/// `embed_base_url` is the CURRENT MODEL's endpoint, so a session on a BYOK
/// or subscription-OAuth model aims memory embeddings at that provider —
/// and the auth middleware stamps `Authorization` unconditionally.
#[tokio::test]
async fn session_credentials_are_withheld_from_a_foreign_endpoint() {
let scoped = EndpointScopedCredentials::for_endpoint(
FOREIGN_ENDPOINT,
false,
Some(Arc::new(ProbeCredentials)),
);
assert!(
scoped.is_empty(),
"an untrusted endpoint must drop both handles"
);
let provider = params_at(FOREIGN_ENDPOINT, scoped)
.make_embedding_provider()
.await;
assert!(
provider.is_none(),
"no credential may ride to a foreign endpoint, and there is no static key to fall back to"
);
}
/// The trusted-endpoint path still builds a credentialed provider.
#[tokio::test]
async fn session_credentials_ride_their_own_endpoint() {
let scoped = EndpointScopedCredentials::for_endpoint(
SESSION_ENDPOINT,
true,
Some(Arc::new(ProbeCredentials)),
);
assert!(
!scoped.is_empty(),
"a trusted endpoint must retain the handles"
);
let provider = params_at(SESSION_ENDPOINT, scoped)
.make_embedding_provider()
.await;
assert!(
provider.is_some(),
"the session endpoint keeps its refresh-capable credential"
);
}
/// The runtime re-check, not construction alone, is what guards the wire:
/// `MemoryBackendParams` is `Clone` and callers rewrite fields on the copy.
#[tokio::test]
async fn a_cloned_param_set_cannot_redirect_scoped_credentials() {
let scoped = EndpointScopedCredentials::for_endpoint(
SESSION_ENDPOINT,
true,
Some(Arc::new(ProbeCredentials)),
);
let redirected = MemoryBackendParams {
embed_base_url: FOREIGN_ENDPOINT.into(),
..params_at(SESSION_ENDPOINT, scoped)
};
assert!(
redirected.make_embedding_provider().await.is_none(),
"credentials scoped to one endpoint must not follow a rewritten base_url"
);
}
/// A platform's own API key is not a session credential: it is resolved for
/// that platform and must keep serving that platform's endpoint.
#[tokio::test]
async fn a_platform_api_key_still_serves_its_own_endpoint() {
let params = MemoryBackendParams {
embed_api_key: Some("platform-key".into()),
..params_at(FOREIGN_ENDPOINT, EndpointScopedCredentials::none())
};
assert!(
params.make_embedding_provider().await.is_some(),
"withholding the session credential must not disable BYOK embeddings"
);
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::index::{MemoryIndex, init_sqlite_vec};
use kigi_config_types::MemoryIndexConfig;
use tempfile::TempDir;
fn setup_index(tmp: &TempDir) -> (PathBuf, MemoryStorage) {
init_sqlite_vec();
let global = tmp.path().join("memory");
let workspace = global.join("test_ws");
let storage = MemoryStorage::with_paths(global, workspace);
let db_path = tmp.path().join("test.sqlite");
let mut idx =
MemoryIndex::open_or_create(&db_path, storage.clone(), MemoryIndexConfig::default(), 4)
.unwrap();
let file_path = tmp.path().join("test.md");
std::fs::write(&file_path, "# Guide\n\nRust programming tutorial.").unwrap();
idx.reindex_file(&file_path, "workspace").unwrap();
(db_path, storage)
}
#[tokio::test]
async fn test_backend_search() {
let tmp = TempDir::new().unwrap();
let (db_path, storage) = setup_index(&tmp);
let backend = MemoryBackendImpl::new(db_path, storage);
let results = backend.search("rust programming", 10, 0.0).await.unwrap();
assert!(!results.is_empty(), "should find indexed content");
assert!(results[0].snippet.contains("Rust"));
}
#[test]
fn test_backend_total_chunks() {
let tmp = TempDir::new().unwrap();
let (db_path, storage) = setup_index(&tmp);
let backend = MemoryBackendImpl::new(db_path, storage);
let count = backend.total_chunks().unwrap();
assert!(count >= 1, "should have at least 1 chunk");
}
#[test]
fn test_backend_is_send_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<MemoryBackendImpl>();
}
#[tokio::test]
async fn test_search_with_punctuation_in_query() {
let tmp = TempDir::new().unwrap();
let (db_path, storage) = setup_index(&tmp);
let backend = MemoryBackendImpl::new(db_path, storage);
// Raw user message with punctuation — should not crash FTS5
let results = backend
.search("what is rust? how to use it!", 10, 0.0)
.await
.unwrap();
assert!(
!results.is_empty(),
"should match 'rust' despite punctuation in query"
);
}
#[tokio::test]
async fn test_search_with_special_chars_only() {
let tmp = TempDir::new().unwrap();
let (db_path, storage) = setup_index(&tmp);
let backend = MemoryBackendImpl::new(db_path, storage);
// Query with only special chars — should return empty, not error
let results = backend.search("???!!!", 10, 0.0).await.unwrap();
assert!(
results.is_empty(),
"special-chars-only query should return empty"
);
}
#[tokio::test]
async fn test_search_hybrid_fts_only_fallback() {
// Without embedding config, hybrid search should degrade to FTS-only
let tmp = TempDir::new().unwrap();
let (db_path, storage) = setup_index(&tmp);
let backend = MemoryBackendImpl::new(db_path, storage);
// Even with high min_score, hybrid search normalizes scores to [0,1]
// so results above the threshold should be returned
let results = backend.search("rust programming", 10, 0.0).await.unwrap();
assert!(
!results.is_empty(),
"FTS-only fallback should still return results"
);
// Scores should be normalized (0,1] range from hybrid scoring
assert!(results[0].score > 0.0, "hybrid scores should be positive");
}
/// The supplemental evergreen query in `search()` adds global/workspace
/// candidates that the base `search_fts` missed due to candidate_limit.
///
/// Tests the mechanism directly at the index level: verifies that with
/// a tight FTS limit, global/workspace chunks are absent from the base
/// results but present in the supplemental source-filtered query. Then
/// confirms the full backend search pipeline surfaces them.
#[tokio::test]
async fn test_search_returns_global_and_workspace_memory() {
let tmp = TempDir::new().unwrap();
init_sqlite_vec();
let global = tmp.path().join("memory");
let workspace = global.join("test_ws");
let storage = MemoryStorage::with_paths(global, workspace);
let db_path = storage.workspace_dir().join("index.sqlite");
let mut idx =
MemoryIndex::open_or_create(&db_path, storage.clone(), MemoryIndexConfig::default(), 4)
.unwrap();
// Index global + workspace with matching content.
let global_file = tmp.path().join("global_mem.md");
std::fs::write(
&global_file,
"# Preferences\n\nAlways use graphite for PRs. Prefer Rust over Python.",
)
.unwrap();
idx.reindex_file(&global_file, "global").unwrap();
let ws_file = tmp.path().join("ws_mem.md");
std::fs::write(
&ws_file,
"# Project Decisions\n\nWe chose graphite for PRs in this project.",
)
.unwrap();
idx.reindex_file(&ws_file, "workspace").unwrap();
// Index session files that also match the query.
for i in 0..5 {
let f = tmp.path().join(format!("session_{i}.md"));
std::fs::write(
&f,
format!("# Session {i}\n\nDiscussed graphite for PRs and item {i}."),
)
.unwrap();
idx.reindex_file(&f, "session").unwrap();
}
// Verify the supplemental query mechanism: with a tight limit the
// base FTS returns a mix, but `search_fts_by_sources` for
// "global"/"workspace" always finds the evergreen chunks.
let evergreen = idx
.search_fts_by_sources("graphite PRs", 10, &["global", "workspace"])
.unwrap();
assert!(
evergreen.len() >= 2,
"supplemental evergreen query must find both global and workspace chunks"
);
let evergreen_sources: Vec<String> = evergreen
.iter()
.filter_map(|r| idx.get_chunk(&r.chunk_id).ok().flatten())
.map(|c| c.source)
.collect();
assert!(
evergreen_sources.contains(&"global".to_string()),
"evergreen query must find global chunk"
);
assert!(
evergreen_sources.contains(&"workspace".to_string()),
"evergreen query must find workspace chunk"
);
drop(idx);
// Full backend search: global/workspace must appear in results.
let backend = MemoryBackendImpl::new(db_path, storage);
let results = backend.search("graphite PRs", 10, 0.0).await.unwrap();
let has_global = results.iter().any(|r| r.source == "global");
let has_workspace = results.iter().any(|r| r.source == "workspace");
assert!(
has_global,
"global MEMORY.md chunks must appear in search results"
);
assert!(
has_workspace,
"workspace MEMORY.md chunks must appear in search results"
);
}
}
#[cfg(test)]
mod index_embedding_tests {
use crate::index::MemoryIndex;
use crate::storage::MemoryStorage;
#[test]
fn test_chunks_without_embeddings() {
let tmp = tempfile::TempDir::new().unwrap();
let global = tmp.path().join("memory");
let workspace = global.join("test_ws");
let storage = MemoryStorage::with_paths(global, workspace);
let db_path = tmp.path().join("test.sqlite");
let mut idx = MemoryIndex::open_or_create(
&db_path,
storage,
kigi_config_types::MemoryIndexConfig::default(),
4,
)
.unwrap();
if !idx.vec_available() {
// sqlite-vec not available — chunks_without_embeddings returns empty
let missing = idx.chunks_without_embeddings().unwrap();
assert!(missing.is_empty(), "no-vec: should return empty");
return;
}
let file_path = tmp.path().join("test.md");
std::fs::write(&file_path, "# Title\n\nSome content here.").unwrap();
idx.reindex_file(&file_path, "workspace").unwrap();
// After reindex, chunks should exist but have no embeddings
let missing = idx.chunks_without_embeddings().unwrap();
assert!(
!missing.is_empty(),
"newly indexed chunks should be missing embeddings"
);
// After upserting an embedding, the chunk should disappear from missing
let (chunk_id, _) = &missing[0];
let dummy_embedding = vec![0.0f32; 4];
idx.upsert_embedding(chunk_id, &dummy_embedding).unwrap();
let missing_after = idx.chunks_without_embeddings().unwrap();
assert_eq!(
missing_after.len(),
missing.len() - 1,
"one fewer chunk should be missing after embedding"
);
}
}