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:
@@ -83,6 +83,35 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider {
|
||||
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)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -253,6 +282,38 @@ mod tests {
|
||||
"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).
|
||||
#[tokio::test]
|
||||
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_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(
|
||||
|| crate::config::PruningConfig {
|
||||
enabled: false,
|
||||
@@ -618,16 +630,11 @@ pub(crate) async fn spawn_session_actor(
|
||||
watcher,
|
||||
stale_claim_secs: watcher_config.stale_claim_secs,
|
||||
search_source: "tool",
|
||||
api_key_provider: api_key_provider.clone(),
|
||||
auth_credentials: auth_manager.as_ref().map(|am| {
|
||||
std::sync::Arc::new(
|
||||
crate::auth::credential_provider::ShellAuthCredentialProvider::new(
|
||||
am.clone(),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
) as std::sync::Arc<dyn kigi_auth::AuthCredentialProvider>
|
||||
}),
|
||||
embedding_credentials: crate::auth::credential_provider::embedding_session_credentials(
|
||||
&embed_base_url,
|
||||
embed_platform,
|
||||
&models_manager.credential_authority(),
|
||||
),
|
||||
};
|
||||
let backend = crate::session::memory::MemoryBackendImpl::from_session_params(
|
||||
storage.clone(),
|
||||
@@ -1287,8 +1294,11 @@ pub(crate) async fn spawn_session_actor(
|
||||
.map(|mc| mc.embedding.clone())
|
||||
.unwrap_or_default();
|
||||
let embed_dims = embed_config.dimensions;
|
||||
let sampling_base_url = embed_base_url.clone();
|
||||
let sampling_api_key = embed_api_key.clone();
|
||||
// The session's own params, so the background reindex embeds through
|
||||
// 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 chunks_added_counter = session.memory.chunks_added.clone();
|
||||
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(),
|
||||
"MEMORY_REINDEX: background reindex complete"
|
||||
);
|
||||
if let Some(api_key) = sampling_api_key
|
||||
&& let Some(provider) =
|
||||
crate::session::memory::embedding::ApiEmbeddingProvider::from_session(
|
||||
&embed_config,
|
||||
sampling_base_url,
|
||||
api_key,
|
||||
)
|
||||
if let Some(ref params) = reindex_params
|
||||
&& let Some(provider) = params.make_embedding_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,
|
||||
stale_claim_secs: 60,
|
||||
search_source: "tool",
|
||||
api_key_provider: None,
|
||||
auth_credentials: None,
|
||||
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||
};
|
||||
let (event_tx, _event_rx) = tokio::sync::mpsc::unbounded_channel::<SessionEvent>();
|
||||
let actor = Arc::new(SessionActor {
|
||||
|
||||
+2
-4
@@ -393,8 +393,7 @@ fn initial_injection_backend_params_use_override_min_score() {
|
||||
watcher: None,
|
||||
stale_claim_secs: 60,
|
||||
search_source: "tool",
|
||||
api_key_provider: None,
|
||||
auth_credentials: None,
|
||||
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||
};
|
||||
let initial_injection = crate::config::MemoryInitialInjectionConfig {
|
||||
enabled: true,
|
||||
@@ -422,8 +421,7 @@ fn initial_injection_backend_params_preserve_default_zero_min_score() {
|
||||
watcher: None,
|
||||
stale_claim_secs: 60,
|
||||
search_source: "tool",
|
||||
api_key_provider: None,
|
||||
auth_credentials: None,
|
||||
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||
};
|
||||
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
|
||||
¶ms,
|
||||
|
||||
@@ -18,8 +18,7 @@ fn initial_injection_backend_params_use_override_min_score() {
|
||||
watcher: None,
|
||||
stale_claim_secs: 60,
|
||||
search_source: "tool",
|
||||
api_key_provider: None,
|
||||
auth_credentials: None,
|
||||
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||
};
|
||||
let initial_injection = crate::config::MemoryInitialInjectionConfig {
|
||||
enabled: true,
|
||||
@@ -47,8 +46,7 @@ fn initial_injection_backend_params_preserve_default_zero_min_score() {
|
||||
watcher: None,
|
||||
stale_claim_secs: 60,
|
||||
search_source: "tool",
|
||||
api_key_provider: None,
|
||||
auth_credentials: None,
|
||||
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||
};
|
||||
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
|
||||
¶ms,
|
||||
@@ -524,8 +522,7 @@ async fn create_injection_ready_actor(
|
||||
watcher: None,
|
||||
stale_claim_secs: 60,
|
||||
search_source: "tool",
|
||||
api_key_provider: None,
|
||||
auth_credentials: None,
|
||||
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||
});
|
||||
actor
|
||||
.chat_state_handle
|
||||
|
||||
Reference in New Issue
Block a user