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.
This commit is contained in:
@@ -18,6 +18,76 @@ use super::embedding::EmbeddingProvider as _;
|
|||||||
use super::storage::MemoryStorage;
|
use super::storage::MemoryStorage;
|
||||||
use super::watcher::MemoryFileWatcher;
|
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.
|
/// 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
|
/// Grouping these in one struct ensures every call site — ToolBridge, first-turn
|
||||||
@@ -30,7 +100,8 @@ pub struct MemoryBackendParams {
|
|||||||
pub session_id: String,
|
pub session_id: String,
|
||||||
/// Embedding provider config — `None` forces FTS-only fallback everywhere.
|
/// Embedding provider config — `None` forces FTS-only fallback everywhere.
|
||||||
pub embed_config: Option<kigi_config_types::MemoryEmbeddingConfig>,
|
pub embed_config: Option<kigi_config_types::MemoryEmbeddingConfig>,
|
||||||
/// Base URL for embedding API calls (CLI proxy).
|
/// 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,
|
pub embed_base_url: String,
|
||||||
/// API key for embedding API calls.
|
/// API key for embedding API calls.
|
||||||
pub embed_api_key: Option<String>,
|
pub embed_api_key: Option<String>,
|
||||||
@@ -47,10 +118,8 @@ pub struct MemoryBackendParams {
|
|||||||
/// - `"injection"` — first-turn memory context injection
|
/// - `"injection"` — first-turn memory context injection
|
||||||
/// - `"compaction_recovery"` — post-compaction context re-injection
|
/// - `"compaction_recovery"` — post-compaction context re-injection
|
||||||
pub search_source: &'static str,
|
pub search_source: &'static str,
|
||||||
/// Dynamic API key provider — when set, `make_embedding_provider()` resolves
|
/// The session credentials, and the single endpoint they may reach.
|
||||||
/// the key per-call instead of using the static `embed_api_key`.
|
pub embedding_credentials: EndpointScopedCredentials,
|
||||||
pub api_key_provider: Option<kigi_tools::types::SharedApiKeyProvider>,
|
|
||||||
pub auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemoryBackendParams {
|
impl MemoryBackendParams {
|
||||||
@@ -59,8 +128,7 @@ impl MemoryBackendParams {
|
|||||||
pub async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
|
pub async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
|
||||||
build_embedding_provider(
|
build_embedding_provider(
|
||||||
self.embed_config.as_ref(),
|
self.embed_config.as_ref(),
|
||||||
self.auth_credentials.as_ref(),
|
&self.embedding_credentials,
|
||||||
self.api_key_provider.as_ref(),
|
|
||||||
self.embed_api_key.as_deref(),
|
self.embed_api_key.as_deref(),
|
||||||
&self.embed_base_url,
|
&self.embed_base_url,
|
||||||
)
|
)
|
||||||
@@ -70,8 +138,7 @@ impl MemoryBackendParams {
|
|||||||
|
|
||||||
async fn build_embedding_provider(
|
async fn build_embedding_provider(
|
||||||
config: Option<&kigi_config_types::MemoryEmbeddingConfig>,
|
config: Option<&kigi_config_types::MemoryEmbeddingConfig>,
|
||||||
auth_credentials: Option<&Arc<dyn kigi_auth::AuthCredentialProvider>>,
|
credentials: &EndpointScopedCredentials,
|
||||||
api_key_provider: Option<&kigi_tools::types::SharedApiKeyProvider>,
|
|
||||||
static_api_key: Option<&str>,
|
static_api_key: Option<&str>,
|
||||||
base_url: &str,
|
base_url: &str,
|
||||||
) -> Option<super::embedding::ApiEmbeddingProvider> {
|
) -> Option<super::embedding::ApiEmbeddingProvider> {
|
||||||
@@ -80,9 +147,19 @@ async fn build_embedding_provider(
|
|||||||
return None;
|
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
|
// Prefer the refresh-capable credential provider — the middleware gives
|
||||||
// 401 retry for free without any per-call key resolution.
|
// 401 retry for free without any per-call key resolution.
|
||||||
if let Some(creds) = auth_credentials {
|
if approved && let Some(creds) = credentials.auth_credentials.as_ref() {
|
||||||
let client = super::embedding::build_middleware_client(creds.clone());
|
let client = super::embedding::build_middleware_client(creds.clone());
|
||||||
return super::embedding::ApiEmbeddingProvider::from_config(
|
return super::embedding::ApiEmbeddingProvider::from_config(
|
||||||
config,
|
config,
|
||||||
@@ -91,14 +168,13 @@ async fn build_embedding_provider(
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Fallback: resolve API key per-call, wrap in a static middleware client
|
// The platform's own configured key, wrapped in a static middleware client
|
||||||
// (no 401 refresh, but auth header is still stamped by middleware).
|
// (no 401 refresh, but the auth header is still stamped by middleware).
|
||||||
let api_key = match api_key_provider {
|
super::embedding::ApiEmbeddingProvider::from_session(
|
||||||
Some(p) => p.current_api_key_async().await,
|
config,
|
||||||
None => None,
|
base_url.to_owned(),
|
||||||
}
|
static_api_key?.to_owned(),
|
||||||
.or_else(|| static_api_key.map(|s| s.to_owned()))?;
|
)
|
||||||
super::embedding::ApiEmbeddingProvider::from_session(config, base_url.to_owned(), api_key)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// `MemoryBackend` implementation backed by hybrid search (FTS5 + vector KNN).
|
/// `MemoryBackend` implementation backed by hybrid search (FTS5 + vector KNN).
|
||||||
@@ -129,10 +205,8 @@ pub struct MemoryBackendImpl {
|
|||||||
/// Only the ToolBridge backend's counter is shared back to the session actor;
|
/// Only the ToolBridge backend's counter is shared back to the session actor;
|
||||||
/// injection and compaction-recovery backends use their own local counters.
|
/// injection and compaction-recovery backends use their own local counters.
|
||||||
pub search_counter: std::sync::Arc<std::sync::atomic::AtomicU64>,
|
pub search_counter: std::sync::Arc<std::sync::atomic::AtomicU64>,
|
||||||
/// Dynamic API key provider for embedding requests.
|
/// The session credentials, and the single endpoint they may reach.
|
||||||
api_key_provider: Option<kigi_tools::types::SharedApiKeyProvider>,
|
embedding_credentials: EndpointScopedCredentials,
|
||||||
/// Refresh-capable credential provider for embedding HTTP middleware.
|
|
||||||
auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemoryBackendImpl {
|
impl MemoryBackendImpl {
|
||||||
@@ -150,8 +224,7 @@ impl MemoryBackendImpl {
|
|||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
session_id: String::new(),
|
session_id: String::new(),
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
search_counter: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)),
|
search_counter: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -200,8 +273,7 @@ impl MemoryBackendImpl {
|
|||||||
async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
|
async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
|
||||||
build_embedding_provider(
|
build_embedding_provider(
|
||||||
self.embed_config.as_ref(),
|
self.embed_config.as_ref(),
|
||||||
self.auth_credentials.as_ref(),
|
&self.embedding_credentials,
|
||||||
self.api_key_provider.as_ref(),
|
|
||||||
self.embed_api_key.as_deref(),
|
self.embed_api_key.as_deref(),
|
||||||
&self.embed_base_url,
|
&self.embed_base_url,
|
||||||
)
|
)
|
||||||
@@ -232,8 +304,7 @@ impl MemoryBackendImpl {
|
|||||||
if let Some(w) = ¶ms.watcher {
|
if let Some(w) = ¶ms.watcher {
|
||||||
backend = backend.with_watcher(w.clone(), params.stale_claim_secs);
|
backend = backend.with_watcher(w.clone(), params.stale_claim_secs);
|
||||||
}
|
}
|
||||||
backend.api_key_provider = params.api_key_provider.clone();
|
backend.embedding_credentials = params.embedding_credentials.clone();
|
||||||
backend.auth_credentials = params.auth_credentials.clone();
|
|
||||||
backend
|
backend
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -479,8 +550,7 @@ mod factory_tests {
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1038,72 +1108,128 @@ mod factory_tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Regression: provider build must use `current_api_key_async`,
|
struct ProbeCredentials;
|
||||||
/// never sync. Prevents memory_search 401s on rotated tokens.
|
impl kigi_auth::HttpAuth for ProbeCredentials {
|
||||||
#[tokio::test]
|
fn apply(
|
||||||
async fn make_embedding_provider_uses_async_api_key_resolution() {
|
&self,
|
||||||
use kigi_tools::types::ApiKeyProvider;
|
builder: reqwest::RequestBuilder,
|
||||||
use std::sync::atomic::{AtomicU32, Ordering};
|
_base_url: &str,
|
||||||
|
) -> reqwest::RequestBuilder {
|
||||||
struct AsyncProbe {
|
builder.bearer_auth("session-bearer")
|
||||||
sync_calls: Arc<AtomicU32>,
|
|
||||||
async_calls: Arc<AtomicU32>,
|
|
||||||
}
|
}
|
||||||
impl ApiKeyProvider for AsyncProbe {
|
}
|
||||||
fn current_api_key(&self) -> Option<String> {
|
#[async_trait::async_trait]
|
||||||
self.sync_calls.fetch_add(1, Ordering::SeqCst);
|
impl kigi_auth::AuthCredentialProvider for ProbeCredentials {
|
||||||
Some("sync-stale".into())
|
fn snapshot(&self) -> kigi_auth::CredentialSnapshot {
|
||||||
}
|
kigi_auth::CredentialSnapshot {
|
||||||
fn current_api_key_async(
|
token: Some("session-bearer".into()),
|
||||||
&self,
|
..Default::default()
|
||||||
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Option<String>> + Send + '_>>
|
|
||||||
{
|
|
||||||
let counter = self.async_calls.clone();
|
|
||||||
Box::pin(async move {
|
|
||||||
counter.fetch_add(1, Ordering::SeqCst);
|
|
||||||
Some("async-fresh".into())
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
async fn refresh_after_unauthorized(&self) -> bool {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let sync_calls = Arc::new(AtomicU32::new(0));
|
const SESSION_ENDPOINT: &str = "https://api.kimi.com/coding/v1";
|
||||||
let async_calls = Arc::new(AtomicU32::new(0));
|
const FOREIGN_ENDPOINT: &str = "https://api.anthropic.com/v1";
|
||||||
let probe: kigi_tools::types::SharedApiKeyProvider = Arc::new(AsyncProbe {
|
|
||||||
sync_calls: sync_calls.clone(),
|
|
||||||
async_calls: async_calls.clone(),
|
|
||||||
});
|
|
||||||
|
|
||||||
let params = MemoryBackendParams {
|
fn params_at(base_url: &str, credentials: EndpointScopedCredentials) -> MemoryBackendParams {
|
||||||
|
MemoryBackendParams {
|
||||||
session_id: "s1".into(),
|
session_id: "s1".into(),
|
||||||
embed_config: Some(MemoryEmbeddingConfig {
|
embed_config: Some(MemoryEmbeddingConfig {
|
||||||
model: Some("test-embed-model".into()),
|
model: Some("test-embed-model".into()),
|
||||||
..Default::default()
|
..Default::default()
|
||||||
}),
|
}),
|
||||||
embed_base_url: "http://example/v1".into(),
|
embed_base_url: base_url.into(),
|
||||||
embed_api_key: Some("static-fallback".into()),
|
embed_api_key: None,
|
||||||
search_config: MemorySearchConfig::default(),
|
search_config: MemorySearchConfig::default(),
|
||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: Some(probe),
|
embedding_credentials: credentials,
|
||||||
// No auth_credentials — forces the api_key_provider fallback path.
|
}
|
||||||
auth_credentials: None,
|
}
|
||||||
};
|
|
||||||
|
|
||||||
let provider = params.make_embedding_provider().await;
|
/// 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!(
|
assert!(
|
||||||
provider.is_some(),
|
provider.is_some(),
|
||||||
"provider must be built when model is set"
|
"the session endpoint keeps its refresh-capable credential"
|
||||||
);
|
);
|
||||||
assert_eq!(
|
}
|
||||||
async_calls.load(Ordering::SeqCst),
|
|
||||||
1,
|
/// The runtime re-check, not construction alone, is what guards the wire:
|
||||||
"must call current_api_key_async exactly once per provider build"
|
/// `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)),
|
||||||
);
|
);
|
||||||
assert_eq!(
|
let redirected = MemoryBackendParams {
|
||||||
sync_calls.load(Ordering::SeqCst),
|
embed_base_url: FOREIGN_ENDPOINT.into(),
|
||||||
0,
|
..params_at(SESSION_ENDPOINT, scoped)
|
||||||
"sync current_api_key must NOT be called — the async path is the contract"
|
};
|
||||||
|
|
||||||
|
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"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -31,7 +31,7 @@ pub mod storage;
|
|||||||
pub mod text_utils;
|
pub mod text_utils;
|
||||||
pub mod watcher;
|
pub mod watcher;
|
||||||
|
|
||||||
pub use backend::{MemoryBackendImpl, MemoryBackendParams};
|
pub use backend::{EndpointScopedCredentials, MemoryBackendImpl, MemoryBackendParams};
|
||||||
pub use index::{MemoryIndex, init_sqlite_vec};
|
pub use index::{MemoryIndex, init_sqlite_vec};
|
||||||
pub use storage::{MemoryScope, MemoryStorage};
|
pub use storage::{MemoryScope, MemoryStorage};
|
||||||
|
|
||||||
|
|||||||
@@ -83,6 +83,35 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider {
|
|||||||
self.auth_manager.try_recover_unauthorized().await
|
self.auth_manager.try_recover_unauthorized().await
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
/// The memory-embedding credentials for `embed_base_url`, decided by the ONE
|
||||||
|
/// authority rather than re-derived here (C1).
|
||||||
|
///
|
||||||
|
/// `embed_base_url` is the session model's own endpoint, so on a BYOK or
|
||||||
|
/// subscription-OAuth model it points at that provider's host — and
|
||||||
|
/// [`kigi_auth::AuthRetryMiddleware`] stamps `Authorization` on every request
|
||||||
|
/// it wraps. Asking [`CredentialAuthority::manager_for`] answers both halves at
|
||||||
|
/// once: whether a session credential may ride there at all, and WHICH manager
|
||||||
|
/// governs it (a subscription-OAuth platform's pooled manager at its own host,
|
||||||
|
/// the primary at the session's coding endpoint). No manager means no session
|
||||||
|
/// credential at all; the platform's own `embed_api_key` is untouched and keeps
|
||||||
|
/// serving its own endpoint.
|
||||||
|
///
|
||||||
|
/// The session's `SharedApiKeyProvider` is deliberately NOT forwarded: it is
|
||||||
|
/// hard-wired to the PRIMARY manager, so at a pooled platform's host — where a
|
||||||
|
/// credential may ride, but only that platform's own — it would resolve the
|
||||||
|
/// wrong bearer.
|
||||||
|
pub(crate) fn embedding_session_credentials(
|
||||||
|
embed_base_url: &str,
|
||||||
|
platform: Option<kigi_models::PlatformId>,
|
||||||
|
authority: &crate::auth::credential_authority::CredentialAuthority,
|
||||||
|
) -> kigi_memory::EndpointScopedCredentials {
|
||||||
|
let auth_credentials = authority.manager_for(platform, embed_base_url).map(|am| {
|
||||||
|
Arc::new(ShellAuthCredentialProvider::new(am, None, None))
|
||||||
|
as Arc<dyn AuthCredentialProvider>
|
||||||
|
});
|
||||||
|
let may_ride = auth_credentials.is_some();
|
||||||
|
kigi_memory::EndpointScopedCredentials::for_endpoint(embed_base_url, may_ride, auth_credentials)
|
||||||
|
}
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -253,6 +282,38 @@ mod tests {
|
|||||||
"snapshot must reflect refreshed token for subsequent apply() calls"
|
"snapshot must reflect refreshed token for subsequent apply() calls"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
/// C1: memory embeddings follow the credential authority, not the session.
|
||||||
|
///
|
||||||
|
/// A session whose model is a BYOK platform aims `embed_base_url` at that
|
||||||
|
/// provider's host; the authority answers "no manager governs a credential
|
||||||
|
/// there", so nothing rides. The session's own coding endpoint still does.
|
||||||
|
#[test]
|
||||||
|
fn embedding_credentials_follow_the_credential_authority() {
|
||||||
|
let _guard = EarlyInvalidationGuard::pin_to_default();
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let mgr = make_manager(
|
||||||
|
&dir,
|
||||||
|
Some(make_auth("session-bearer", ChronoDuration::hours(1))),
|
||||||
|
);
|
||||||
|
let endpoints = crate::agent::config::EndpointsConfig::default();
|
||||||
|
let coding_endpoint = endpoints.proxy_url();
|
||||||
|
let authority =
|
||||||
|
crate::auth::credential_authority::CredentialAuthority::new(endpoints, Some(mgr));
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
!embedding_session_credentials(&coding_endpoint, None, &authority).is_empty(),
|
||||||
|
"the session's own coding endpoint keeps its credential"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
embedding_session_credentials(
|
||||||
|
"https://api.anthropic.com/v1",
|
||||||
|
kigi_models::PlatformId::parse("anthropic"),
|
||||||
|
&authority,
|
||||||
|
)
|
||||||
|
.is_empty(),
|
||||||
|
"an API-key platform's host must receive no session credential"
|
||||||
|
);
|
||||||
|
}
|
||||||
/// Deployment-key path has no recovery (operator owns the bearer).
|
/// Deployment-key path has no recovery (operator owns the bearer).
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn refresh_after_unauthorized_is_noop_for_deployment_key() {
|
async fn refresh_after_unauthorized_is_noop_for_deployment_key() {
|
||||||
|
|||||||
@@ -311,6 +311,18 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
};
|
};
|
||||||
let embed_base_url = sampling_config.base_url.clone();
|
let embed_base_url = sampling_config.base_url.clone();
|
||||||
let embed_api_key = sampling_config.api_key.clone();
|
let embed_api_key = sampling_config.api_key.clone();
|
||||||
|
// The platform behind the endpoint memory embeddings will call, resolved
|
||||||
|
// through the SAME session-key disambiguation the actor is seeded with, so
|
||||||
|
// a slug that collides across platforms cannot resolve to the twin (H-b).
|
||||||
|
let embed_platform = {
|
||||||
|
let models = models_manager.models();
|
||||||
|
crate::agent::models::platform_for_slug(
|
||||||
|
&models,
|
||||||
|
crate::agent::models::selected_catalog_key_for_spawn(&models, &session_model_id)
|
||||||
|
.as_deref(),
|
||||||
|
&sampling_config.model,
|
||||||
|
)
|
||||||
|
};
|
||||||
let session_pruning_config: crate::config::PruningConfig = memory_config.as_ref().map_or_else(
|
let session_pruning_config: crate::config::PruningConfig = memory_config.as_ref().map_or_else(
|
||||||
|| crate::config::PruningConfig {
|
|| crate::config::PruningConfig {
|
||||||
enabled: false,
|
enabled: false,
|
||||||
@@ -618,16 +630,11 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
watcher,
|
watcher,
|
||||||
stale_claim_secs: watcher_config.stale_claim_secs,
|
stale_claim_secs: watcher_config.stale_claim_secs,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: api_key_provider.clone(),
|
embedding_credentials: crate::auth::credential_provider::embedding_session_credentials(
|
||||||
auth_credentials: auth_manager.as_ref().map(|am| {
|
&embed_base_url,
|
||||||
std::sync::Arc::new(
|
embed_platform,
|
||||||
crate::auth::credential_provider::ShellAuthCredentialProvider::new(
|
&models_manager.credential_authority(),
|
||||||
am.clone(),
|
),
|
||||||
None,
|
|
||||||
None,
|
|
||||||
),
|
|
||||||
) as std::sync::Arc<dyn kigi_auth::AuthCredentialProvider>
|
|
||||||
}),
|
|
||||||
};
|
};
|
||||||
let backend = crate::session::memory::MemoryBackendImpl::from_session_params(
|
let backend = crate::session::memory::MemoryBackendImpl::from_session_params(
|
||||||
storage.clone(),
|
storage.clone(),
|
||||||
@@ -1287,8 +1294,11 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
.map(|mc| mc.embedding.clone())
|
.map(|mc| mc.embedding.clone())
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
let embed_dims = embed_config.dimensions;
|
let embed_dims = embed_config.dimensions;
|
||||||
let sampling_base_url = embed_base_url.clone();
|
// The session's own params, so the background reindex embeds through
|
||||||
let sampling_api_key = embed_api_key.clone();
|
// the same endpoint-scoped credential as every foreground path — a
|
||||||
|
// second locally-built provider would re-derive credentials outside
|
||||||
|
// the chokepoint and would carry a static key with no 401 refresh.
|
||||||
|
let reindex_params = session.memory.backend_params.clone();
|
||||||
let session_id_for_reindex = session_info.id.to_string();
|
let session_id_for_reindex = session_info.id.to_string();
|
||||||
let chunks_added_counter = session.memory.chunks_added.clone();
|
let chunks_added_counter = session.memory.chunks_added.clone();
|
||||||
tokio::task::spawn_local(async move {
|
tokio::task::spawn_local(async move {
|
||||||
@@ -1311,13 +1321,8 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
target : kigi_log::memory_log::TARGET, files = files.len(),
|
target : kigi_log::memory_log::TARGET, files = files.len(),
|
||||||
"MEMORY_REINDEX: background reindex complete"
|
"MEMORY_REINDEX: background reindex complete"
|
||||||
);
|
);
|
||||||
if let Some(api_key) = sampling_api_key
|
if let Some(ref params) = reindex_params
|
||||||
&& let Some(provider) =
|
&& let Some(provider) = params.make_embedding_provider().await
|
||||||
crate::session::memory::embedding::ApiEmbeddingProvider::from_session(
|
|
||||||
&embed_config,
|
|
||||||
sampling_base_url,
|
|
||||||
api_key,
|
|
||||||
)
|
|
||||||
{
|
{
|
||||||
crate::session::memory::embed_missing_chunks(&index, &provider).await;
|
crate::session::memory::embed_missing_chunks(&index, &provider).await;
|
||||||
}
|
}
|
||||||
|
|||||||
+1
-2
@@ -555,8 +555,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
};
|
};
|
||||||
let (event_tx, _event_rx) = tokio::sync::mpsc::unbounded_channel::<SessionEvent>();
|
let (event_tx, _event_rx) = tokio::sync::mpsc::unbounded_channel::<SessionEvent>();
|
||||||
let actor = Arc::new(SessionActor {
|
let actor = Arc::new(SessionActor {
|
||||||
|
|||||||
+2
-4
@@ -393,8 +393,7 @@ fn initial_injection_backend_params_use_override_min_score() {
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
};
|
};
|
||||||
let initial_injection = crate::config::MemoryInitialInjectionConfig {
|
let initial_injection = crate::config::MemoryInitialInjectionConfig {
|
||||||
enabled: true,
|
enabled: true,
|
||||||
@@ -422,8 +421,7 @@ fn initial_injection_backend_params_preserve_default_zero_min_score() {
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
};
|
};
|
||||||
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
|
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
|
||||||
¶ms,
|
¶ms,
|
||||||
|
|||||||
@@ -18,8 +18,7 @@ fn initial_injection_backend_params_use_override_min_score() {
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
};
|
};
|
||||||
let initial_injection = crate::config::MemoryInitialInjectionConfig {
|
let initial_injection = crate::config::MemoryInitialInjectionConfig {
|
||||||
enabled: true,
|
enabled: true,
|
||||||
@@ -47,8 +46,7 @@ fn initial_injection_backend_params_preserve_default_zero_min_score() {
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
};
|
};
|
||||||
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
|
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
|
||||||
¶ms,
|
¶ms,
|
||||||
@@ -524,8 +522,7 @@ async fn create_injection_ready_actor(
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
});
|
});
|
||||||
actor
|
actor
|
||||||
.chat_state_handle
|
.chat_state_handle
|
||||||
|
|||||||
Reference in New Issue
Block a user