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:
2026-07-26 23:41:45 -04:00
parent 867b3e110b
commit ed8049cf77
7 changed files with 295 additions and 109 deletions
@@ -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;
}
@@ -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 {
@@ -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(
&params,
@@ -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(
&params,
@@ -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