refactor(auth): centralize inference-credential routing in CredentialAuthority
One authority answers 'which credential may ride this request': credential_class / manager_for / credential_for / bearer_resolver_for, keyed by (platform, base_url). SessionCredential is an opaque type with no production constructor, so a new call site cannot re-introduce the session-bearer leak. Platform-scoped tests extended across all bearer channels (session, aux, summary, subagent override). Verified: cargo check --workspace --all-targets clean; kigi-shell and kigi-tui suites green (6611+ tests).
This commit is contained in:
@@ -0,0 +1,568 @@
|
||||
//! THE credential chokepoint.
|
||||
//!
|
||||
//! One authority answers, for every outgoing inference request, the only
|
||||
//! question that matters: **which credential — if any — may ride it?**
|
||||
//! ([`CredentialAuthority::credential_class`]). Before this module the answer was
|
||||
//! re-derived — differently — at every call site (`ModelsManager`, `MvpAgent`,
|
||||
//! `SessionActor`, the aux/summary/subagent paths), and three separate rounds
|
||||
//! of fixes each closed some sites and missed others.
|
||||
//!
|
||||
//! # How omission is structurally prevented
|
||||
//!
|
||||
//! 1. [`SessionCredential`] wraps the bearer and has **no production
|
||||
//! constructor outside this module**. The only function in the crate that
|
||||
//! can build one is [`CredentialAuthority::credential_for`], which *requires*
|
||||
//! `(platform, base_url)` and holds the session's `EndpointsConfig` and
|
||||
//! primary [`AuthManager`] privately.
|
||||
//! 2. Every API that stamps a session credential onto a request —
|
||||
//! `resolve_credentials`, `resolve_aux_model_sampling_config`,
|
||||
//! `try_resolve_model_credentials`,
|
||||
//! `resolve_chat_state_auth_type` — takes `Option<&SessionCredential>`,
|
||||
//! never `Option<&str>`. A new call site therefore *cannot compile* a leak:
|
||||
//! there is no way to produce the value without going through the rule.
|
||||
//! 3. The authority owns the primary manager privately and exposes it only via
|
||||
//! [`CredentialAuthority::manager_for`] /
|
||||
//! [`CredentialAuthority::bearer_resolver_for`], which take the same
|
||||
//! `(platform, base_url)` pair — so the `bearer_resolver` sink is funnelled
|
||||
//! through the identical rule as the `api_key` sink.
|
||||
//! 4. A guard asks [`CredentialAuthority::credential_class`] and MATCHES on the
|
||||
//! answer. There is no second, similarly-named boolean to pick by mistake:
|
||||
//! the round-3 defect (C1) was `takes_session_credential` — *may **a**
|
||||
//! session credential ride?* — paired with a hand-carried PRIMARY bearer,
|
||||
//! and the two predicates that made that pairing expressible are gone.
|
||||
//!
|
||||
//! SECURITY: no token is ever logged, `Debug`-printed or `Display`ed here.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::agent::config::{EndpointsConfig, ModelEntry};
|
||||
use crate::auth::AuthManager;
|
||||
|
||||
/// A session bearer this authority has cleared for one specific request
|
||||
/// endpoint.
|
||||
///
|
||||
/// Opaque by construction: the inner `String` is private, the type is not
|
||||
/// `Debug`/`Clone`-into-`String`, and the only production constructor is
|
||||
/// [`CredentialAuthority::credential_for`]. See the module docs for why that
|
||||
/// matters.
|
||||
pub(crate) struct SessionCredential(String);
|
||||
|
||||
impl SessionCredential {
|
||||
/// The raw bearer. SECURITY: callers stamp this straight onto a request —
|
||||
/// never log it.
|
||||
pub(crate) fn expose(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
|
||||
/// Test-only forgery, so unit tests can exercise the *downstream*
|
||||
/// credential plumbing (`resolve_credentials`' BYOK-vs-session precedence,
|
||||
/// aux config shapes) without standing up an `AuthManager`. Deliberately
|
||||
/// `#[cfg(test)]`: production code has no way to build one.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn for_test(key: &str) -> Self {
|
||||
Self(key.to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
/// WHICH credential — if any — may ride a request routed to a given
|
||||
/// `(platform, base_url)` pair.
|
||||
///
|
||||
/// ONE question with three answers, replacing the two look-alike booleans
|
||||
/// `takes_session_credential` / `takes_primary_credential` (identical
|
||||
/// signatures, near-identical names, opposite answers on a subscription host).
|
||||
/// C1 was caused by asking the first and stamping the credential the second
|
||||
/// describes; with a single classifier a call site must MATCH on the answer, so
|
||||
/// that mistake is no longer expressible.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(crate) enum CredentialClass {
|
||||
/// `platform`'s OWN pooled subscription-OAuth token, at its own registry
|
||||
/// host. NEVER the primary session bearer and never the house key.
|
||||
Pooled,
|
||||
/// The credential that authorizes the SESSION's own coding endpoint:
|
||||
/// the primary (`kimi-code` / platform-less) bearer.
|
||||
///
|
||||
/// Deliberately NOT split into a separate `HouseKey` variant: the house
|
||||
/// `KIGI_API_KEY` is accepted by exactly this endpoint and no other, so it
|
||||
/// rides precisely this class. A fourth variant would re-create the
|
||||
/// two-similar-answers hazard this enum exists to remove.
|
||||
Primary,
|
||||
/// Nothing rides: every API-key registry platform, an OAuth platform
|
||||
/// redirected off its own host, and any endpoint that is not the session's.
|
||||
None,
|
||||
}
|
||||
|
||||
impl CredentialClass {
|
||||
/// Stable label for structured logs. SECURITY: names a channel, never a
|
||||
/// token.
|
||||
pub(crate) fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Pooled => "pooled",
|
||||
Self::Primary => "primary",
|
||||
Self::None => "none",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The single authority over inference-time session credentials.
|
||||
///
|
||||
/// Construct one from the session's EFFECTIVE endpoints plus its primary
|
||||
/// (first-party / Kimi) manager, then ask it about a request. Cheap to build
|
||||
/// (a handful of `Option<String>` clones + an `Arc` clone).
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CredentialAuthority {
|
||||
/// The session's effective `[endpoints]` — config.toml layered over env.
|
||||
/// H3: `EndpointsConfig::proxy_url()` prefers `[endpoints]
|
||||
/// coding_api_base_url` from **config.toml** and only then falls back to
|
||||
/// `KIGI_CODE_BASE_URL`. A predicate that knows only the env var makes a
|
||||
/// managed/enterprise deployment lose its session bearer entirely (401 on
|
||||
/// every turn), so the endpoints are part of the authority's identity, not
|
||||
/// an afterthought.
|
||||
endpoints: EndpointsConfig,
|
||||
/// The primary session manager. PRIVATE: nothing hands it back, so a path
|
||||
/// holding a `CredentialAuthority` cannot reach `current_or_expired()`
|
||||
/// without naming an endpoint.
|
||||
primary: Option<Arc<AuthManager>>,
|
||||
}
|
||||
|
||||
impl CredentialAuthority {
|
||||
pub(crate) fn new(endpoints: EndpointsConfig, primary: Option<Arc<AuthManager>>) -> Self {
|
||||
Self { endpoints, primary }
|
||||
}
|
||||
|
||||
/// THE rule, stated once.
|
||||
///
|
||||
/// - a subscription-OAuth platform (claude-pro-max, openai-codex,
|
||||
/// github-copilot, xai-grok) rides ITS OWN pooled manager — never the
|
||||
/// primary — and only to its own registry host (L10: a
|
||||
/// `[model."claude-pro-max/x"]` override keeps `info.id` but can point
|
||||
/// `base_url` anywhere, and used to ship the Claude OAuth bearer there);
|
||||
/// - `kimi-code` — the one `uses_oauth` platform with no `OAuthConfig` —
|
||||
/// rides the PRIMARY session, and only at the session's own effective
|
||||
/// coding endpoint;
|
||||
/// - every API-key registry platform (deepseek, openai, anthropic,
|
||||
/// moonshot-*, …) rides NOTHING: its credential is that platform's API
|
||||
/// key, already resolved into the catalog entry;
|
||||
/// - a platform-less model (a bare slug or a `[model.*]` block) is decided
|
||||
/// purely by the ENDPOINT — BYOK detection probes `std::env::var` at call
|
||||
/// time, so an unset/mistyped `env_key` must not turn into "send the
|
||||
/// subscription bearer to `api.openai.com`".
|
||||
pub(crate) fn credential_class(
|
||||
&self,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> CredentialClass {
|
||||
match platform {
|
||||
Some(platform) => match platform.oauth() {
|
||||
Some(_) if self.endpoint_is_platform_host(platform, base_url) => {
|
||||
CredentialClass::Pooled
|
||||
}
|
||||
// An OAuth platform pointed at a host that is NOT its own.
|
||||
Some(_) => CredentialClass::None,
|
||||
None if platform.uses_oauth() && self.is_session_coding_endpoint(base_url) => {
|
||||
CredentialClass::Primary
|
||||
}
|
||||
// `kimi-code` off the session's endpoint, and every API-key
|
||||
// registry platform.
|
||||
None => CredentialClass::None,
|
||||
},
|
||||
None if self.is_session_coding_endpoint(base_url) => CredentialClass::Primary,
|
||||
None => CredentialClass::None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The manager behind [`Self::credential_class`]. Derived from the class, so
|
||||
/// the rule is stated exactly once and the two can never disagree.
|
||||
fn governing_manager(
|
||||
&self,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> Option<Arc<AuthManager>> {
|
||||
match self.credential_class(platform, base_url) {
|
||||
CredentialClass::Pooled => {
|
||||
platform
|
||||
.and_then(kigi_models::PlatformId::oauth)
|
||||
.map(|oauth| {
|
||||
crate::auth::oauth_registry::global_manager_for(
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
oauth,
|
||||
)
|
||||
})
|
||||
}
|
||||
CredentialClass::Primary => self.primary.clone(),
|
||||
CredentialClass::None => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether `base_url` is the SESSION's own coding endpoint: the effective
|
||||
/// `[endpoints] coding_api_base_url` from **config.toml** (what a managed /
|
||||
/// enterprise deployment actually sets — H3), the `models_base_url`
|
||||
/// custom-endpoint mode, the `KIGI_CODE_BASE_URL` env override, a loopback
|
||||
/// dev proxy, or the compiled production endpoint.
|
||||
///
|
||||
/// Deliberately NOT [`crate::util::is_first_party_url`], which is
|
||||
/// production-only and would break every custom deployment.
|
||||
fn is_session_coding_endpoint(&self, base_url: &str) -> bool {
|
||||
if crate::util::is_effective_coding_endpoint_url(base_url) {
|
||||
return true;
|
||||
}
|
||||
if crate::util::matches_trusted_base_url(base_url, &self.endpoints.proxy_url()) {
|
||||
return true;
|
||||
}
|
||||
self.endpoints
|
||||
.models_base_url
|
||||
.as_deref()
|
||||
.is_some_and(|models_base| crate::util::matches_trusted_base_url(base_url, models_base))
|
||||
}
|
||||
|
||||
/// Whether `base_url` is `platform`'s own registry host — the guard that
|
||||
/// keeps a subscription-OAuth bearer from riding a redirected `[model.*]`
|
||||
/// override to a third party (L10).
|
||||
fn endpoint_is_platform_host(&self, platform: kigi_models::PlatformId, base_url: &str) -> bool {
|
||||
crate::util::matches_trusted_base_url(base_url, &platform.base_url())
|
||||
}
|
||||
|
||||
/// The `AuthManager` that governs this request's bearer resolution,
|
||||
/// mid-session refresh and 401 recovery — or `None` when no session
|
||||
/// credential may ride (fail fast; never a silent fallback to the primary).
|
||||
pub(crate) fn manager_for(
|
||||
&self,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> Option<Arc<AuthManager>> {
|
||||
self.governing_manager(platform, base_url)
|
||||
}
|
||||
|
||||
/// The session bearer to stamp as this request's `api_key`, or `None`.
|
||||
///
|
||||
/// The ONLY production constructor of [`SessionCredential`].
|
||||
pub(crate) fn credential_for(
|
||||
&self,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> Option<SessionCredential> {
|
||||
self.governing_manager(platform, base_url)
|
||||
.and_then(|am| am.current_or_expired())
|
||||
.map(|auth| SessionCredential(auth.key))
|
||||
}
|
||||
|
||||
/// A live sampler `bearer_resolver` over the governing manager, so the
|
||||
/// request keeps mid-session refresh / 401 recovery against the credential
|
||||
/// that actually belongs to its host.
|
||||
pub(crate) fn bearer_resolver_for(
|
||||
&self,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> Option<kigi_sampler::SharedBearerResolver> {
|
||||
self.manager_for(platform, base_url)
|
||||
.map(crate::session::acp_session::sampler_turn::auth_manager_bearer_resolver)
|
||||
}
|
||||
|
||||
/// [`Self::credential_for`] for a resolved catalog entry: derives the
|
||||
/// platform and the base URL from the SAME entry, so the two can never be
|
||||
/// mismatched by a call site.
|
||||
pub(crate) fn credential_for_model(&self, entry: &ModelEntry) -> Option<SessionCredential> {
|
||||
let info = entry.info();
|
||||
self.credential_for(entry_platform(entry), &info.base_url)
|
||||
}
|
||||
|
||||
/// [`Self::credential_for`] for the catalog model a routing slug resolves
|
||||
/// to. `current_key` is the SESSION's own selected catalog key (see
|
||||
/// [`crate::agent::models::entry_for_slug`]); pass `None` for aux /
|
||||
/// override slugs, which are not the session's selection.
|
||||
///
|
||||
/// M5: a slug that is NOT in the catalog resolves through the SAME endpoint
|
||||
/// rule against the aux fallback endpoint
|
||||
/// (`EndpointsConfig::resolve_inference_base_url`, which is exactly where
|
||||
/// `resolve_aux_model_sampling_config`'s Tier-2 entry routes) instead of
|
||||
/// being handed the primary unconditionally — the old "first-party by
|
||||
/// construction" justification was false once `models_base_url` could point
|
||||
/// anywhere.
|
||||
///
|
||||
/// M6: the platform and the base URL come from ONE
|
||||
/// [`crate::agent::models::entry_for_slug`] lookup, so they can no longer
|
||||
/// disagree (the aux path used to resolve the platform with `current_key`
|
||||
/// and the credential with a separate `find_model_by_id`).
|
||||
pub(crate) fn credential_for_slug(
|
||||
&self,
|
||||
models: &indexmap::IndexMap<String, ModelEntry>,
|
||||
current_key: Option<&str>,
|
||||
slug: &str,
|
||||
) -> Option<SessionCredential> {
|
||||
match crate::agent::models::entry_for_slug(models, current_key, slug) {
|
||||
Some(entry) => self.credential_for_model(entry),
|
||||
None => self.credential_for(None, &self.endpoints.resolve_inference_base_url()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The registry platform a catalog entry belongs to (`info.id` is the managed
|
||||
/// key `{platform}/{model}`). `None` for a bare / `[model.*]` entry.
|
||||
pub(crate) fn entry_platform(entry: &ModelEntry) -> Option<kigi_models::PlatformId> {
|
||||
entry
|
||||
.info()
|
||||
.id
|
||||
.as_deref()
|
||||
.and_then(kigi_models::parse_managed_model_key)
|
||||
.map(|(platform, _)| platform)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::auth::{AuthMode, KimiAuth, KimiCodeConfig};
|
||||
|
||||
/// A primary holding a fixed in-memory bearer. The `TempDir` is returned so
|
||||
/// the caller keeps it alive; the token is read from memory, so on-disk
|
||||
/// contents are irrelevant.
|
||||
fn primary(key: &str) -> (tempfile::TempDir, Arc<AuthManager>) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let manager = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
|
||||
manager.hot_swap(KimiAuth {
|
||||
key: key.to_string(),
|
||||
auth_mode: AuthMode::OAuth,
|
||||
..KimiAuth::test_default()
|
||||
});
|
||||
(dir, manager)
|
||||
}
|
||||
|
||||
fn authority(endpoints: EndpointsConfig, primary: Arc<AuthManager>) -> CredentialAuthority {
|
||||
CredentialAuthority::new(endpoints, Some(primary))
|
||||
}
|
||||
|
||||
fn platform(id: &str) -> kigi_models::PlatformId {
|
||||
kigi_models::PlatformId::parse(id).expect("known platform")
|
||||
}
|
||||
|
||||
/// H3 (REGRESSION): the effective coding endpoint is
|
||||
/// `EndpointsConfig::proxy_url()`, which prefers `[endpoints]
|
||||
/// coding_api_base_url` from **config.toml** — the key the managed-config
|
||||
/// sync writes. A predicate that knows only `KIGI_CODE_BASE_URL` classifies
|
||||
/// such a deployment as third-party, withholds the api_key AND the
|
||||
/// resolver, and 401s on every turn.
|
||||
///
|
||||
/// Revert-to-red: drop the `proxy_url()` arm from
|
||||
/// `is_session_coding_endpoint` (leaving only
|
||||
/// `is_effective_coding_endpoint_url`) and every assertion here fails —
|
||||
/// with NO env var set anywhere in the test.
|
||||
#[test]
|
||||
fn config_toml_coding_endpoint_still_rides_the_session_bearer() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let managed = "https://proxy.acme.com/v1";
|
||||
let auth = authority(
|
||||
EndpointsConfig {
|
||||
coding_api_base_url: Some(managed.to_string()),
|
||||
..EndpointsConfig::default()
|
||||
},
|
||||
kimi.clone(),
|
||||
);
|
||||
assert_eq!(
|
||||
auth.credential_class(None, managed),
|
||||
CredentialClass::Primary,
|
||||
"a [model.*] entry inheriting the managed coding endpoint takes the session bearer"
|
||||
);
|
||||
assert_eq!(
|
||||
auth.credential_for(None, managed)
|
||||
.map(|c| c.expose().to_owned()),
|
||||
Some("kimi-tok".to_string()),
|
||||
"the managed deployment must still receive the session bearer"
|
||||
);
|
||||
assert!(
|
||||
auth.manager_for(None, managed).is_some(),
|
||||
"and must keep a live manager, or it loses refresh and 401 recovery"
|
||||
);
|
||||
// kimi-code entries route to `proxy_url()` too (models_fetch's
|
||||
// `platform_fetch_base`), so the platform arm must honour it as well.
|
||||
assert_eq!(
|
||||
auth.credential_for(Some(platform("kimi-code")), managed)
|
||||
.map(|c| c.expose().to_owned()),
|
||||
Some("kimi-tok".to_string()),
|
||||
);
|
||||
// A DIFFERENT authority (no managed key configured) must NOT trust it.
|
||||
let default_auth = authority(EndpointsConfig::default(), kimi);
|
||||
assert_eq!(
|
||||
default_auth.credential_class(None, managed),
|
||||
CredentialClass::None,
|
||||
"the managed host is only trusted for the session that configured it"
|
||||
);
|
||||
}
|
||||
|
||||
/// The `models_base_url` custom-endpoint mode is equally invisible to the
|
||||
/// env-var-only predicate.
|
||||
#[test]
|
||||
fn config_toml_models_base_url_still_rides_the_session_bearer() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let custom = "https://models.acme.internal/v1";
|
||||
let auth = authority(
|
||||
EndpointsConfig {
|
||||
models_base_url: Some(custom.to_string()),
|
||||
..EndpointsConfig::default()
|
||||
},
|
||||
kimi,
|
||||
);
|
||||
assert_eq!(
|
||||
auth.credential_for(None, custom)
|
||||
.map(|c| c.expose().to_owned()),
|
||||
Some("kimi-tok".to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
/// The compiled production endpoint and loopback proxies are unchanged.
|
||||
#[test]
|
||||
fn production_and_loopback_endpoints_are_unchanged() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let auth = authority(EndpointsConfig::default(), kimi);
|
||||
for url in [
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
"http://127.0.0.1:8080/v1",
|
||||
"http://localhost:3000/v1",
|
||||
"http://[::1]:9000/v1",
|
||||
] {
|
||||
assert_eq!(
|
||||
auth.credential_class(None, url),
|
||||
CredentialClass::Primary,
|
||||
"{url}: the session's own endpoint is byte-identical"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// LEAK guard: every API-key registry platform, and any platform-less model
|
||||
/// on a third-party host, gets NO session credential and NO manager.
|
||||
#[test]
|
||||
fn third_party_endpoints_never_receive_the_primary() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let auth = authority(EndpointsConfig::default(), kimi);
|
||||
for id in [
|
||||
"deepseek",
|
||||
"openai",
|
||||
"anthropic",
|
||||
"moonshot-cn",
|
||||
"moonshot-ai",
|
||||
] {
|
||||
let p = platform(id);
|
||||
assert!(
|
||||
auth.credential_for(Some(p), &p.base_url()).is_none(),
|
||||
"LEAK: {id} is an API-key platform — no session bearer may ride there"
|
||||
);
|
||||
assert!(auth.manager_for(Some(p), &p.base_url()).is_none());
|
||||
}
|
||||
for url in [
|
||||
"https://api.openai.com/v1",
|
||||
"https://api.deepseek.com/v1",
|
||||
"https://api.moonshot.cn/v1",
|
||||
"",
|
||||
] {
|
||||
assert!(
|
||||
auth.credential_for(None, url).is_none(),
|
||||
"LEAK: {url} is a third-party host"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// L10: an OAuth platform whose `[model.*]` override redirects `base_url`
|
||||
/// to a third-party host keeps `info.id` — and must NOT ship that
|
||||
/// platform's pooled OAuth bearer there.
|
||||
#[tokio::test]
|
||||
async fn oauth_platform_redirected_to_a_third_party_host_gets_nothing() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let auth = authority(EndpointsConfig::default(), kimi);
|
||||
for id in [
|
||||
"claude-pro-max",
|
||||
"openai-codex",
|
||||
"github-copilot",
|
||||
"xai-grok",
|
||||
] {
|
||||
let p = platform(id);
|
||||
assert!(
|
||||
auth.manager_for(Some(p), &p.base_url()).is_some(),
|
||||
"{id} keeps its own pooled manager on its own host"
|
||||
);
|
||||
assert!(
|
||||
auth.manager_for(Some(p), "https://third.party/v1")
|
||||
.is_none(),
|
||||
"LEAK: {id} redirected to a third-party host must ship no bearer"
|
||||
);
|
||||
assert_eq!(
|
||||
auth.credential_class(Some(p), "https://third.party/v1"),
|
||||
CredentialClass::None,
|
||||
"LEAK: {id} redirected to a third-party host takes no session credential"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// C1 — a subscription platform's own host DOES take a session credential,
|
||||
/// but it is that platform's POOLED token, never the primary / house key.
|
||||
/// Two look-alike booleans used to encode this, and picking the wrong one is
|
||||
/// the whole defect; one classifier makes the distinction impossible to
|
||||
/// mis-read.
|
||||
#[test]
|
||||
fn a_subscription_host_classifies_pooled_never_primary() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let auth = authority(EndpointsConfig::default(), kimi);
|
||||
for id in [
|
||||
"claude-pro-max",
|
||||
"openai-codex",
|
||||
"github-copilot",
|
||||
"xai-grok",
|
||||
] {
|
||||
let p = platform(id);
|
||||
assert_eq!(
|
||||
auth.credential_class(Some(p), &p.base_url()),
|
||||
CredentialClass::Pooled,
|
||||
"LEAK: {id}'s own host takes its POOLED token — never the primary / house key"
|
||||
);
|
||||
}
|
||||
// Every API-key registry platform: nothing at all.
|
||||
for id in ["deepseek", "openai", "anthropic", "moonshot-cn"] {
|
||||
let p = platform(id);
|
||||
assert_eq!(
|
||||
auth.credential_class(Some(p), &p.base_url()),
|
||||
CredentialClass::None
|
||||
);
|
||||
}
|
||||
// The primary channel is unchanged: kimi-code and a platform-less model
|
||||
// on the session's own endpoint, and nothing on a third-party host.
|
||||
for url in [
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
"http://127.0.0.1:8080/v1",
|
||||
] {
|
||||
assert_eq!(auth.credential_class(None, url), CredentialClass::Primary);
|
||||
assert_eq!(
|
||||
auth.credential_class(Some(platform("kimi-code")), url),
|
||||
CredentialClass::Primary
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
auth.credential_class(None, "https://api.openai.com/v1"),
|
||||
CredentialClass::None
|
||||
);
|
||||
}
|
||||
|
||||
/// The four subscription-OAuth platforms draw from their OWN pooled
|
||||
/// managers — never the primary Kimi one, even under a Kimi session.
|
||||
#[tokio::test]
|
||||
async fn oauth_platforms_never_resolve_the_primary() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let auth = authority(EndpointsConfig::default(), kimi.clone());
|
||||
for id in [
|
||||
"claude-pro-max",
|
||||
"openai-codex",
|
||||
"github-copilot",
|
||||
"xai-grok",
|
||||
] {
|
||||
let p = platform(id);
|
||||
let resolved = auth
|
||||
.manager_for(Some(p), &p.base_url())
|
||||
.expect("pooled manager");
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &kimi),
|
||||
"{id} must NOT resolve the primary Kimi manager"
|
||||
);
|
||||
assert_ne!(
|
||||
auth.credential_for(Some(p), &p.base_url())
|
||||
.map(|c| c.expose().to_owned()),
|
||||
Some("kimi-tok".to_string()),
|
||||
"{id} must never receive the primary Kimi bearer"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
pub(crate) mod attribution;
|
||||
mod config;
|
||||
pub(crate) mod credential_authority;
|
||||
pub mod credential_provider;
|
||||
pub(crate) mod device;
|
||||
pub mod device_code;
|
||||
|
||||
@@ -15,8 +15,13 @@
|
||||
//! on-disk token stays fresh and a 401 recovers via the provider's own manager.
|
||||
//! Managers are built ON DEMAND from the on-disk token ([`global_manager_for`]),
|
||||
//! so a login landing AFTER a session spawned self-heals — no frozen per-session
|
||||
//! snapshot. [`manager_for_model`] routes a managed catalog key to the pool
|
||||
//! (oauth platform) or to the session's primary (everything else).
|
||||
//! snapshot.
|
||||
//!
|
||||
//! ROUTING LIVES ELSEWHERE. This module is only the pool; the decision of which
|
||||
//! credential governs a request belongs to the single chokepoint,
|
||||
//! [`crate::auth::credential_authority::CredentialAuthority`]. Keeping the two
|
||||
//! apart is deliberate: three rounds of leaks came from routing rules being
|
||||
//! re-derived per call site.
|
||||
//!
|
||||
//! SECURITY: access/refresh tokens and resolved bearers are NEVER logged here.
|
||||
|
||||
@@ -38,26 +43,53 @@ fn oauth_manager_pool() -> &'static Mutex<HashMap<&'static str, Arc<AuthManager>
|
||||
POOL.get_or_init(|| Mutex::new(HashMap::new()))
|
||||
}
|
||||
|
||||
/// The kigi home every OAuth-pool call site resolves from. Single definition so
|
||||
/// the pool, the aux/summary token routing and the session's inference manager
|
||||
/// can never read different homes.
|
||||
/// The kigi home EVERY OAuth-provider construction in this crate resolves from
|
||||
/// — the pool, the catalog fetch's per-platform token resolution, the
|
||||
/// aux/summary token routing and the session's inference manager — so they can
|
||||
/// never read different homes.
|
||||
///
|
||||
/// Production: [`crate::util::kigi_home::kigi_home`]. LIB TESTS: a
|
||||
/// process-lifetime `TempDir`, unconditionally — the pool is process-global and
|
||||
/// every manager it builds starts a never-cancelled proactive-refresh loop, so
|
||||
/// a unit test resolving the real `~/.kigi` would read the developer's stored
|
||||
/// OAuth tokens and, 60 s later, fire REAL refresh requests against them.
|
||||
/// Deliberately not a per-test opt-in that can be forgotten: `kigi_home()` is
|
||||
/// itself a `OnceLock` an earlier test has usually already resolved to the real
|
||||
/// home, so setting `KIGI_SHARE_DIR` in a test cannot pin it after the fact.
|
||||
/// per-process path under the system temp dir that is deliberately **never
|
||||
/// created**. The pool is process-global and every manager it builds starts a
|
||||
/// never-cancelled proactive-refresh loop, so a unit test resolving the real
|
||||
/// `~/.kigi` would read the developer's stored OAuth tokens and, 60 s later,
|
||||
/// fire REAL refresh requests against them. Deliberately not a per-test opt-in
|
||||
/// that can be forgotten: `kigi_home()` is itself a `OnceLock` an earlier test
|
||||
/// has usually already resolved to the real home, so setting `KIGI_SHARE_DIR`
|
||||
/// in a test cannot pin it after the fact.
|
||||
///
|
||||
/// M4 — THE LIMIT, STATED: `cfg(test)` is set only for THIS crate's `--lib`
|
||||
/// tests. `crates/codegen/kigi-shell/tests/*.rs` link the library built WITHOUT
|
||||
/// it, so for an integration test this resolves the real home unless that test
|
||||
/// binary itself isolates one, which it must do through the two overrides the
|
||||
/// auth stack already honours and BEFORE anything resolves `kigi_home()`:
|
||||
/// `KIGI_SHARE_DIR` (read by `kigi_home()`, a `OnceLock`) or `KIGI_AUTH_PATH`
|
||||
/// (read by [`AuthManager::new_oauth_provider`], which pins the token file
|
||||
/// outright and so overrides this home entirely). 12 of the 28 integration
|
||||
/// binaries under `crates/codegen/kigi-shell/tests/` set `KIGI_SHARE_DIR`; the
|
||||
/// other 16 never reach an OAuth-platform inference path today, which is a
|
||||
/// property of those tests, not a guarantee of this function. No
|
||||
/// production-readable env override is added here on purpose: a knob that
|
||||
/// redirects where OAuth tokens are read from is not worth a test convenience.
|
||||
///
|
||||
/// M8: this used to be a `static OnceLock<TempDir>`. Statics are never dropped,
|
||||
/// so that leaked one temp directory per test binary — against the project's
|
||||
/// "tests are TempDir self-cleaning" discipline. Nothing is created here
|
||||
/// instead, and nothing in the lib-test suite creates it: a manager reads a
|
||||
/// missing `auth.json` as "no session", and the only two paths that WRITE one
|
||||
/// are a successful token refresh (which needs a stored refresh token that by
|
||||
/// construction does not exist here) and a completed device login
|
||||
/// ([`crate::agent::mvp_agent::MvpAgent::authenticate_oauth_platform`], which
|
||||
/// M4 repointed at this same home). Both require the network, so no lib test
|
||||
/// performs either. That is an observation about the suite, not an invariant of
|
||||
/// this function — [`tests::test_pool_home_is_disposable_and_never_the_real_home`]
|
||||
/// asserts the directory does not exist and is the tripwire if one ever does
|
||||
/// (the path is per-PROCESS under the system temp dir, so the blast radius of a
|
||||
/// future login-driving test is one disposable directory, never `~/.kigi`).
|
||||
pub(crate) fn pool_home() -> std::path::PathBuf {
|
||||
#[cfg(test)]
|
||||
{
|
||||
static TEST_HOME: OnceLock<tempfile::TempDir> = OnceLock::new();
|
||||
TEST_HOME
|
||||
.get_or_init(|| tempfile::tempdir().expect("tempdir for the test OAuth pool"))
|
||||
.path()
|
||||
.to_path_buf()
|
||||
std::env::temp_dir().join(format!("kigi-oauth-pool-test-{}", std::process::id()))
|
||||
}
|
||||
#[cfg(not(test))]
|
||||
crate::util::kigi_home::kigi_home()
|
||||
@@ -89,411 +121,73 @@ pub(crate) fn global_manager_for(
|
||||
manager
|
||||
}
|
||||
|
||||
/// The `AuthManager` that governs INFERENCE auth for `managed_key`
|
||||
/// (`{platform}/{model}`, e.g. `xai-grok/grok-4-latest`).
|
||||
///
|
||||
/// A generic device-code OAuth platform routes to ITS OWN scope-keyed manager
|
||||
/// from the process-global pool ([`global_manager_for`], built on demand from
|
||||
/// the on-disk token); every other key (Kimi, API-key platforms, `[model.*]`
|
||||
/// entries, or an unprefixed bare id) routes to `primary`.
|
||||
///
|
||||
/// The pool is the single source of truth — there is no per-session snapshot to
|
||||
/// freeze at spawn, so a grok login that happens AFTER a session spawned is
|
||||
/// resolved correctly on the next grok turn. A grok key NEVER resolves to
|
||||
/// `primary`: even before the user logs into grok the pooled manager simply
|
||||
/// holds no token (its bearer / api_key is then `None`), so the Kimi
|
||||
/// subscription bearer can never reach a third-party host — fail-fast, never a
|
||||
/// silent fallback to the Kimi manager.
|
||||
pub(crate) fn manager_for_model(
|
||||
kigi_home: &Path,
|
||||
managed_key: &str,
|
||||
primary: Option<&Arc<AuthManager>>,
|
||||
) -> Option<Arc<AuthManager>> {
|
||||
if let Some((platform, _)) = kigi_models::parse_managed_model_key(managed_key)
|
||||
&& let Some(oauth) = platform.oauth()
|
||||
{
|
||||
return Some(global_manager_for(kigi_home, oauth));
|
||||
}
|
||||
primary.cloned()
|
||||
}
|
||||
|
||||
/// The SESSION token (the raw bearer/key string) that may ride an INFERENCE
|
||||
/// request routed to `platform` at `base_url`. Used by the aux-model, summary
|
||||
/// and subagent-override wire paths, where the result is stamped straight into
|
||||
/// [`crate::agent::config::resolve_credentials`] as the request's `api_key`.
|
||||
///
|
||||
/// - a generic device-code OAuth platform (xai-grok, claude-pro-max,
|
||||
/// github-copilot, openai-codex) draws from ITS OWN pooled manager; when that
|
||||
/// provider has no stored session the result is `None` — never `primary`;
|
||||
/// - `kimi-code`, and a platform-less model whose endpoint IS the session's own
|
||||
/// coding endpoint (incl. a `KIGI_CODE_BASE_URL` deployment or a loopback
|
||||
/// proxy), yield the primary's current-or-expired token — byte-identical to
|
||||
/// reading it directly;
|
||||
/// - every API-key registry platform, and every `[model.*]` block pointed at a
|
||||
/// third-party host, yields `None`. Handing them `primary` put the user's
|
||||
/// Kimi subscription bearer on `api.deepseek.com` / `api.moonshot.cn` / …
|
||||
/// as the request's `api_key`.
|
||||
///
|
||||
/// SECURITY: the resolved token is never logged.
|
||||
pub(crate) fn session_key_for_endpoint(
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
primary: Option<&Arc<AuthManager>>,
|
||||
) -> Option<String> {
|
||||
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
|
||||
return global_manager_for(&pool_home(), oauth)
|
||||
.current_or_expired()
|
||||
.map(|a| a.key);
|
||||
}
|
||||
if !crate::agent::auth_method::platform_takes_session_credential(platform, base_url) {
|
||||
return None;
|
||||
}
|
||||
primary
|
||||
.and_then(|am| am.current_or_expired())
|
||||
.map(|a| a.key)
|
||||
}
|
||||
|
||||
/// [`session_key_for_endpoint`] for the catalog model whose routing slug (or
|
||||
/// catalog key) is `slug`.
|
||||
///
|
||||
/// A slug absent from the catalog keeps the pre-registry behaviour: the aux
|
||||
/// resolver's Tier-2 fallback builds its entry against
|
||||
/// `EndpointsConfig::resolve_inference_base_url` (first-party), so the primary
|
||||
/// still governs.
|
||||
pub(crate) fn session_key_for_catalog_model(
|
||||
models: &indexmap::IndexMap<String, crate::agent::config::ModelEntry>,
|
||||
slug: &str,
|
||||
primary: Option<&Arc<AuthManager>>,
|
||||
) -> Option<String> {
|
||||
let Some(entry) = crate::agent::config::find_model_by_id(models, slug) else {
|
||||
return primary
|
||||
.and_then(|am| am.current_or_expired())
|
||||
.map(|a| a.key);
|
||||
};
|
||||
let info = entry.info();
|
||||
session_key_for_endpoint(
|
||||
info.id
|
||||
.as_deref()
|
||||
.and_then(kigi_models::parse_managed_model_key)
|
||||
.map(|(platform, _)| platform),
|
||||
&info.base_url,
|
||||
primary,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::auth::KimiCodeConfig;
|
||||
use crate::auth::{AuthMode, KimiAuth};
|
||||
|
||||
/// A Kimi manager holding a fixed in-memory bearer, standing in for a
|
||||
/// session's primary. The `TempDir` is returned so the caller keeps it
|
||||
/// alive; the token is read from memory (`current_or_expired`), so disk
|
||||
/// contents are irrelevant to the assertion.
|
||||
fn primary_with_token(key: &str) -> (tempfile::TempDir, Arc<AuthManager>) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let manager = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
|
||||
manager.hot_swap(KimiAuth {
|
||||
key: key.to_string(),
|
||||
auth_mode: AuthMode::OAuth,
|
||||
..KimiAuth::test_default()
|
||||
});
|
||||
(dir, manager)
|
||||
}
|
||||
|
||||
fn xai_oauth() -> &'static kigi_models::OAuthConfig {
|
||||
kigi_models::PlatformId::XaiGrok
|
||||
fn oauth_for(id: &str) -> &'static kigi_models::OAuthConfig {
|
||||
kigi_models::PlatformId::parse(id)
|
||||
.expect("known platform")
|
||||
.oauth()
|
||||
.expect("xai-grok carries an OAuthConfig")
|
||||
.expect("subscription-OAuth platform carries an OAuthConfig")
|
||||
}
|
||||
|
||||
fn claude_oauth() -> &'static kigi_models::OAuthConfig {
|
||||
kigi_models::PlatformId::ClaudeProMax
|
||||
.oauth()
|
||||
.expect("claude-pro-max carries an OAuthConfig")
|
||||
}
|
||||
|
||||
fn copilot_oauth() -> &'static kigi_models::OAuthConfig {
|
||||
kigi_models::PlatformId::GithubCopilot
|
||||
.oauth()
|
||||
.expect("github-copilot carries an OAuthConfig")
|
||||
}
|
||||
|
||||
fn codex_oauth() -> &'static kigi_models::OAuthConfig {
|
||||
kigi_models::PlatformId::OpenaiCodex
|
||||
.oauth()
|
||||
.expect("openai-codex carries an OAuthConfig")
|
||||
}
|
||||
|
||||
/// `session_key_for_endpoint` for a managed catalog key, resolving the
|
||||
/// platform and its base URL from the registry exactly as the catalog entry
|
||||
/// would.
|
||||
fn session_key_for_key(
|
||||
managed_key: &str,
|
||||
primary: Option<&Arc<AuthManager>>,
|
||||
) -> Option<String> {
|
||||
let platform = kigi_models::parse_managed_model_key(managed_key).map(|(p, _)| p);
|
||||
let base_url = platform
|
||||
.map(kigi_models::PlatformId::base_url)
|
||||
.unwrap_or_default();
|
||||
session_key_for_endpoint(platform, &base_url, primary)
|
||||
}
|
||||
|
||||
/// An `openai-codex/<model>` turn resolves to the process-global pooled
|
||||
/// openai-codex manager (its OWN `oauth/openai-codex` scope), NEVER the
|
||||
/// primary Kimi manager — the same leak-safe routing as the other OAuth
|
||||
/// platforms, and a DISTINCT pool entry from each. Fail-fast: even with a
|
||||
/// Kimi primary, a codex turn never yields the Kimi bearer.
|
||||
/// Each subscription-OAuth platform gets its OWN process-global pooled
|
||||
/// manager, and no two share one. (Which credential governs a REQUEST is
|
||||
/// not decided here — see
|
||||
/// [`crate::auth::credential_authority::CredentialAuthority`].)
|
||||
#[tokio::test]
|
||||
async fn openai_codex_model_resolves_to_its_own_manager_not_kimi() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
async fn every_oauth_scope_gets_its_own_pooled_manager() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let resolved = manager_for_model(home.path(), "openai-codex/gpt-5.5", Some(&kimi))
|
||||
.expect("openai-codex model resolves to its pooled manager");
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &kimi),
|
||||
"openai-codex must NOT resolve to the Kimi manager"
|
||||
);
|
||||
assert!(
|
||||
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), codex_oauth())),
|
||||
"openai-codex must resolve to its OWN process-global pooled manager"
|
||||
);
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &global_manager_for(home.path(), copilot_oauth())),
|
||||
"openai-codex and github-copilot must not share a pooled manager"
|
||||
);
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &global_manager_for(home.path(), claude_oauth())),
|
||||
"openai-codex and claude-pro-max must not share a pooled manager"
|
||||
);
|
||||
assert_ne!(
|
||||
session_key_for_key("openai-codex/gpt-5.5", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"an openai-codex model must never receive the primary Kimi token"
|
||||
);
|
||||
}
|
||||
|
||||
/// A `github-copilot/<model>` turn resolves to the process-global pooled
|
||||
/// github-copilot manager (its OWN `oauth/github-copilot` scope), NEVER the
|
||||
/// primary Kimi manager — the same leak-safe routing as xai-grok /
|
||||
/// claude-pro-max, and a DISTINCT pool entry from either.
|
||||
#[tokio::test]
|
||||
async fn github_copilot_model_resolves_to_its_own_manager_not_kimi() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let resolved = manager_for_model(home.path(), "github-copilot/gpt-4.1", Some(&kimi))
|
||||
.expect("github-copilot model resolves to its pooled manager");
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &kimi),
|
||||
"github-copilot must NOT resolve to the Kimi manager"
|
||||
);
|
||||
assert!(
|
||||
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), copilot_oauth())),
|
||||
"github-copilot must resolve to its OWN process-global pooled manager"
|
||||
);
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &global_manager_for(home.path(), claude_oauth())),
|
||||
"github-copilot and claude-pro-max must not share a pooled manager"
|
||||
);
|
||||
// Fail-fast: even with a Kimi primary, a copilot turn never yields the
|
||||
// Kimi bearer — it draws from the copilot pool (its own token, or None).
|
||||
assert_ne!(
|
||||
session_key_for_key("github-copilot/gpt-4.1", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"a github-copilot model must never receive the primary Kimi token"
|
||||
);
|
||||
}
|
||||
|
||||
/// A `claude-pro-max/<model>` turn resolves to the process-global pooled
|
||||
/// claude-pro-max manager (its OWN `oauth/claude-pro-max` scope), NEVER the
|
||||
/// primary Kimi manager — the same leak-safe routing as xai-grok, and a
|
||||
/// DISTINCT pool entry from the xai manager.
|
||||
#[tokio::test]
|
||||
async fn claude_pro_max_model_resolves_to_its_own_manager_not_kimi() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let resolved =
|
||||
manager_for_model(home.path(), "claude-pro-max/claude-opus-4-8", Some(&kimi))
|
||||
.expect("claude-pro-max model resolves to its pooled manager");
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &kimi),
|
||||
"claude-pro-max must NOT resolve to the Kimi manager"
|
||||
);
|
||||
assert!(
|
||||
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), claude_oauth())),
|
||||
"claude-pro-max must resolve to its OWN process-global pooled manager"
|
||||
);
|
||||
// And it is a DIFFERENT manager than xai-grok's pooled one.
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &global_manager_for(home.path(), xai_oauth())),
|
||||
"claude-pro-max and xai-grok must not share a pooled manager"
|
||||
);
|
||||
}
|
||||
|
||||
/// Fail-fast (no Kimi fallback): a claude-pro-max key with a Kimi primary
|
||||
/// never yields the Kimi session token — it draws from the claude pool (its
|
||||
/// own token, or `None`), so the Kimi bearer can never reach api.anthropic.
|
||||
#[tokio::test]
|
||||
async fn session_key_for_claude_pro_max_is_never_the_kimi_primary() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
assert_ne!(
|
||||
session_key_for_key("claude-pro-max/claude-opus-4-8", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"a claude-pro-max model must never receive the primary Kimi session token"
|
||||
);
|
||||
}
|
||||
|
||||
/// A non-OAuth managed key (moonshot-cn/…) and an unprefixed bare id both
|
||||
/// route to the primary Kimi manager — the Kimi / first-party path is
|
||||
/// untouched and never consults the pool (no runtime needed).
|
||||
#[test]
|
||||
fn non_oauth_and_bare_models_route_to_primary() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
for key in ["moonshot-cn/kimi-k2", "kimi-k2-0905-preview"] {
|
||||
let resolved = manager_for_model(home.path(), key, Some(&kimi))
|
||||
.expect("non-oauth key routes to the primary");
|
||||
let ids = [
|
||||
"xai-grok",
|
||||
"claude-pro-max",
|
||||
"github-copilot",
|
||||
"openai-codex",
|
||||
];
|
||||
let managers: Vec<_> = ids
|
||||
.iter()
|
||||
.map(|id| global_manager_for(home.path(), oauth_for(id)))
|
||||
.collect();
|
||||
for (i, a) in managers.iter().enumerate() {
|
||||
assert!(
|
||||
Arc::ptr_eq(&resolved, &kimi),
|
||||
"{key} must resolve to the primary manager"
|
||||
Arc::ptr_eq(a, &global_manager_for(home.path(), oauth_for(ids[i]))),
|
||||
"{}: the pool must return the SAME manager for a scope",
|
||||
ids[i]
|
||||
);
|
||||
assert_eq!(resolved.current_or_expired().unwrap().key, "kimi-tok");
|
||||
for (j, b) in managers.iter().enumerate() {
|
||||
if i != j {
|
||||
assert!(
|
||||
!Arc::ptr_eq(a, b),
|
||||
"{} and {} must not share a pooled manager",
|
||||
ids[i],
|
||||
ids[j]
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The primary being `None` (test / BYOK sessions) still yields `None` for a
|
||||
/// non-oauth key, never a panic — and without touching the pool.
|
||||
#[test]
|
||||
fn none_primary_is_passed_through_for_non_oauth() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
assert!(manager_for_model(home.path(), "kimi-k2", None).is_none());
|
||||
}
|
||||
|
||||
/// An `xai-grok/<model>` turn resolves to the process-global pooled xai
|
||||
/// manager, NEVER the primary Kimi manager — the pool is the single source.
|
||||
#[tokio::test]
|
||||
async fn grok_model_resolves_to_pooled_xai_manager_not_kimi() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi))
|
||||
.expect("grok model resolves to the pooled xai manager");
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &kimi),
|
||||
"grok model must NOT resolve to the Kimi manager"
|
||||
);
|
||||
assert!(
|
||||
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), xai_oauth())),
|
||||
"grok model must resolve to the process-global pooled xai manager"
|
||||
);
|
||||
}
|
||||
|
||||
/// Facet B guard: the resolver routes purely by the model's platform, with
|
||||
/// no auth-method input — so even when the session's primary is a Kimi
|
||||
/// (session) manager holding "kimi-tok", a grok model never resolves that
|
||||
/// Kimi token.
|
||||
#[tokio::test]
|
||||
async fn grok_model_under_kimi_primary_never_yields_kimi_token() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-fast", Some(&kimi))
|
||||
.expect("grok model resolves to its own pooled manager regardless of primary");
|
||||
assert!(!Arc::ptr_eq(&resolved, &kimi));
|
||||
assert_ne!(
|
||||
resolved.current_or_expired().map(|a| a.key),
|
||||
Some("kimi-tok".to_string()),
|
||||
"the Kimi bearer must never be what a grok turn resolves"
|
||||
);
|
||||
}
|
||||
|
||||
/// Fail-fast: a grok key resolves to the pooled xai manager (never the Kimi
|
||||
/// primary) even with no stored grok session in the pool — the pooled
|
||||
/// manager then simply holds no token, so nothing (least of all the Kimi
|
||||
/// bearer) is sent to api.x.ai.
|
||||
#[tokio::test]
|
||||
async fn grok_never_falls_back_to_kimi_primary() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi))
|
||||
.expect("grok routes to the pooled xai manager, not None");
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &kimi),
|
||||
"an OAuth platform must never fall back to the primary Kimi manager"
|
||||
);
|
||||
}
|
||||
|
||||
/// `session_key_for_endpoint`: the endpoints that genuinely ride the
|
||||
/// PRIMARY session — `kimi-code` (the subscription channel) and a
|
||||
/// platform-less model routed at the session's own coding endpoint (a
|
||||
/// `KIGI_CODE_BASE_URL` deployment or a loopback dev proxy) — yield the
|
||||
/// primary token exactly as reading it directly would. No runtime / pool
|
||||
/// touched.
|
||||
#[test]
|
||||
fn session_key_for_the_sessions_own_endpoint_is_the_primary_token() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
assert_eq!(
|
||||
session_key_for_key("kimi-code/kimi-for-coding", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"kimi-code rides the primary session, unchanged"
|
||||
);
|
||||
for url in [
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
"http://127.0.0.1:4000/v1",
|
||||
] {
|
||||
assert_eq!(
|
||||
session_key_for_endpoint(None, url, Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"{url}: a platform-less model on the session's own endpoint is unchanged"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// LEAK guard (aux / summary / subagent-override `api_key` channel): an
|
||||
/// API-key registry platform, and a `[model.*]` block pointed at a
|
||||
/// third-party host, must yield NO session token. Handing them the primary
|
||||
/// stamped the user's Kimi subscription bearer onto `api.moonshot.cn` /
|
||||
/// `api.deepseek.com` as the request's `api_key` — the channel the
|
||||
/// `bearer_resolver` guard alone does not close.
|
||||
/// M8: the test pool home is a per-process path that is never created, so a
|
||||
/// test binary leaves nothing behind (and never resolves the developer's
|
||||
/// real `~/.kigi`, whose stored OAuth tokens the pool would otherwise read
|
||||
/// and proactively refresh over the network).
|
||||
///
|
||||
/// Revert-to-red: dropping the `platform_takes_session_credential` term
|
||||
/// from `session_key_for_endpoint` returns `Some("kimi-tok")` here.
|
||||
/// This is also the tripwire for the cleanup claim in [`pool_home`]: a
|
||||
/// completed device login through `authenticate_oauth_platform` WOULD create
|
||||
/// this directory, so if a lib test ever drives one, this assertion fires
|
||||
/// and the cleanup has to be added rather than silently regressing the
|
||||
/// "tests are TempDir self-cleaning" discipline.
|
||||
#[test]
|
||||
fn session_key_for_a_third_party_endpoint_is_never_the_primary_token() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
for key in [
|
||||
"moonshot-cn/kimi-k2",
|
||||
"deepseek/deepseek-chat",
|
||||
"openai/gpt-5",
|
||||
] {
|
||||
assert_eq!(
|
||||
session_key_for_key(key, Some(&kimi)),
|
||||
None,
|
||||
"LEAK: {key} is an API-key platform — no session token may ride there"
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
session_key_for_endpoint(None, "https://api.openai.com/v1", Some(&kimi)),
|
||||
None,
|
||||
"LEAK: a [model.*] block on a third-party host gets no session token"
|
||||
fn test_pool_home_is_disposable_and_never_the_real_home() {
|
||||
let home = pool_home();
|
||||
assert!(
|
||||
home.starts_with(std::env::temp_dir()),
|
||||
"the test pool home must live under the system temp dir, got {home:?}"
|
||||
);
|
||||
}
|
||||
|
||||
/// LEAK guard (aux-model + subagent-override token routing): a grok key with
|
||||
/// a Kimi primary NEVER yields the primary Kimi token — it draws from the
|
||||
/// pooled xai manager (its own token, or `None`). This is the exact source
|
||||
/// the aux `session_key` and the override `session_key` now use.
|
||||
#[tokio::test]
|
||||
async fn session_key_for_grok_is_never_the_kimi_primary() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
assert_ne!(
|
||||
session_key_for_key("xai-grok/grok-4-latest", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"a grok aux/override model must never receive the primary Kimi session token"
|
||||
);
|
||||
// Even with `None` primary the routing is unchanged: grok → pool, never a panic.
|
||||
assert_ne!(
|
||||
session_key_for_key("xai-grok/grok-4-fast", None),
|
||||
Some("kimi-tok".to_string()),
|
||||
assert!(
|
||||
!home.exists(),
|
||||
"the test pool home must not be created — nothing to clean up"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user