ZacharyZhang-NY
ed8049cf77
`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.
103 lines
3.2 KiB
Rust
103 lines
3.2 KiB
Rust
//! Markdown-based memory storage that persists knowledge across sessions.
|
|
//!
|
|
//! ## Data Layout
|
|
//!
|
|
//! ```text
|
|
//! ~/.kigi/memory/
|
|
//! ├── MEMORY.md # Global curated knowledge
|
|
//! └── {workspace_hash}/ # Per-workspace (blake3(cwd)[..16])
|
|
//! ├── MEMORY.md # Project-level curated knowledge
|
|
//! └── sessions/
|
|
//! └── YYYY-MM-DD-{slug}-{sid8}.md # Session logs
|
|
//! ```
|
|
//!
|
|
//! ## Feature Flag
|
|
//!
|
|
//! Memory is gated behind the `--experimental-memory` CLI flag or
|
|
//! `KIGI_MEMORY=1`; when disabled the host never initializes this crate.
|
|
|
|
pub mod archive;
|
|
pub mod backend;
|
|
pub mod chunker;
|
|
pub mod dream;
|
|
pub mod dream_lock;
|
|
pub mod embedding;
|
|
pub mod index;
|
|
pub mod mmr;
|
|
pub mod query_expansion;
|
|
pub mod schema;
|
|
pub mod search;
|
|
pub mod storage;
|
|
pub mod text_utils;
|
|
pub mod watcher;
|
|
|
|
pub use backend::{EndpointScopedCredentials, MemoryBackendImpl, MemoryBackendParams};
|
|
pub use index::{MemoryIndex, init_sqlite_vec};
|
|
pub use storage::{MemoryScope, MemoryStorage};
|
|
|
|
/// Embeds every chunk that has no embedding yet, returning how many succeeded.
|
|
///
|
|
/// The async glue between the sync `MemoryIndex` and the async
|
|
/// `EmbeddingProvider`. Call after reindex, flush writes, or session-end
|
|
/// writes. Failures are logged and skipped rather than propagated, so a dead
|
|
/// embedding endpoint degrades search instead of breaking the session.
|
|
pub async fn embed_missing_chunks(
|
|
index: &MemoryIndex,
|
|
provider: &dyn embedding::EmbeddingProvider,
|
|
) -> usize {
|
|
let chunks = match index.chunks_without_embeddings() {
|
|
Ok(c) if c.is_empty() => return 0,
|
|
Ok(c) => c,
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
target: kigi_log::memory_log::TARGET,
|
|
error = %e,
|
|
"failed to query chunks without embeddings"
|
|
);
|
|
return 0;
|
|
}
|
|
};
|
|
|
|
let total = chunks.len();
|
|
let mut embedded = 0;
|
|
|
|
// 32 matches the typical provider max batch size.
|
|
for batch in chunks.chunks(32) {
|
|
let texts: Vec<&str> = batch.iter().map(|(_, text)| text.as_str()).collect();
|
|
match provider.embed_batch(&texts).await {
|
|
Ok(embeddings) => {
|
|
for ((chunk_id, _), embedding) in batch.iter().zip(embeddings.iter()) {
|
|
if let Err(e) = index.upsert_embedding(chunk_id, embedding) {
|
|
tracing::warn!(
|
|
target: kigi_log::memory_log::TARGET,
|
|
chunk_id,
|
|
error = %e,
|
|
"failed to upsert embedding"
|
|
);
|
|
} else {
|
|
embedded += 1;
|
|
}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
target: kigi_log::memory_log::TARGET,
|
|
error = %e,
|
|
batch_size = texts.len(),
|
|
"embedding batch failed, skipping"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
if embedded > 0 {
|
|
tracing::info!(
|
|
target: kigi_log::memory_log::TARGET,
|
|
embedded,
|
|
total,
|
|
"embedded missing chunks"
|
|
);
|
|
}
|
|
embedded
|
|
}
|