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:
2026-07-22 15:12:00 -04:00
parent 2d00a4e6e6
commit 48d89c7830
35 changed files with 3158 additions and 1099 deletions
+76 -26
View File
@@ -259,40 +259,90 @@ edges stay deterministic Rust. The harness appends a terminal
token (`resolve_generic_oauth_tokens`, refreshed on expiry) and routes token (`resolve_generic_oauth_tokens`, refreshed on expiry) and routes
`platform.oauth().is_some()``platform.base_url()` (kimi-code alone → `platform.oauth().is_some()``platform.base_url()` (kimi-code alone →
`proxy_url()`). Tokens/codes/verifiers are NEVER logged. `proxy_url()`). Tokens/codes/verifiers are NEVER logged.
- INFERENCE-AUTH INVARIANT (security): a SESSION bearer may ride a request ONLY - INFERENCE-AUTH CHOKEPOINT (security): `auth::credential_authority::CredentialAuthority`
when the endpoint's platform `uses_oauth()``kimi-code` (the primary is the ONE authority answering "WHICH credential — if any — may ride this
session) or one of the four subscription-OAuth platforms (their OWN pooled request?". It holds the session's effective `EndpointsConfig` and
`AuthManager` via `oauth_registry::manager_for_model`, which is why they keep its primary `AuthManager` PRIVATELY and answers only
a live `bearer_resolver` and mid-session refresh despite non-first-party base `(platform, base_url)` questions: `credential_class``CredentialClass::{Pooled,
URLs). Every API-key registry platform is refused, because Primary, None}` (the outer term of `auth_method::session_token_auth_gate`),
`manager_for_model` falls through to the PRIMARY manager for a non-OAuth key `manager_for` (the governing manager
and `SamplingClient::post` REPLACES the request's auth header from the for refresh / 401 recovery), `credential_for` (the request's `api_key`) and
resolver. A model with NO platform (a bare slug / `[model.*]` entry) is `bearer_resolver_for` (the aux/summary/session resolver). Do NOT re-derive this
decided by the ENDPOINT — `util::is_effective_coding_endpoint_url` (the rule anywhere else — three rounds of leaks came from exactly that.
effective `KIGI_CODE_BASE_URL` deployment, loopback dev proxies, and the The rule: a subscription-OAuth platform rides ITS OWN pooled `AuthManager` (so
compiled production endpoint), never a blanket allow: BYOK is it keeps a live `bearer_resolver` and mid-session refresh despite a
`has_own_credentials()`, which probes `std::env::var` at call time, so a non-first-party base URL) and ONLY at its own `platform.base_url()` a
`[model.*]` block with an unset `env_key` classifies `NotByok`. Both are `[model."claude-pro-max/x"]` override keeps `info.id` but can point `base_url`
`auth_method::platform_takes_session_credential(platform, base_url)`, the anywhere. `kimi-code` and a platform-less model (a bare slug / `[model.*]`
single predicate enforced at BOTH channels that reach the wire — the entry) ride the PRIMARY session, and ONLY at the session's own effective coding
`bearer_resolver` (outer term of `auth_method::session_token_auth_gate`) and endpoint: `EndpointsConfig::proxy_url()` (which prefers `[endpoints]
the `api_key` (`MvpAgent::session_token_for_model`, coding_api_base_url` from **config.toml** — what the managed-config sync writes
`oauth_registry::session_key_for_endpoint` for the aux / summary / — over `KIGI_CODE_BASE_URL`), `models_base_url`, loopback, or the compiled
subagent-override paths, and `sampler_turn::aux_bearer_resolver` for the production endpoint. Never a blanket allow: BYOK is `has_own_credentials()`,
stamped aux configs). which probes `std::env::var` at call time, so a `[model.*]` block with an unset
`env_key` classifies `NotByok`. Every API-key registry platform rides NOTHING.
STRUCTURAL ENFORCEMENT: the authority is the only producer of
`SessionCredential`, an opaque type with no production constructor, and every
API that stamps a session bearer onto a request (`resolve_credentials`,
`resolve_aux_model_sampling_config`, `try_resolve_model_credentials`,
`resolve_chat_state_auth_type`) takes
`Option<&SessionCredential>` rather than `Option<&str>` — a new call site
cannot express the leak. `stamp_session_local_sampler_fields` likewise takes the
aux `bearer_resolver` explicitly instead of copying the session's and relying on
the caller to re-point it, and `sampler_turn::aux_bearer_resolver_for` is the
ONE definition of the aux/summary resolver rule (the session actor and
`MvpAgent::build_summary_client` both call it; a private second copy is how the
summary client stayed ungated after M3). NEVER HAND-CARRY A CREDENTIAL TO A
GUARD (C1): "may **a** session credential ride here" is `true` for
a subscription-OAuth platform at its own host — where the credential that may
ride is that platform's POOLED token, never the primary. Ask
`credential_for(platform, base_url)` and stamp what it returns, so the question
and the credential are the same object; where you cannot (a credential the
authority does not own), MATCH on `credential_class` rather than reach for a
boolean. There is deliberately no second, similarly-named predicate to pick
wrongly. The shared `sampling_config.api_key`
(the subagent baseline and the unresolved-model fallback) has exactly TWO
production writers, both guarded by the authority:
`MvpAgent::stamp_session_credential` (the `cached_token` / `kimi.com/oidc`
login handlers and the `new_session` / `load_session` seed), which asks
`credential_for`; and the `xai.api_key` handler in `acp_agent.rs`, which stamps
the house `KIGI_API_KEY` read from the environment — a credential the authority
does not own — only when `credential_class` is `Primary`, the class the house
key rides and the one every OAuth platform's own host is NOT.
- MODEL→PLATFORM LOOKUP (security): `SamplingConfig::model` is the BARE routing - MODEL→PLATFORM LOOKUP (security): `SamplingConfig::model` is the BARE routing
slug, and duplicate slugs across platforms are BY DESIGN — an API-key platform slug, and duplicate slugs across platforms are BY DESIGN — an API-key platform
and its subscription-OAuth twin list identical ids (`xai`/`xai-grok`, and its subscription-OAuth twin list identical ids (`xai`/`xai-grok`,
`anthropic`/`claude-pro-max`, `openai`/`openai-codex`), with the API-key `anthropic`/`claude-pro-max`, `openai`/`openai-codex`), with the API-key
platform FIRST in `PlatformId::ALL`. The auth layer therefore resolves the platform FIRST in `PlatformId::ALL`. The auth layer therefore resolves the
platform from the catalog KEY the picker selected platform from the catalog KEY the picker selected, held PER SESSION in
(`ModelsManager::current_model_id`), via `SessionActor::selected_catalog_key` (seeded at spawn by
`agent::models::managed_key_for_slug`/`platform_for_slug`; anything else falls `agent::models::selected_catalog_key_for_spawn`, rewritten by `SetSessionModel`,
back to the picker's own `resolve_catalog_key`, and CLEARED by `OverrideModelName` when the rename makes it stale), via
`agent::models::entry_for_slug`/`platform_for_slug`.
NEVER `ModelsManager::current_model_id()`: that cell is process-global,
last-writer-wins across concurrent sessions, and Leader mode never writes it at
all (`agent/handlers/model_switch.rs`). The shared `MvpAgent::sampling_config`
is BUILT from that cell — but ONCE, at startup, and never rebuilt, while the
cell moves on every non-Leader switch. Its guards therefore read
`MvpAgent::sampling_config_platform`, the platform captured WITH the config by
the same `ModelsManager::sampling_config()` call, never a fresh lookup against
the live cell: once the two drift, re-resolving the config's bare slug falls
through to `resolve_catalog_key`'s `.rev()` scan, answers the API-key twin, and
a post-expiry `kigi login` silently leaves the EXPIRED bearer in the config
that seeds every subagent (H-a). REFUSE RATHER THAN GUESS (H-b):
when the per-session key does not name the slug, `platform_for_slug` returns
`None` for a slug that collides across platforms rather than trusting
`resolve_catalog_key`'s `.rev()` last match — which is the subscription-OAuth
twin, so the guess hands an API-key session the pooled bearer that REPLACES its
own key on the wire. `None` then routes purely by the ENDPOINT, which for an
OAuth host means no credential, no resolver and no adaptation.
Anything else (aux models, subagent
overrides) falls back to the picker's own `resolve_catalog_key`, and
`config::find_model_by_id`'s slug scan takes the LAST match so the two can `config::find_model_by_id`'s slug scan takes the LAST match so the two can
never disagree. Resolving the wrong twin costs the OAuth platform its live never disagree. Resolving the wrong twin costs the OAuth platform its live
`bearer_resolver` (unrecoverable 401 ~1h in), its Messages adaptation and its `bearer_resolver` (unrecoverable 401 ~1h in), its Messages adaptation and its
Copilot/Codex identity headers. Copilot/Codex identity headers — and hands the API-key twin's session a pooled
OAuth bearer stamped over the user's own key.
- CATALOG VISIBILITY: `platform_wire_model_to_entry` stamps - CATALOG VISIBILITY: `platform_wire_model_to_entry` stamps
`supported_in_api = platform != KimiCode`. `ModelInfo::visible_for_auth` `supported_in_api = platform != KimiCode`. `ModelInfo::visible_for_auth`
reads only the PRIMARY manager's auth mode, so gating the other OAuth reads only the PRIMARY manager's auth mode, so gating the other OAuth
+1 -1
View File
@@ -33,7 +33,7 @@ pub enum AuthScheme {
/// ///
/// `SamplerConfig` is the single source of truth for sampler /// `SamplerConfig` is the single source of truth for sampler
/// configuration. The shell builds it directly (see /// configuration. The shell builds it directly (see
/// `agent::config::resolve_model_to_sampling_config` and /// `agent::config::sampling_config_for_model` and
/// `session::acp_session::SessionActor::reconstruct_full_config`) by /// `session::acp_session::SessionActor::reconstruct_full_config`) by
/// composing chat-state's `kigi_sampling_types::SamplingConfig` /// composing chat-state's `kigi_sampling_types::SamplingConfig`
/// with `Credentials` (api key, client version). /// with `Credentials` (api key, client version).
+11 -1
View File
@@ -34,7 +34,17 @@ pub fn random_f64() -> f64 {
pub fn probabilistic_sample(rate: f64) -> bool { pub fn probabilistic_sample(rate: f64) -> bool {
random_f64() < rate random_f64() < rate
} }
fn matches_trusted_base_url(candidate: &str, trusted_base: &str) -> bool { /// True when `candidate` is `trusted_base` or a path below it, comparing
/// scheme, host and effective port exactly (so suffix attacks such as
/// `api.kimi.com.evil.example` never match).
///
/// Public because the credential chokepoint
/// ([`kigi_shell::auth::credential_authority`](../../../kigi_shell/auth/credential_authority/index.html))
/// must compare a request URL against the SESSION's effective endpoints
/// (`[endpoints] coding_api_base_url` from config.toml, `models_base_url`, a
/// platform's own registry host) — none of which the env-var-only predicates
/// below can see.
pub fn matches_trusted_base_url(candidate: &str, trusted_base: &str) -> bool {
let Ok(candidate) = reqwest::Url::parse(candidate) else { let Ok(candidate) = reqwest::Url::parse(candidate) else {
return false; return false;
}; };
@@ -1,6 +1,7 @@
use agent_client_protocol as acp; use agent_client_protocol as acp;
use crate::agent::config::ModelEntry; use crate::agent::config::ModelEntry;
use crate::auth::credential_authority::CredentialClass;
/// Shared, live handle to the agent's current ACP auth method id. /// Shared, live handle to the agent's current ACP auth method id.
/// ///
@@ -228,7 +229,7 @@ impl AuthMethodKind {
/// session, so the per-turn refresh / 401-recovery gate must be ACTIVE for /// session, so the per-turn refresh / 401-recovery gate must be ACTIVE for
/// it. This is correct ONLY because the session routes a model's /// it. This is correct ONLY because the session routes a model's
/// bearer/refresh/recovery to that model's OWN scope-keyed `AuthManager` /// bearer/refresh/recovery to that model's OWN scope-keyed `AuthManager`
/// (see `SessionActor::auth_manager_for_model`) — the gate being active /// (see `SessionActor::auth_manager_for_endpoint`) — the gate being active
/// wraps the grok manager for a grok turn, never the Kimi one. /// wraps the grok manager for a grok turn, never the Kimi one.
pub fn is_session_based(self) -> bool { pub fn is_session_based(self) -> bool {
matches!( matches!(
@@ -305,23 +306,30 @@ impl ModelByok {
/// third-party BYOK endpoint. A definite `NotByok` refreshes (within the /// third-party BYOK endpoint. A definite `NotByok` refreshes (within the
/// session's own providers); a definite `Byok` never does. /// session's own providers); a definite `Byok` never does.
/// ///
/// `endpoint_takes_session_credential` ([`platform_takes_session_credential`]) /// `credential_class` is the outer, non-negotiable guard and the reason the
/// is the outer, non-negotiable guard and the reason the `NotByok` arm is safe: /// `NotByok` arm is safe. It is NOT computed here: it comes from the single
/// it is `false` for every API-key registry platform (deepseek, openai, /// credential chokepoint,
/// anthropic, moonshot-*, …), whose models classify `NotByok` (they carry no /// [`crate::auth::credential_authority::CredentialAuthority::credential_class`],
/// `[model.*]` key) yet route to a THIRD-PARTY inference host while /// which is `None` for every API-key registry platform (deepseek, openai,
/// `oauth_registry::manager_for_model` falls through to the primary (Kimi) /// anthropic, moonshot-*, …) — whose models classify `NotByok` (they carry no
/// manager. Without this term the mainstream "session auth + API-key-platform /// `[model.*]` key) yet route to a THIRD-PARTY inference host — and for any
/// model" configuration stamps the user's Kimi subscription bearer on every /// `[model.*]` block pointed away from the session's own coding endpoint.
/// request to that host. /// Without this term the mainstream "session auth + API-key-platform model"
pub fn session_token_auth_gate( /// configuration stamps the user's Kimi subscription bearer on every request to
/// that host. `Pooled` and `Primary` both gate open: each names a credential
/// that *is* refreshable on that host, just not the same one.
pub(crate) fn session_token_auth_gate(
is_session_based_method: bool, is_session_based_method: bool,
model_byok: ModelByok, model_byok: ModelByok,
endpoint_is_first_party: bool, endpoint_is_first_party: bool,
endpoint_takes_session_credential: bool, credential_class: CredentialClass,
) -> bool { ) -> bool {
let endpoint_takes_a_session_credential = match credential_class {
CredentialClass::Pooled | CredentialClass::Primary => true,
CredentialClass::None => false,
};
is_session_based_method is_session_based_method
&& endpoint_takes_session_credential && endpoint_takes_a_session_credential
&& match model_byok { && match model_byok {
ModelByok::NotByok => true, ModelByok::NotByok => true,
ModelByok::Byok => false, ModelByok::Byok => false,
@@ -329,42 +337,6 @@ pub fn session_token_auth_gate(
} }
} }
/// Whether a session bearer may EVER be stamped on a request routed to
/// `platform` at `base_url` — i.e. whether the credential
/// `oauth_registry::manager_for_model` resolves for such a model belongs to the
/// host that receives it.
///
/// - a `uses_oauth` platform: `kimi-code` rides the PRIMARY session; each of the
/// four subscription-OAuth platforms (claude-pro-max, openai-codex,
/// github-copilot, xai-grok) rides its OWN pooled `AuthManager`. In both cases
/// the resolved bearer belongs to the host being called, and mid-session
/// refresh/401-recovery must stay live — so this is `true` even though those
/// four have non-first-party base URLs.
/// - every API-key registry platform: its credential is that platform's API key
/// (resolved into the catalog entry), never a session bearer, and
/// `manager_for_model` has no platform manager to route to. `false`.
/// - `None` — a bare slug or a `[model.*]` config entry, which carries no
/// platform at all — is decided by the ENDPOINT, never blanket-allowed: BYOK
/// is `has_own_credentials()`, which probes `std::env::var` at call time, so a
/// `[model.gpt-4o]` block with `base_url = "https://api.openai.com/v1"` and an
/// unset / mistyped `env_key` classifies `NotByok` and would otherwise hand
/// the Kimi subscription bearer to `api.openai.com`. Only the SESSION's own
/// coding endpoint qualifies: [`crate::util::is_effective_coding_endpoint_url`]
/// = the *effective* `KIGI_CODE_BASE_URL` deployment (so custom Kimi
/// deployments keep the session bearer) plus loopback (local dev proxies and
/// test mocks) plus the compiled production endpoint. Deliberately NOT
/// `is_first_party_url`, which is production-only and would break every
/// `KIGI_CODE_BASE_URL` deployment.
pub fn platform_takes_session_credential(
platform: Option<kigi_models::PlatformId>,
base_url: &str,
) -> bool {
match platform {
Some(platform) => platform.uses_oauth(),
None => crate::util::is_effective_coding_endpoint_url(base_url),
}
}
pub const AUTH_ERROR_SESSION_EXPIRED: &str = pub const AUTH_ERROR_SESSION_EXPIRED: &str =
"Session expired. Run `kigi login` to re-authenticate."; "Session expired. Run `kigi login` to re-authenticate.";
@@ -731,42 +703,47 @@ mod tests {
#[test] #[test]
fn session_token_auth_gate_matrix() { fn session_token_auth_gate_matrix() {
// Session method + NotByok → refresh (endpoint takes the session cred). // Both credential-bearing classes gate open on the same arms: `Pooled`
assert!(session_token_auth_gate( // is a subscription platform's own token on its own host, `Primary` the
true, // session's own bearer on the session's own endpoint.
ModelByok::NotByok, for class in [CredentialClass::Pooled, CredentialClass::Primary] {
false, // Session method + NotByok → refresh.
true assert!(session_token_auth_gate(
)); true,
// Session method + Byok → never. ModelByok::NotByok,
assert!(!session_token_auth_gate(true, ModelByok::Byok, true, true)); false,
// Session method + Unknown → only on first-party endpoints. class
assert!(session_token_auth_gate( ));
true, // Session method + Byok → never.
ModelByok::Unknown, assert!(!session_token_auth_gate(true, ModelByok::Byok, true, class));
true, // Session method + Unknown → only on first-party endpoints.
true assert!(session_token_auth_gate(
)); true,
assert!(!session_token_auth_gate( ModelByok::Unknown,
true, true,
ModelByok::Unknown, class
false, ));
true assert!(!session_token_auth_gate(
)); true,
// Non-session method → never. ModelByok::Unknown,
assert!(!session_token_auth_gate( false,
false, class
ModelByok::NotByok, ));
true, // Non-session method → never.
true assert!(!session_token_auth_gate(
)); false,
// An endpoint that does not take the session credential (every API-key ModelByok::NotByok,
// registry platform) is refused on EVERY arm — this is the outer guard true,
// that keeps the Kimi subscription bearer off third-party hosts. class
));
}
// An endpoint whose class is `None` (every API-key registry platform) is
// refused on EVERY arm — this is the outer guard that keeps the Kimi
// subscription bearer off third-party hosts.
for byok in [ModelByok::NotByok, ModelByok::Byok, ModelByok::Unknown] { for byok in [ModelByok::NotByok, ModelByok::Byok, ModelByok::Unknown] {
for first_party in [false, true] { for first_party in [false, true] {
assert!( assert!(
!session_token_auth_gate(true, byok, first_party, false), !session_token_auth_gate(true, byok, first_party, CredentialClass::None),
"byok={byok:?} first_party={first_party}: an API-key-platform \ "byok={byok:?} first_party={first_party}: an API-key-platform \
endpoint must never receive a session bearer" endpoint must never receive a session bearer"
); );
@@ -774,76 +751,6 @@ mod tests {
} }
} }
/// The platform → "may a session bearer ride here?" classification.
/// `kimi-code` rides the PRIMARY session; the four subscription-OAuth
/// platforms ride their OWN pooled managers (so they keep a live
/// bearer_resolver despite non-first-party base URLs); every API-key
/// registry platform is refused, whatever the endpoint.
#[test]
fn platform_takes_session_credential_matrix() {
let first_party = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
for id in [
"kimi-code",
"claude-pro-max",
"openai-codex",
"github-copilot",
"xai-grok",
] {
let platform = kigi_models::PlatformId::parse(id).expect("known platform");
for url in [first_party, "https://api.anthropic.com/v1"] {
assert!(
platform_takes_session_credential(Some(platform), url),
"{id} rides a session credential (primary or its own pool)"
);
}
}
for platform in kigi_models::PlatformId::ALL {
if platform.uses_oauth() {
continue;
}
for url in [first_party, "https://api.deepseek.com/v1"] {
assert!(
!platform_takes_session_credential(Some(platform), url),
"{} is an API-key platform — no session bearer may ride to it",
platform.as_str()
);
}
}
}
/// C2 regression: a platform-less model (a bare slug or a `[model.*]` entry)
/// is decided by the ENDPOINT, never blanket-allowed. A `[model.gpt-4o]`
/// block whose `env_key` is unset classifies `NotByok`, so before this the
/// `None` arm handed the Kimi subscription bearer to `api.openai.com`.
/// Custom `KIGI_CODE_BASE_URL` deployments and local dev proxies must still
/// keep it (that is why the predicate is not `is_first_party_url`).
#[test]
fn platform_less_model_takes_the_session_credential_only_on_its_own_endpoint() {
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!(
platform_takes_session_credential(None, url),
"{url} is the session's own endpoint (or a local proxy) — unchanged"
);
}
for url in [
"https://api.openai.com/v1",
"https://api.deepseek.com/v1",
"https://api.anthropic.com/v1",
"https://api.moonshot.cn/v1",
"",
] {
assert!(
!platform_takes_session_credential(None, url),
"LEAK: {url} is a third-party host — no session bearer may ride there"
);
}
}
/// RAII guard restoring an env var on drop (panic-safe). /// RAII guard restoring an env var on drop (panic-safe).
struct EnvGuard { struct EnvGuard {
key: &'static str, key: &'static str,
+69 -61
View File
@@ -1,5 +1,6 @@
use crate::agent::auth_method::ModelByok; use crate::agent::auth_method::ModelByok;
use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW; use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW;
use crate::auth::credential_authority::SessionCredential;
use crate::auth::{AuthManager, KimiCodeConfig}; use crate::auth::{AuthManager, KimiCodeConfig};
use crate::{config::StorageMode, sampling::ApiBackend, tools::config::ShellToolsetConfig}; use crate::{config::StorageMode, sampling::ApiBackend, tools::config::ShellToolsetConfig};
use agent_client_protocol as acp; use agent_client_protocol as acp;
@@ -3800,7 +3801,18 @@ pub(crate) fn first_own_credential(
/// Priority: model api_key/env_key > session token > XAI_API_KEY. /// Priority: model api_key/env_key > session token > XAI_API_KEY.
/// ///
/// When `env_key` lists multiple names, the first set non-empty value is used. /// When `env_key` lists multiple names, the first set non-empty value is used.
pub fn resolve_credentials(model: &ModelEntry, session_key: Option<&str>) -> ResolvedCredentials { ///
/// `session_key` is a [`SessionCredential`], NOT a bare string: the only way to
/// obtain one is
/// [`CredentialAuthority::credential_for`](crate::auth::credential_authority::CredentialAuthority::credential_for),
/// which requires the request's `(platform, base_url)` and the session's
/// effective endpoints. That is what makes the credential rule structural — a
/// new call site cannot stamp an unsanctioned bearer here even by omission,
/// because it cannot construct the value.
pub(crate) fn resolve_credentials(
model: &ModelEntry,
session_key: Option<&SessionCredential>,
) -> ResolvedCredentials {
let info = model.info(); let info = model.info();
let (api_key, base_url, auth_type) = if let Some(key) = model.own_credential() { let (api_key, base_url, auth_type) = if let Some(key) = model.own_credential() {
( (
@@ -3810,7 +3822,7 @@ pub fn resolve_credentials(model: &ModelEntry, session_key: Option<&str>) -> Res
) )
} else if let Some(key) = session_key { } else if let Some(key) = session_key {
( (
Some(key.to_owned()), Some(key.expose().to_owned()),
info.base_url.clone(), info.base_url.clone(),
kigi_chat_state::AuthType::SessionToken, kigi_chat_state::AuthType::SessionToken,
) )
@@ -3850,11 +3862,11 @@ pub fn resolve_credentials(model: &ModelEntry, session_key: Option<&str>) -> Res
/// Try to resolve credentials for a model by loading the effective config. /// Try to resolve credentials for a model by loading the effective config.
/// Returns `None` (with a warning) if config loading, parsing, or model /// Returns `None` (with a warning) if config loading, parsing, or model
/// lookup fails. `session_key` should only be passed when `auth_type` is /// lookup fails. `session_key` is a [`SessionCredential`], so only a bearer the
/// `SessionToken` — callers must guard this. /// credential chokepoint cleared for this endpoint can reach the wire.
pub fn try_resolve_model_credentials( pub(crate) fn try_resolve_model_credentials(
model_id: &str, model_id: &str,
session_key: Option<&str>, session_key: Option<&SessionCredential>,
) -> Option<ResolvedCredentials> { ) -> Option<ResolvedCredentials> {
let raw = crate::config::load_effective_config() let raw = crate::config::load_effective_config()
.map_err(|e| tracing::warn!(error = % e, "config load failed for credential resolution")) .map_err(|e| tracing::warn!(error = % e, "config load failed for credential resolution"))
@@ -3928,11 +3940,11 @@ fn with_resolved_model<T>(model_id: &str, f: impl FnOnce(ModelLookup) -> T) -> T
/// description, session summary, ...), resolved through the catalog so a /// description, session summary, ...), resolved through the catalog so a
/// `[model.*]` override redirects it to its own endpoint, credentials, and /// `[model.*]` override redirects it to its own endpoint, credentials, and
/// routing `model`. `None` → caller falls back to the active session's model. /// routing `model`. `None` → caller falls back to the active session's model.
pub fn resolve_aux_model_sampling_config( pub(crate) fn resolve_aux_model_sampling_config(
model_id: &str, model_id: &str,
models: &IndexMap<String, ModelEntry>, models: &IndexMap<String, ModelEntry>,
endpoints: &EndpointsConfig, endpoints: &EndpointsConfig,
session_key: Option<&str>, session_key: Option<&SessionCredential>,
alpha_test_key: Option<String>, alpha_test_key: Option<String>,
) -> Option<SamplerConfig> { ) -> Option<SamplerConfig> {
let catalog_entry = find_model_by_id(models, model_id).cloned(); let catalog_entry = find_model_by_id(models, model_id).cloned();
@@ -3944,7 +3956,7 @@ pub fn resolve_aux_model_sampling_config(
} }
} }
let xai_bearer = session_key let xai_bearer = session_key
.map(|s| s.to_owned()) .map(|s| s.expose().to_owned())
.or_else(|| crate::agent::auth_method::read_xai_api_key_env().ok()) .or_else(|| crate::agent::auth_method::read_xai_api_key_env().ok())
.or_else(|| endpoints.deployment_key.clone()); .or_else(|| endpoints.deployment_key.clone());
if let Some(bearer) = xai_bearer { if let Some(bearer) = xai_bearer {
@@ -4010,18 +4022,30 @@ pub fn resolve_aux_model_sampling_config(
/// from the active session onto a routed aux `SamplerConfig` so a /// from the active session onto a routed aux `SamplerConfig` so a
/// helper model keeps the session's auth/attribution. Shared by image-describe /// helper model keeps the session's auth/attribution. Shared by image-describe
/// and the auto-mode classifier so the two can't drift. /// and the auto-mode classifier so the two can't drift.
pub fn stamp_session_local_sampler_fields( /// Stamp the session-local fields onto a routed aux `SamplerConfig`.
///
/// `bearer_resolver` is an EXPLICIT parameter, not a copy of the session's:
/// this helper used to clone `active_session_config.bearer_resolver`
/// unconditionally and rely on every call site remembering to re-point it
/// afterwards. `SamplingClient::post` REPLACES the request's auth header from
/// the resolver, so a forgotten re-point overwrote the aux model's own key on
/// the AUX host with the session bearer. Callers obtain the value from
/// [`SessionActor::aux_bearer_resolver`](crate::session::acp_session::SessionActor)
/// — the chokepoint — so "forgot to re-point" is no longer expressible.
pub(crate) fn stamp_session_local_sampler_fields(
cfg: &mut SamplerConfig, cfg: &mut SamplerConfig,
active_session_config: &SamplerConfig, active_session_config: &SamplerConfig,
bearer_resolver: Option<kigi_sampler::SharedBearerResolver>,
max_retries: Option<u32>, max_retries: Option<u32>,
) { ) {
cfg.attribution_callback = active_session_config.attribution_callback.clone(); cfg.attribution_callback = active_session_config.attribution_callback.clone();
cfg.bearer_resolver = active_session_config.bearer_resolver.clone(); cfg.bearer_resolver = bearer_resolver;
cfg.max_retries = max_retries; cfg.max_retries = max_retries;
} }
pub fn finalize_image_describe_sampler_config( pub(crate) fn finalize_image_describe_sampler_config(
resolved_aux: Option<SamplerConfig>, resolved_aux: Option<SamplerConfig>,
active_session_config: &SamplerConfig, active_session_config: &SamplerConfig,
bearer_resolver: Option<kigi_sampler::SharedBearerResolver>,
max_retries: Option<u32>, max_retries: Option<u32>,
) -> (String, SamplerConfig) { ) -> (String, SamplerConfig) {
match resolved_aux { match resolved_aux {
@@ -4029,6 +4053,7 @@ pub fn finalize_image_describe_sampler_config(
stamp_session_local_sampler_fields( stamp_session_local_sampler_fields(
&mut describe_cfg, &mut describe_cfg,
active_session_config, active_session_config,
bearer_resolver,
max_retries, max_retries,
); );
let model = describe_cfg.model.clone(); let model = describe_cfg.model.clone();
@@ -4043,9 +4068,9 @@ pub fn finalize_image_describe_sampler_config(
/// Re-derive `auth_type` from the model's own credentials so BYOK env-key /// Re-derive `auth_type` from the model's own credentials so BYOK env-key
/// models stay on `ApiKey` even when a session token is present. Falls /// models stay on `ApiKey` even when a session token is present. Falls
/// back to `fallback` when the model isn't in the on-disk catalog. /// back to `fallback` when the model isn't in the on-disk catalog.
pub fn resolve_chat_state_auth_type( pub(crate) fn resolve_chat_state_auth_type(
model_id: &str, model_id: &str,
session_key: Option<&str>, session_key: Option<&SessionCredential>,
fallback: kigi_chat_state::AuthType, fallback: kigi_chat_state::AuthType,
) -> kigi_chat_state::AuthType { ) -> kigi_chat_state::AuthType {
try_resolve_model_credentials(model_id, session_key) try_resolve_model_credentials(model_id, session_key)
@@ -4185,23 +4210,6 @@ pub fn inject_url_derived_headers(
} }
let _ = (alpha_test_key, base_url); let _ = (alpha_test_key, base_url);
} }
pub fn resolve_model_to_sampling_config(
model_id: &str,
models: &IndexMap<String, ModelEntry>,
session_key: Option<&str>,
alpha_test_key: Option<String>,
fallback_entry: Option<ModelEntry>,
) -> Option<SamplerConfig> {
let entry = find_model_by_id(models, model_id)
.cloned()
.or(fallback_entry)?;
let credentials = resolve_credentials(&entry, session_key);
Some(sampling_config_for_model(
&entry,
credentials,
alpha_test_key,
))
}
pub fn to_acp_model_info( pub fn to_acp_model_info(
models: &IndexMap<String, ModelEntry>, models: &IndexMap<String, ModelEntry>,
) -> IndexMap<acp::ModelId, acp::ModelInfo> { ) -> IndexMap<acp::ModelId, acp::ModelInfo> {
@@ -4598,7 +4606,7 @@ reasoning_effort = "low"
model: "composer-session-model".into(), model: "composer-session-model".into(),
..Default::default() ..Default::default()
}; };
let (model, cfg) = finalize_image_describe_sampler_config(None, &active, Some(3)); let (model, cfg) = finalize_image_describe_sampler_config(None, &active, None, Some(3));
assert_eq!(model, "composer-session-model"); assert_eq!(model, "composer-session-model");
assert_eq!(cfg.model, "composer-session-model"); assert_eq!(cfg.model, "composer-session-model");
assert_ne!(cfg.model, "kigi"); assert_ne!(cfg.model, "kigi");
@@ -4613,7 +4621,8 @@ reasoning_effort = "low"
model: "kigi".into(), model: "kigi".into(),
..Default::default() ..Default::default()
}; };
let (model, cfg) = finalize_image_describe_sampler_config(Some(aux), &active, Some(7)); let (model, cfg) =
finalize_image_describe_sampler_config(Some(aux), &active, None, Some(7));
assert_eq!(model, "kigi"); assert_eq!(model, "kigi");
assert_eq!(cfg.model, "kigi"); assert_eq!(cfg.model, "kigi");
assert_eq!(cfg.max_retries, Some(7)); assert_eq!(cfg.max_retries, Some(7));
@@ -4661,9 +4670,9 @@ reasoning_effort = "low"
/// stamped as the `api_key` on an `api.moonshot.cn` request); /// stamped as the `api_key` on an `api.moonshot.cn` request);
/// - a `kimi-code` aux model → the primary, byte-identical. /// - a `kimi-code` aux model → the primary, byte-identical.
/// ///
/// Revert-to-red: dropping the `platform_takes_session_credential` term /// Revert-to-red: make `CredentialAuthority::governing_manager`'s
/// from `session_key_for_endpoint` makes the moonshot assertion see /// `Some(platform) => None` arm return `self.primary.clone()` and the
/// `Some("kimi-tok")`. /// moonshot assertion sees `Some("kimi-tok")`.
#[tokio::test] #[tokio::test]
async fn aux_model_session_key_is_platform_scoped_never_leaking_kimi() { async fn aux_model_session_key_is_platform_scoped_never_leaking_kimi() {
let (_kd, kimi) = kimi_primary("kimi-tok"); let (_kd, kimi) = kimi_primary("kimi-tok");
@@ -4673,16 +4682,16 @@ reasoning_effort = "low"
grok.info.id = Some("xai-grok/grok-4-latest".to_string()); grok.info.id = Some("xai-grok/grok-4-latest".to_string());
let mut grok_catalog = IndexMap::new(); let mut grok_catalog = IndexMap::new();
grok_catalog.insert("grok".to_string(), grok); grok_catalog.insert("grok".to_string(), grok);
let grok_key = crate::auth::oauth_registry::session_key_for_catalog_model( let authority = crate::auth::credential_authority::CredentialAuthority::new(
&grok_catalog, endpoints.clone(),
"grok", Some(kimi.clone()),
Some(&kimi),
); );
let grok_key = authority.credential_for_slug(&grok_catalog, None, "grok");
let grok_cfg = resolve_aux_model_sampling_config( let grok_cfg = resolve_aux_model_sampling_config(
"grok", "grok",
&grok_catalog, &grok_catalog,
&endpoints, &endpoints,
grok_key.as_deref(), grok_key.as_ref(),
None, None,
); );
assert_ne!( assert_ne!(
@@ -4703,11 +4712,9 @@ reasoning_effort = "low"
let mut k2_catalog = IndexMap::new(); let mut k2_catalog = IndexMap::new();
k2_catalog.insert("k2".to_string(), k2); k2_catalog.insert("k2".to_string(), k2);
assert_eq!( assert_eq!(
crate::auth::oauth_registry::session_key_for_catalog_model( authority
&k2_catalog, .credential_for_slug(&k2_catalog, None, "k2")
"k2", .map(|c| c.expose().to_owned()),
Some(&kimi),
),
None, None,
"LEAK: an API-key-platform aux model must receive no session token", "LEAK: an API-key-platform aux model must receive no session token",
); );
@@ -4722,16 +4729,12 @@ reasoning_effort = "low"
kimi_code.info.id = Some("kimi-code/kimi-for-coding".to_string()); kimi_code.info.id = Some("kimi-code/kimi-for-coding".to_string());
let mut kimi_catalog = IndexMap::new(); let mut kimi_catalog = IndexMap::new();
kimi_catalog.insert("kimi-for-coding".to_string(), kimi_code); kimi_catalog.insert("kimi-for-coding".to_string(), kimi_code);
let kimi_key = crate::auth::oauth_registry::session_key_for_catalog_model( let kimi_key = authority.credential_for_slug(&kimi_catalog, None, "kimi-for-coding");
&kimi_catalog,
"kimi-for-coding",
Some(&kimi),
);
let kimi_cfg = resolve_aux_model_sampling_config( let kimi_cfg = resolve_aux_model_sampling_config(
"kimi-for-coding", "kimi-for-coding",
&kimi_catalog, &kimi_catalog,
&endpoints, &endpoints,
kimi_key.as_deref(), kimi_key.as_ref(),
None, None,
) )
.expect("a kimi-code aux model resolves via the primary session token"); .expect("a kimi-code aux model resolves via the primary session token");
@@ -4870,7 +4873,8 @@ reasoning_effort = "low"
if entry.api_base_url.is_none() { if entry.api_base_url.is_none() {
continue; continue;
} }
let session_creds = resolve_credentials(&entry, Some("tok")); let session_creds =
resolve_credentials(&entry, Some(&SessionCredential::for_test("tok")));
assert_eq!( assert_eq!(
session_creds.base_url, session_creds.base_url,
endpoints.proxy_url(), endpoints.proxy_url(),
@@ -5028,7 +5032,7 @@ reasoning_effort = "low"
let mut model = test_model_entry("m", "https://inference.example/v1", None, None, None); let mut model = test_model_entry("m", "https://inference.example/v1", None, None, None);
model.env_key = Some(EnvKeys::new([primary, alias])); model.env_key = Some(EnvKeys::new([primary, alias]));
assert!(!model.has_own_credentials()); assert!(!model.has_own_credentials());
let creds = resolve_credentials(&model, Some("session-jwt")); let creds = resolve_credentials(&model, Some(&SessionCredential::for_test("session-jwt")));
assert_eq!(creds.auth_type, AuthType::SessionToken); assert_eq!(creds.auth_type, AuthType::SessionToken);
assert_eq!(creds.api_key.as_deref(), Some("session-jwt")); assert_eq!(creds.api_key.as_deref(), Some("session-jwt"));
} }
@@ -5062,7 +5066,7 @@ reasoning_effort = "low"
use kigi_chat_state::AuthType; use kigi_chat_state::AuthType;
let model = test_model_entry("m", "https://inference.example/v1", Some(""), None, None); let model = test_model_entry("m", "https://inference.example/v1", Some(""), None, None);
assert!(!model.has_own_credentials()); assert!(!model.has_own_credentials());
let creds = resolve_credentials(&model, Some("session-jwt")); let creds = resolve_credentials(&model, Some(&SessionCredential::for_test("session-jwt")));
assert_eq!(creds.auth_type, AuthType::SessionToken); assert_eq!(creds.auth_type, AuthType::SessionToken);
assert_eq!(creds.api_key.as_deref(), Some("session-jwt")); assert_eq!(creds.api_key.as_deref(), Some("session-jwt"));
} }
@@ -5091,10 +5095,10 @@ reasoning_effort = "low"
fn resolve_credentials_sets_auth_type() { fn resolve_credentials_sets_auth_type() {
use kigi_chat_state::AuthType; use kigi_chat_state::AuthType;
let model = test_model_entry("m", "https://example.com/v1", None, None, None); let model = test_model_entry("m", "https://example.com/v1", None, None, None);
let creds = resolve_credentials(&model, Some("tok")); let creds = resolve_credentials(&model, Some(&SessionCredential::for_test("tok")));
assert_eq!(creds.auth_type, AuthType::SessionToken); assert_eq!(creds.auth_type, AuthType::SessionToken);
let byok = test_model_entry("m", "https://example.com/v1", Some("key"), None, None); let byok = test_model_entry("m", "https://example.com/v1", Some("key"), None, None);
let creds = resolve_credentials(&byok, Some("tok")); let creds = resolve_credentials(&byok, Some(&SessionCredential::for_test("tok")));
assert_eq!(creds.auth_type, AuthType::ApiKey); assert_eq!(creds.auth_type, AuthType::ApiKey);
} }
/// Regression: BYOK env-var auth must stay ApiKey even when signed in, /// Regression: BYOK env-var auth must stay ApiKey even when signed in,
@@ -5115,7 +5119,7 @@ reasoning_effort = "low"
None, None,
); );
assert!(model.has_own_credentials()); assert!(model.has_own_credentials());
let creds = resolve_credentials(&model, Some("session-jwt")); let creds = resolve_credentials(&model, Some(&SessionCredential::for_test("session-jwt")));
assert_eq!( assert_eq!(
creds.auth_type, creds.auth_type,
AuthType::ApiKey, AuthType::ApiKey,
@@ -5140,8 +5144,11 @@ reasoning_effort = "low"
None, None,
); );
model.info.api_backend = ApiBackend::Messages; model.info.api_backend = ApiBackend::Messages;
let config = let config = sampling_config_for_model(
sampling_config_for_model(&model, resolve_credentials(&model, Some("tok")), None); &model,
resolve_credentials(&model, Some(&SessionCredential::for_test("tok"))),
None,
);
assert_eq!(config.api_backend, ApiBackend::Messages); assert_eq!(config.api_backend, ApiBackend::Messages);
assert_eq!(config.auth_scheme, AuthScheme::Bearer); assert_eq!(config.auth_scheme, AuthScheme::Bearer);
assert_eq!(config.api_key, Some("tok".to_string())); assert_eq!(config.api_key, Some("tok".to_string()));
@@ -6377,7 +6384,8 @@ reasoning_effort = "low"
(cfg, resolved) (cfg, resolved)
} }
fn resolve_sampling(model: &ModelEntry, session_key: Option<&str>) -> SamplerConfig { fn resolve_sampling(model: &ModelEntry, session_key: Option<&str>) -> SamplerConfig {
let credentials = resolve_credentials(model, session_key); let session_key = session_key.map(SessionCredential::for_test);
let credentials = resolve_credentials(model, session_key.as_ref());
sampling_config_for_model(model, credentials, None) sampling_config_for_model(model, credentials, None)
} }
#[test] #[test]
@@ -192,9 +192,18 @@ pub(crate) async fn apply(
model.map(|e| &e.info), model.map(|e| &e.info),
) )
}; };
// H4: hand the session the catalog KEY the picker actually resolved. The
// slug in `model_sampling.model` cannot distinguish `xai/grok-*` from
// `xai-grok/grok-*` (duplicate ids across an API-key platform and its
// subscription-OAuth twin are by design), and the process-global
// `current_model_id()` below is not written at all in Leader mode.
let catalog_key =
crate::agent::models::resolve_catalog_key(&agent.models_manager.models(), &model_id)
.map(|k| k.0.to_string());
let (tx, rx) = oneshot::channel(); let (tx, rx) = oneshot::channel();
let _ = handle.cmd_tx.send(SessionCommand::SetSessionModel { let _ = handle.cmd_tx.send(SessionCommand::SetSessionModel {
sampling_config: model_sampling, sampling_config: model_sampling,
catalog_key,
use_concise, use_concise,
apply_prompt_override, apply_prompt_override,
skip_prompt_rewrite: did_rebuild || model_unchanged, skip_prompt_rewrite: did_rebuild || model_unchanged,
+190 -45
View File
@@ -16,6 +16,19 @@ use crate::sampling::SamplerConfig as SamplingConfig;
use globset::{Glob, GlobSet, GlobSetBuilder}; use globset::{Glob, GlobSet, GlobSetBuilder};
use kigi_sampling_types::{ReasoningEffort, ReasoningEffortOption}; use kigi_sampling_types::{ReasoningEffort, ReasoningEffortOption};
/// The agent-wide baseline sampling config together with the registry platform
/// of the catalog entry it was BUILT from.
///
/// Returned as ONE value by [`ModelsManager::sampling_config`] so a holder
/// cannot end up with the config but not the platform its credential was
/// resolved against — the drift that made the login stamp guard answer about a
/// model the config no longer represents (H-a).
pub struct BaselineSamplingConfig {
pub config: SamplingConfig,
/// `None` for a bare / `[model.*]` entry.
pub platform: Option<kigi_models::PlatformId>,
}
// ── Auth method for model fetching ────────────────────────────────────────── // ── Auth method for model fetching ──────────────────────────────────────────
/// How the model catalog is fetched (PRD F4). The old xAI tier-gated proxy /// How the model catalog is fetched (PRD F4). The old xAI tier-gated proxy
@@ -1060,8 +1073,17 @@ impl ModelsManager {
.store(false, Ordering::Relaxed); .store(false, Ordering::Relaxed);
} }
/// Build a `SamplingConfig` from the current model + auth state. /// Build the agent-wide baseline `SamplingConfig` from the current model +
pub fn sampling_config(&self) -> SamplingConfig { /// auth state, together with the registry platform of the catalog entry it
/// was BUILT from.
///
/// H-a: the platform is returned WITH the config, never re-derived later.
/// The shared config is built exactly once (`MvpAgent::with_models`) and
/// never rebuilt, while [`Self::current_model_id`] moves on every non-Leader
/// model switch and on catalog reselection — so a guard that re-resolved
/// `config.model` against the LIVE cell answers about a DIFFERENT entry as
/// soon as the two drift.
pub fn sampling_config(&self) -> BaselineSamplingConfig {
let config = self.inner.cfg.read().clone(); let config = self.inner.cfg.read().clone();
let current_model_id = self.current_model_id(); let current_model_id = self.current_model_id();
let all_models = self.models(); let all_models = self.models();
@@ -1079,35 +1101,45 @@ impl ModelsManager {
} }
}; };
// Resolve the session bearer PER the current model's platform: a // H1: the session bearer comes from the ONE credential chokepoint,
// generic device-code OAuth platform (xai-grok) resolves from ITS OWN // resolved against the CURRENT MODEL's own platform AND endpoint. This
// scope-keyed store — never the Kimi session. Kimi and every other // used to be a local re-derivation that fell through to
// model keep the primary manager's token (path byte-identical). // `auth_manager.current_or_expired()` for every non-OAuth platform —
let session_key = self.session_key_for_catalog_key(current_model_id.0.as_ref()); // byte-for-byte the round-1 defect, reachable with ZERO configuration
let credentials = resolve_credentials(current_model, session_key.as_deref()); // (`default_models.json` bundles `moonshot-cn/*` entries a Kimi
// subscription user sees on first launch / offline, and this config is
sampling_config_for_model( // the `MvpAgent` baseline that seeds subagents and the
// unresolved-model fallback).
let credentials = resolve_credentials(
current_model, current_model,
credentials, self.credential_authority()
config.endpoints.alpha_test_key.clone(), .credential_for_model(current_model)
) .as_ref(),
);
BaselineSamplingConfig {
// The SAME entry the credential above was resolved against, so the
// config's guards ask about the model the config represents.
platform: crate::auth::credential_authority::entry_platform(current_model),
config: sampling_config_for_model(
current_model,
credentials,
config.endpoints.alpha_test_key.clone(),
),
}
} }
/// The session bearer for a model catalog key. Generic device-code OAuth /// The credential chokepoint for this manager: the session's EFFECTIVE
/// platforms (xai-grok) resolve from their own persisted scope; every other /// endpoints (so a managed `[endpoints] coding_api_base_url` deployment is
/// key resolves from the primary (Kimi) `AuthManager`. Reads the persisted /// recognised) plus the primary session manager, which the authority keeps
/// token (refresh happens in the async catalog/refresh paths); never logs /// private.
/// it. pub(crate) fn credential_authority(
fn session_key_for_catalog_key(&self, catalog_key: &str) -> Option<String> { &self,
if let Some((platform, _)) = kigi_models::parse_managed_model_key(catalog_key) ) -> crate::auth::credential_authority::CredentialAuthority {
&& let Some(oauth) = platform.oauth() crate::auth::credential_authority::CredentialAuthority::new(
{ self.inner.cfg.read().endpoints.clone(),
let kigi_home = crate::util::kigi_home::kigi_home(); Some(self.inner.auth_manager.clone()),
return AuthManager::new_oauth_provider(&kigi_home, oauth) )
.current_or_expired()
.map(|a| a.key);
}
self.inner.auth_manager.current_or_expired().map(|a| a.key)
} }
/// Disk-cache origin key for this manager's current endpoints/auth shape /// Disk-cache origin key for this manager's current endpoints/auth shape
@@ -1184,7 +1216,11 @@ impl ModelsManager {
// (refreshed on expiry) from its own scope — independent of the Kimi // (refreshed on expiry) from its own scope — independent of the Kimi
// session above. // session above.
let oauth_tokens = crate::agent::models_fetch::resolve_generic_oauth_tokens( let oauth_tokens = crate::agent::models_fetch::resolve_generic_oauth_tokens(
&crate::util::kigi_home::kigi_home(), // M7: the SAME home the OAuth pool resolves from. Under `cargo
// test` that is a disposable path, so a catalog fetch can never
// read the developer's real `~/.kigi` tokens — nor fire a real
// refresh request against them via `configure_refresher()`.
&crate::auth::oauth_registry::pool_home(),
) )
.await; .await;
let outcome = fetch_models_async( let outcome = fetch_models_async(
@@ -1212,7 +1248,11 @@ impl ModelsManager {
} }
let auth = self.inner.auth_manager.auth().await.ok(); let auth = self.inner.auth_manager.auth().await.ok();
let oauth_tokens = crate::agent::models_fetch::resolve_generic_oauth_tokens( let oauth_tokens = crate::agent::models_fetch::resolve_generic_oauth_tokens(
&crate::util::kigi_home::kigi_home(), // M7: the SAME home the OAuth pool resolves from. Under `cargo
// test` that is a disposable path, so a catalog fetch can never
// read the developer's real `~/.kigi` tokens — nor fire a real
// refresh request against them via `configure_refresher()`.
&crate::auth::oauth_registry::pool_home(),
) )
.await; .await;
let retry = let retry =
@@ -1917,7 +1957,11 @@ pub(crate) fn resolve_catalog_key(
.map(|(key, _)| acp::ModelId::new(key.clone())) .map(|(key, _)| acp::ModelId::new(key.clone()))
} }
/// The managed catalog key (`{platform}/{model}`) a routing slug belongs to. /// The catalog ENTRY a routing slug resolves to. The single lookup behind both
/// the managed key and the endpoint, so a caller can never take the platform
/// from one entry and the `base_url` from another (M6: the aux path used to
/// resolve the platform with `current_key` and the credential with a separate
/// `find_model_by_id`).
/// ///
/// H5: `SamplingConfig::model` is the BARE routing slug, never the catalog key, /// H5: `SamplingConfig::model` is the BARE routing slug, never the catalog key,
/// and duplicate slugs across platforms are BY DESIGN — the registry guarantees /// and duplicate slugs across platforms are BY DESIGN — the registry guarantees
@@ -1929,37 +1973,138 @@ pub(crate) fn resolve_catalog_key(
/// unrecoverable 401 ~1h in), its Messages adaptation and its Copilot/Codex /// unrecoverable 401 ~1h in), its Messages adaptation and its Copilot/Codex
/// identity headers. /// identity headers.
/// ///
/// `current_key` — [`ModelsManager::current_model_id`], the catalog key the /// `current_key` — the SESSION's own selected catalog key
/// picker actually selected — is therefore authoritative whenever it names this /// ([`crate::session::acp_session::SessionActor::selected_catalog_key`], seeded
/// slug. Anything else (aux models, subagent overrides, unlisted slugs) falls /// by [`selected_catalog_key_for_spawn`] and rewritten by `SetSessionModel`) —
/// back to the picker's OWN lookup, [`resolve_catalog_key`], so the auth layer /// is therefore authoritative whenever it names this slug. It is NOT
/// and the picker can never resolve different entries. /// [`ModelsManager::current_model_id`]: that cell is process-global,
pub(crate) fn managed_key_for_slug( /// last-writer-wins across concurrent sessions and never written at all in
models: &IndexMap<String, ModelEntry>, /// Leader mode (H4). Anything else (aux models, subagent overrides, unlisted
/// slugs) falls back to the picker's OWN lookup, [`resolve_catalog_key`], so the
/// auth layer and the picker can never resolve different entries.
///
/// The one caller that legitimately passes `current_model_id` is the SHARED
/// `MvpAgent::sampling_config`, which `ModelsManager::sampling_config` builds
/// from exactly that key — there the two lookups must agree (H-a).
///
/// (L: this rule used to be stated on a `managed_key_for_slug` wrapper that no
/// caller needed once [`platform_for_slug`] resolved the entry itself. The
/// crate-level `#![allow(dead_code)]` in `lib.rs` means an unused helper on a
/// credential path raises no warning, so dead ones are deleted on sight rather
/// than left as a second, unexercised way to answer the same question.)
pub(crate) fn entry_for_slug<'a>(
models: &'a IndexMap<String, ModelEntry>,
current_key: Option<&str>, current_key: Option<&str>,
slug: &str, slug: &str,
) -> Option<String> { ) -> Option<&'a ModelEntry> {
entry_for_slug_resolution(models, current_key, slug).map(|(entry, _)| entry)
}
/// [`entry_for_slug`] plus whether the SESSION's own selected catalog key is
/// what resolved it. `false` means the entry came from the picker's slug scan —
/// a resolution that is only a GUESS when the slug collides across platforms
/// (H-b).
fn entry_for_slug_resolution<'a>(
models: &'a IndexMap<String, ModelEntry>,
current_key: Option<&str>,
slug: &str,
) -> Option<(&'a ModelEntry, bool)> {
if let Some(entry) = current_key.and_then(|key| models.get(key)) if let Some(entry) = current_key.and_then(|key| models.get(key))
&& (entry.info.model == slug || current_key == Some(slug)) && (entry.info.model == slug || current_key == Some(slug))
{ {
return entry.info.id.clone(); return Some((entry, true));
} }
let key = resolve_catalog_key(models, &acp::ModelId::new(slug.to_string()))?; let key = resolve_catalog_key(models, &acp::ModelId::new(slug.to_string()))?;
models.get(key.0.as_ref())?.info.id.clone() models.get(key.0.as_ref()).map(|entry| (entry, false))
} }
/// The registry platform a routing slug resolves to, via [`managed_key_for_slug`]. /// Whether `slug` is carried as a routing slug by catalog entries belonging to
/// MORE THAN ONE platform — the dual-credential collision the registry creates
/// BY DESIGN (`anthropic` + `claude-pro-max` both list `claude-opus-4-8`;
/// likewise `xai`/`xai-grok`, `openai`/`openai-codex`).
///
/// An exact catalog-key match is never ambiguous: it names exactly one entry.
fn slug_collides_across_platforms(models: &IndexMap<String, ModelEntry>, slug: &str) -> bool {
if models.contains_key(slug) {
return false;
}
let mut seen: Option<Option<kigi_models::PlatformId>> = None;
for entry in models.values().filter(|entry| entry.info.model == slug) {
let platform = crate::auth::credential_authority::entry_platform(entry);
match seen {
None => seen = Some(platform),
Some(first) if first == platform => {}
Some(_) => return true,
}
}
false
}
/// The registry platform a routing slug resolves to, via [`entry_for_slug`].
/// `None` for a bare / `[model.*]` / unlisted model. Single definition shared by /// `None` for a bare / `[model.*]` / unlisted model. Single definition shared by
/// the session actor's inference-auth chokepoints and the aux/summary paths, so /// the session actor's inference-auth chokepoints and the aux/summary paths, so
/// the gate, the manager and the wire adaptations can never disagree. /// the gate, the manager and the wire adaptations can never disagree.
///
/// H-b — REFUSE RATHER THAN GUESS. When `current_key` does not name this slug
/// (`None`, or stale after an `OverrideModelName` rename) the resolution falls
/// through to [`resolve_catalog_key`]'s `.rev()` scan, whose LAST match is the
/// subscription-OAuth twin because `PlatformId::ALL` orders every API-key
/// platform first. Trusting that guess hands an API-KEY session the twin's
/// POOLED bearer — `SamplingClient::post` REPLACES the user's own `sk-ant-…` on
/// the wire — plus the OAuth Messages adaptation, which is precisely what H4
/// exists to prevent. A collided slug the session did not disambiguate
/// therefore resolves to NO platform, which the chokepoint then decides purely
/// by the ENDPOINT: the OAuth host is not the session's coding endpoint, so no
/// session credential, no resolver and no adaptation ride. First-party is
/// untouched — a platform-less model and a `kimi-code` model take the identical
/// endpoint arm.
pub(crate) fn platform_for_slug( pub(crate) fn platform_for_slug(
models: &IndexMap<String, ModelEntry>, models: &IndexMap<String, ModelEntry>,
current_key: Option<&str>, current_key: Option<&str>,
slug: &str, slug: &str,
) -> Option<kigi_models::PlatformId> { ) -> Option<kigi_models::PlatformId> {
let key = managed_key_for_slug(models, current_key, slug); let (key, disambiguated) = match entry_for_slug_resolution(models, current_key, slug) {
kigi_models::parse_managed_model_key(key.as_deref().unwrap_or(slug)) Some((entry, disambiguated)) => (entry.info.id.clone(), disambiguated),
.map(|(platform, _)| platform) None => (None, false),
};
let platform = kigi_models::parse_managed_model_key(key.as_deref().unwrap_or(slug))
.map(|(platform, _)| platform)?;
if platform.oauth().is_some() && !disambiguated && slug_collides_across_platforms(models, slug)
{
tracing::warn!(
slug,
platform = platform.as_str(),
"auth: routing slug collides across platforms and this session did not select \
one — refusing to guess the subscription-OAuth twin (no session credential)"
);
return None;
}
Some(platform)
}
/// The per-session selected catalog key (H4) a session spawned on
/// `session_model_id` must record: the picker's OWN lookup, never the
/// process-global [`ModelsManager::current_model_id`].
///
/// `session_model_id` is a catalog key on the `new_session` path (the picker's
/// `current_model_id()`), so this is idempotent there. On `load_session` it is
/// the RAW persisted `summary.current_model_id` (`acp_agent.rs`'s spawn call) —
/// `resolve_catalog_key` runs LATER, on the model-state / availability path, not
/// before the spawn — and after any `SetSessionModel` that persisted value is a
/// BARE routing slug, because `handle_set_session_model` persists
/// `sampling_config.model`. The slug branch below is therefore the live resume
/// path, not a theoretical one: it resolves through the picker's OWN lookup, and
/// for a slug that collides across platforms the `.rev()` last match is the
/// resume default (H-b then applies to anything that lookup could not pin).
///
/// Named — rather than inlined at the spawn site — so it is reachable from a
/// test: it is one of only two production writers of the field the whole
/// model→platform rule keys on, and a wrong seed is silent (H-c).
pub(crate) fn selected_catalog_key_for_spawn(
models: &IndexMap<String, ModelEntry>,
session_model_id: &acp::ModelId,
) -> Option<String> {
resolve_catalog_key(models, session_model_id).map(|key| key.0.to_string())
} }
/// Catalog key for a persisted session model id, restricted to **selectable** /// Catalog key for a persisted session model id, restricted to **selectable**
@@ -330,10 +330,28 @@ impl acp::Agent for MvpAgent {
); );
match arguments.method_id.0.as_ref() { match arguments.method_id.0.as_ref() {
auth_method::XAI_API_KEY_METHOD_ID => { auth_method::XAI_API_KEY_METHOD_ID => {
// C1: the SECOND writer of the shared `sampling_config.api_key`.
// The house `KIGI_API_KEY` is the user's own key for the
// session's own endpoint, so the stamp requires the authority to
// classify that endpoint `CredentialClass::Primary` — NOT merely
// "takes some session credential", which is also true on a
// subscription-OAuth platform's host, where this key has no
// business (the shared config is the subagent baseline and the
// unresolved-model fallback). The key is still persisted to
// auth.json either way; only the stamp is guarded.
let takes_house_key = self.shared_config_takes_house_key();
let mut sampling_config = self.sampling_config.borrow_mut(); let mut sampling_config = self.sampling_config.borrow_mut();
if sampling_config.api_key.is_none() { if sampling_config.api_key.is_none() {
if let Ok(api_key) = auth_method::read_xai_api_key_env() { if let Ok(api_key) = auth_method::read_xai_api_key_env() {
sampling_config.api_key = Some(api_key.clone()); if takes_house_key {
sampling_config.api_key = Some(api_key.clone());
} else {
tracing::debug!(
model = sampling_config.model.as_str(),
"auth: house api key withheld from the shared sampling config \
(its endpoint is not this session's own)"
);
}
if let Err(e) = crate::auth::store_api_key( if let Err(e) = crate::auth::store_api_key(
&crate::util::kigi_home::kigi_home(), &crate::util::kigi_home::kigi_home(),
&api_key, &api_key,
@@ -389,7 +407,7 @@ impl acp::Agent for MvpAgent {
), ),
), ),
); );
let Some(auth) = self.auth_manager.current() else { let Some(_auth) = self.auth_manager.current() else {
let message = if self.auth_manager.is_expired() { let message = if self.auth_manager.is_expired() {
"Session expired, re-authentication required" "Session expired, re-authentication required"
} else { } else {
@@ -408,9 +426,15 @@ impl acp::Agent for MvpAgent {
.await; .await;
}; };
self.emit_settings_update_notification(); self.emit_settings_update_notification();
{ // H2/C1: route the stamp through the ONE guard, which asks the
let mut sampling_config = self.sampling_config.borrow_mut(); // authority which credential governs the shared config rather
sampling_config.api_key = Some(auth.key); // than being handed this one. That config may already point at
// a third-party model or at another provider's subscription
// host (it is the subagent baseline and the unresolved-model
// fallback), and `auth.key` authorizes only the session's own
// coding endpoint. The manager already holds this token, so the
// authority reads it back where it belongs.
if self.stamp_session_credential(true) {
tracing::debug!( tracing::debug!(
"auth: cached_token handler set api_key (SessionToken)" "auth: cached_token handler set api_key (SessionToken)"
); );
@@ -484,9 +508,15 @@ impl acp::Agent for MvpAgent {
err.message = e.to_string(); err.message = e.to_string();
err err
})?; })?;
{ // C1: hot-swap FIRST, then let the authority read the fresh
let mut sampling_config = self.sampling_config.borrow_mut(); // token back where it belongs. Nothing hand-carries `auth.key`
sampling_config.api_key = Some(auth.key.clone()); // to the shared config any more — the stamp is whatever
// credential governs that config's own model + endpoint, which
// for a session whose current model is another provider's
// subscription model is that provider's pooled token, and for a
// third-party host is nothing at all.
self.auth_manager.hot_swap(auth.clone());
if self.stamp_session_credential(true) {
tracing::debug!( tracing::debug!(
"auth: kimi.com/oidc handler set api_key (SessionToken)" "auth: kimi.com/oidc handler set api_key (SessionToken)"
); );
@@ -496,7 +526,6 @@ impl acp::Agent for MvpAgent {
None, None,
); );
} }
self.auth_manager.hot_swap(auth.clone());
self.emit_settings_update_notification(); self.emit_settings_update_notification();
self.set_auth_method(arguments.method_id.clone()); self.set_auth_method(arguments.method_id.clone());
self.models_manager.on_auth_changed().await; self.models_manager.on_auth_changed().await;
@@ -33,33 +33,27 @@ impl MvpAgent {
// would otherwise stamp onto an api.x.ai / api.deepseek.com request). // would otherwise stamp onto an api.x.ai / api.deepseek.com request).
// The first-party subscription channel still gets the primary // The first-party subscription channel still gets the primary
// (byte-identical). // (byte-identical).
let session_key = crate::auth::oauth_registry::session_key_for_catalog_model( let authority = self.credential_authority();
&models, let session_key = authority.credential_for_slug(&models, None, &slug);
&slug,
Some(&self.auth_manager),
);
let endpoints = self.models_manager.endpoints(); let endpoints = self.models_manager.endpoints();
let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone(); let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone();
let config = match crate::agent::config::resolve_aux_model_sampling_config( let config = match crate::agent::config::resolve_aux_model_sampling_config(
&slug, &slug,
&models, &models,
&endpoints, &endpoints,
session_key.as_deref(), session_key.as_ref(),
alpha_test_key, alpha_test_key,
) { ) {
Some(mut cfg) => { Some(mut cfg) => {
cfg.attribution_callback = primary.attribution_callback.clone(); cfg.attribution_callback = primary.attribution_callback.clone();
// H4: the SESSION model's bearer_resolver must not ride to a // The SESSION model's bearer_resolver must not ride to a summary
// summary model on a different provider — // model on a different provider — `SamplingClient::post`
// `SamplingClient::post` REPLACES the request's auth header // REPLACES the request's auth header from it, overwriting the
// from it, overwriting the summary model's own resolved key on // summary model's own resolved key on ITS host. The chokepoint
// ITS host. Route through the one aux-resolver decision. // resolves the resolver from the SUMMARY model's platform +
// endpoint instead; the session's is never even read.
cfg.bearer_resolver = cfg.bearer_resolver =
crate::session::acp_session::sampler_turn::aux_bearer_resolver( self.summary_bearer_resolver(&models, &slug, &cfg.base_url);
primary.bearer_resolver.clone(),
crate::agent::models::platform_for_slug(&models, None, &slug),
&cfg.base_url,
);
cfg.max_retries = primary.max_retries; cfg.max_retries = primary.max_retries;
cfg cfg
} }
@@ -73,6 +67,36 @@ impl MvpAgent {
let client = OaiCompatClient::new(config).map_err(map_sampling_err_to_acp)?; let client = OaiCompatClient::new(config).map_err(map_sampling_err_to_acp)?;
Ok((client, model)) Ok((client, model))
} }
/// The `bearer_resolver` the SESSION-SUMMARY client may carry — the SAME
/// rule the session actor's aux path applies
/// ([`crate::session::acp_session::sampler_turn::aux_bearer_resolver_for`]),
/// not a second copy of it.
///
/// M3 completed: the aux path was gated on the session-token gate and this
/// one was not, so an api-key / house-key session whose
/// `[model.session-summary]` block carries its OWN `env_key` on the
/// session's own coding endpoint had that key REPLACED on the wire by the
/// primary bearer on every summary request. Named (rather than inlined
/// above) so the gate is reachable from a test — the ungated version stayed
/// green because the resolver is consumed by `OaiCompatClient::new`.
///
/// Aux slugs are not the session's selection, so `current_key = None`
/// (H-b then refuses a collided slug rather than guessing its OAuth twin).
pub(super) fn summary_bearer_resolver(
&self,
models: &indexmap::IndexMap<String, ModelEntry>,
slug: &str,
base_url: &str,
) -> Option<kigi_sampler::SharedBearerResolver> {
let auth_method = self.auth_method_id.load();
crate::session::acp_session::sampler_turn::aux_bearer_resolver_for(
&self.credential_authority(),
auth_method.as_deref(),
crate::agent::models::platform_for_slug(models, None, slug),
crate::agent::config::resolve_model_auth_facts(slug).byok,
base_url,
)
}
/// `true` for session-based ACP auth methods. /// `true` for session-based ACP auth methods.
fn is_session_based_auth(&self) -> bool { fn is_session_based_auth(&self) -> bool {
self.auth_method_id self.auth_method_id
@@ -559,7 +583,12 @@ impl MvpAgent {
reauth = auth_meta.reauth, reauth = auth_meta.reauth,
"auth: generic oauth device login", "auth: generic oauth device login",
); );
let kigi_home = crate::util::kigi_home::kigi_home(); // M4: the SAME home the pool reads from
// (`oauth_registry::pool_home()`), not `kigi_home()` directly —
// identical in production, but a login driven from a lib test would
// otherwise write into the developer's real `~/.kigi` while every
// inference-time lookup read the disposable test pool home.
let kigi_home = crate::auth::oauth_registry::pool_home();
let auth_manager = let auth_manager =
std::sync::Arc::new(crate::auth::AuthManager::new_oauth_provider(&kigi_home, oauth)); std::sync::Arc::new(crate::auth::AuthManager::new_oauth_provider(&kigi_home, oauth));
auth_manager.configure_refresher(); auth_manager.configure_refresher();
@@ -681,56 +710,44 @@ impl MvpAgent {
); );
Ok(entry.clone()) Ok(entry.clone())
} }
/// Resolve the SESSION token for `model` by the model's OWN platform AND /// This agent's credential chokepoint: the EFFECTIVE endpoints (so a
/// endpoint — the single guard against the api_key-channel token leak. /// managed `[endpoints] coding_api_base_url` deployment keeps the session
/// bearer — H3) plus the primary manager, which the authority keeps private.
pub(crate) fn credential_authority(
&self,
) -> crate::auth::credential_authority::CredentialAuthority {
crate::auth::credential_authority::CredentialAuthority::new(
self.models_manager.endpoints(),
Some(self.auth_manager.clone()),
)
}
/// The SESSION credential for `model`, resolved by the model's OWN platform
/// AND endpoint at the chokepoint — the single guard against the
/// api_key-channel token leak (C1).
/// ///
/// - an oauth-platform model (xai-grok, claude-pro-max, github-copilot, /// A subscription-OAuth model draws from ITS OWN pooled manager (`None`,
/// openai-codex) draws its session token from ITS OWN process-global pool /// never the Kimi key, when that provider has no stored session). Every
/// manager (built on demand from the on-disk token, proactively /// API-key registry platform and every `[model.*]` block pointed at a
/// refreshed), INDEPENDENT of the primary `auth_method`; when that /// third-party host gets `None`: `resolve_credentials` would otherwise take
/// provider has no stored session the token is `None` — NEVER the primary /// the `else if let Some(key) = session_key` arm and set
/// Kimi key; /// `api_key = <Kimi bearer>` with the THIRD-PARTY `base_url`, which
/// - an endpoint that does not take a session credential at all /// `SamplingClient` builds into `Authorization: Bearer …` — reachable with
/// ([`crate::agent::auth_method::platform_takes_session_credential`] — /// ZERO configuration, since `default_models.json` bundles `moonshot-cn/*`
/// every API-key registry platform, and any `[model.*]` block pointed at a /// and `moonshot-ai/*` entries a Kimi-subscription user sees on first
/// third-party host) gets `None`. This is C1: `resolve_credentials` takes /// launch / offline. The first-party subscription channel uses the primary,
/// the `else if let Some(key) = session_key` arm and sets /// and only under a session-based auth method — byte-identical.
/// `api_key = <Kimi bearer>` with the THIRD-PARTY `base_url`, which
/// `SamplingClient` then builds into `Authorization: Bearer …`. It is
/// reachable with ZERO configuration: `default_models.json` bundles
/// `moonshot-cn/*` + `moonshot-ai/*` entries that a Kimi-subscription user
/// sees on first launch / offline;
/// - the first-party subscription channel (kimi-code, a `KIGI_CODE_BASE_URL`
/// deployment, a loopback proxy) uses the primary session manager, and
/// only under a session-based auth method — byte-identical to the pre-fix
/// path.
/// ///
/// SECURITY: the resolved token is never logged. /// SECURITY: the resolved token is never logged.
fn session_token_for_model(&self, model: &ModelEntry) -> Option<crate::auth::KimiAuth> { fn session_token_for_model(
let info = model.info(); &self,
let platform = info model: &ModelEntry,
.id ) -> Option<crate::auth::credential_authority::SessionCredential> {
.as_deref() let is_primary_channel = crate::auth::credential_authority::entry_platform(model)
.and_then(kigi_models::parse_managed_model_key) .is_none_or(|platform| platform.oauth().is_none());
.map(|(platform, _)| platform); if is_primary_channel && !self.is_session_based_auth() {
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
return crate::auth::oauth_registry::global_manager_for(
&crate::auth::oauth_registry::pool_home(),
oauth,
)
.current_or_expired();
}
if !crate::agent::auth_method::platform_takes_session_credential(
platform,
&info.base_url,
) {
return None; return None;
} }
if self.is_session_based_auth() { self.credential_authority().credential_for_model(model)
self.auth_manager.current_or_expired()
} else {
None
}
} }
pub(crate) fn prepare_sampling_config_for_model( pub(crate) fn prepare_sampling_config_for_model(
&self, &self,
@@ -746,10 +763,7 @@ impl MvpAgent {
// api_key is its own grok token or `None`, never the primary Kimi key. // api_key is its own grok token or `None`, never the primary Kimi key.
let session = self.session_token_for_model(model); let session = self.session_token_for_model(model);
let has_session_key = session.is_some(); let has_session_key = session.is_some();
let mut credentials = resolve_credentials( let mut credentials = resolve_credentials(model, session.as_ref());
model,
session.as_ref().map(|a| a.key.as_str()),
);
if !has_session_key && credentials.auth_type == kigi_chat_state::AuthType::ApiKey if !has_session_key && credentials.auth_type == kigi_chat_state::AuthType::ApiKey
&& !model.has_own_credentials() && self.is_session_based_auth() && !model.has_own_credentials() && self.is_session_based_auth()
{ {
@@ -936,7 +950,8 @@ impl MvpAgent {
models_manager: crate::agent::models::ModelsManager, models_manager: crate::agent::models::ModelsManager,
) -> Self { ) -> Self {
models_manager.set_gateway(gateway.clone()); models_manager.set_gateway(gateway.clone());
let sampling_config = models_manager.sampling_config(); // H-a: the config AND the platform it was built from, from ONE call.
let baseline = models_manager.sampling_config();
let storage_mode = cfg.storage_mode; let storage_mode = cfg.storage_mode;
let default_yolo_mode = cfg.default_yolo_mode; let default_yolo_mode = cfg.default_yolo_mode;
let default_auto_mode = cfg.default_auto_mode; let default_auto_mode = cfg.default_auto_mode;
@@ -991,7 +1006,8 @@ impl MvpAgent {
models_manager, models_manager,
cfg: RefCell::new(cfg.clone()), cfg: RefCell::new(cfg.clone()),
auth_method_id: crate::agent::auth_method::new_shared_auth_method_id(None), auth_method_id: crate::agent::auth_method::new_shared_auth_method_id(None),
sampling_config: RefCell::new(sampling_config), sampling_config: RefCell::new(baseline.config),
sampling_config_platform: std::cell::Cell::new(baseline.platform),
auth_manager, auth_manager,
auth_code_tx: RefCell::new(None), auth_code_tx: RefCell::new(None),
auth_url_rx: RefCell::new(None), auth_url_rx: RefCell::new(None),
@@ -1581,38 +1597,133 @@ impl MvpAgent {
); );
(serde_json::json!({ "options" : config_options }), serde_json::json!(detail)) (serde_json::json!({ "options" : config_options }), serde_json::json!(detail))
} }
/// The registry platform the SHARED `sampling_config` routes to: the one
/// captured WITH the config when it was built, never a fresh lookup.
///
/// H-a: `ModelsManager::sampling_config` builds the shared config from the
/// catalog entry `current_model_id()` named AT THAT MOMENT, and the config
/// is never rebuilt. A guard that re-resolves `config.model` (the BARE
/// routing slug) against the LIVE cell therefore answers about a DIFFERENT
/// entry the moment the two drift — and they drift on every non-Leader model
/// switch (`handlers/model_switch.rs`) and on catalog reselection.
/// `entry_for_slug_resolution`'s `entry.info.model == slug` test then fails
/// and the resolution falls through to `resolve_catalog_key`'s `.rev()`
/// scan; `PlatformId::ALL` lists `kimi-code` first and its API-key twin
/// `kimi-coding` 19th, so the LAST match is the twin, which takes no session
/// credential. For a Kimi-subscription user who also has `KIMI_API_KEY` set,
/// a post-expiry `kigi login` then found no governing manager and — because
/// the stamp only overwrites on success — left the EXPIRED bearer in place:
/// every unresolved-model fallback and every subagent baseline turn 401'd
/// until restart.
fn shared_config_platform(&self) -> Option<kigi_models::PlatformId> {
self.sampling_config_platform.get()
}
/// H2/C1: stamp the shared `sampling_config`'s OWN governing session
/// credential — whatever the authority says that is for this config's model
/// and endpoint.
///
/// The shared config is not inert: `resolve_sampling_config_for_model`
/// returns it verbatim whenever a model id fails to resolve, and
/// `SubagentSpawnContext` clones it as every subagent's baseline — so an
/// `api_key` stamped here reaches the wire against whatever `base_url` the
/// config carries. The login/seed sites used to stamp it unconditionally,
/// exactly the mistake `authenticate_oauth_platform` already documents ("Do
/// NOT stamp this token onto the shared sampling_config").
///
/// C1: this function does NOT take a credential. The previous shape took
/// `key: String` — always the primary Kimi bearer — and guarded it with a
/// predicate asking whether **a** session credential may ride. For a
/// subscription-OAuth platform at its own host that is correctly `true`, but
/// the credential that may ride there is that platform's POOLED token: a
/// Claude Pro/Max user running `kigi login` stamped the Kimi subscription
/// bearer onto a config routed at `api.anthropic.com`. Asking
/// `credential_for` instead makes the question and the credential the same
/// object, so the pairing cannot be wrong — and there is no primary handle
/// here to hand-carry (M2).
///
/// `overwrite = false` keeps the historical "only if missing" seeding
/// behaviour; the login handlers pass `true` because a fresh login must
/// replace a stale bearer, and they call this AFTER the manager holds the
/// new token so it is read back through the authority.
///
/// SECURITY: the token is never logged.
pub(super) fn stamp_session_credential(&self, overwrite: bool) -> bool {
let (model, base_url) = {
let sampling_config = self.sampling_config.borrow();
if !overwrite && sampling_config.api_key.is_some() {
return false;
}
(
sampling_config.model.clone(),
sampling_config.base_url.clone(),
)
};
let platform = self.shared_config_platform();
let Some(credential) = self.credential_authority().credential_for(platform, &base_url)
else {
tracing::debug!(
model = model.as_str(),
"auth: no session credential governs the shared sampling config \
(its model + endpoint take none, or the governing provider has no session)"
);
return false;
};
self.sampling_config.borrow_mut().api_key = Some(credential.expose().to_owned());
true
}
/// Whether the shared `sampling_config` may receive a credential that
/// authorizes only the SESSION's own first-party endpoint — the house
/// `KIGI_API_KEY` the `xai.api_key` login handler reads from the
/// environment, which the authority does not own and so cannot produce.
///
/// The house key rides the [`CredentialClass::Primary`] channel and no
/// other. `Pooled` is deliberately excluded: a subscription-OAuth platform's
/// own host DOES take a session credential, but that credential is the
/// platform's pooled token, where the house key has no business (C1).
pub(super) fn shared_config_takes_house_key(&self) -> bool {
let base_url = self.sampling_config.borrow().base_url.clone();
matches!(
self.credential_authority()
.credential_class(self.shared_config_platform(), &base_url),
crate::auth::credential_authority::CredentialClass::Primary
)
}
/// Seed the global sampling config with login auth when available. /// Seed the global sampling config with login auth when available.
/// ///
/// Only sets the `api_key` if missing. Does NOT resolve `base_url` from /// Only sets the `api_key` if missing, and only with the credential that
/// governs the config's own model + endpoint (see
/// [`Self::stamp_session_credential`]). Does NOT resolve `base_url` from
/// `current_model_id` — that's deferred to session creation time to avoid /// `current_model_id` — that's deferred to session creation time to avoid
/// cross-client contamination in leader mode (where `current_model_id` is /// cross-client contamination in leader mode (where `current_model_id` is
/// shared mutable state). /// shared mutable state).
pub(super) fn seed_client_config_auth_if_available(&self) { pub(super) fn seed_client_config_auth_if_available(&self) {
let mut sampling_config = self.sampling_config.borrow_mut(); if self.sampling_config.borrow().api_key.is_some() {
if sampling_config.api_key.is_none() { return;
if let Some(auth) = self.auth_manager.current_or_expired() { }
sampling_config.api_key = Some(auth.key); if self.stamp_session_credential(false) {
tracing::debug!("auth: seed_client_config set auth (SessionToken)"); tracing::debug!("auth: seed_client_config set auth (SessionToken)");
kigi_log::unified_log::debug( kigi_log::unified_log::debug(
"auth: seed_client_config set auth (SessionToken)", "auth: seed_client_config set auth (SessionToken)",
None, None,
None, None,
); );
} else if !self return;
}
// No credential was stamped. Only the total-absence case is worth
// warning about: a withheld-by-endpoint seed is the rule working.
if self.auth_manager.current_or_expired().is_none()
&& !self
.models_manager .models_manager
.models() .models()
.values() .values()
.any(|m| m.has_own_credentials()) .any(|m| m.has_own_credentials())
{ {
tracing::warn!( tracing::warn!("No credentials found: no login token and no model api_key/env_key");
"No credentials found: no login token and no model api_key/env_key" kigi_log::unified_log::warn(
); "No credentials found: no login token and no model api_key/env_key",
kigi_log::unified_log::warn( None,
"No credentials found: no login token and no model api_key/env_key", None,
None, );
None,
);
}
} }
} }
/// Allocate the next monotonic telemetry turn number for a session. /// Allocate the next monotonic telemetry turn number for a session.
@@ -2334,12 +2445,34 @@ impl MvpAgent {
} }
let (mut handle, agent_system_prompt, session_thread) = { let (mut handle, agent_system_prompt, session_thread) = {
let _timer = crate::instrumentation_timer!("session.spawn_actor_call"); let _timer = crate::instrumentation_timer!("session.spawn_actor_call");
let session_key = self.auth_manager.current_or_expired().map(|a| a.key); // Classify the credential's auth_type against the bearer this
// model's endpoint would ACTUALLY receive, not the raw primary:
// an API-key-platform model has no session credential at all, and
// reading the primary here reported one.
//
// H-a: the platform is resolved with the SAME catalog key the
// actor about to be spawned seeds itself with
// (`selected_catalog_key_for_spawn`, spawn.rs), so this
// classification and every per-turn decision that session makes
// agree by construction instead of re-resolving the bare slug.
let models_for_key = self.models_manager.models();
let session_key = self.credential_authority().credential_for(
crate::agent::models::platform_for_slug(
&models_for_key,
crate::agent::models::selected_catalog_key_for_spawn(
&models_for_key,
&session_model_id,
)
.as_deref(),
sampling_config.model.as_str(),
),
&sampling_config.base_url,
);
let credentials = kigi_chat_state::Credentials { let credentials = kigi_chat_state::Credentials {
api_key: sampling_config.api_key.clone(), api_key: sampling_config.api_key.clone(),
auth_type: crate::agent::config::resolve_chat_state_auth_type( auth_type: crate::agent::config::resolve_chat_state_auth_type(
sampling_config.model.as_str(), sampling_config.model.as_str(),
session_key.as_deref(), session_key.as_ref(),
self.auth_type(), self.auth_type(),
), ),
alpha_test_key: self.alpha_test_key(), alpha_test_key: self.alpha_test_key(),
@@ -501,6 +501,19 @@ pub struct MvpAgent {
/// only api_key is written here (same for all clients). Per-session base_url /// only api_key is written here (same for all clients). Per-session base_url
/// is resolved at session creation time in `new_session` / `load_session`. /// is resolved at session creation time in `new_session` / `load_session`.
pub(crate) sampling_config: RefCell<SamplingConfig>, pub(crate) sampling_config: RefCell<SamplingConfig>,
/// The registry platform [`Self::sampling_config`] was BUILT from, captured
/// in the same `ModelsManager::sampling_config()` call that produced it.
///
/// H-a: the shared config is built ONCE (`Self::with_models`) and never
/// rebuilt, but `ModelsManager::current_model_id()` moves on every
/// non-Leader model switch and on catalog reselection. Its guards
/// (`stamp_session_credential`, `shared_config_takes_house_key`)
/// therefore read THIS cell, never the live one: after a switch the live
/// cell names a different entry, and re-resolving the config's bare slug
/// against it fell through to `resolve_catalog_key`'s `.rev()` scan —
/// answering the API-key twin, which takes no session credential, so a
/// successful `kigi login` silently failed to replace the expired bearer.
pub(crate) sampling_config_platform: std::cell::Cell<Option<kigi_models::PlatformId>>,
pub(crate) auth_manager: Arc<AuthManager>, pub(crate) auth_manager: Arc<AuthManager>,
pub(crate) models_manager: crate::agent::models::ModelsManager, pub(crate) models_manager: crate::agent::models::ModelsManager,
/// Forwards pasted codes from `handle_auth_submit_code` to the auth flow. /// Forwards pasted codes from `handle_auth_submit_code` to the auth flow.
@@ -1291,6 +1291,7 @@ mod subagent_spawn_context_tests;
/// LEAK guard for the `api_key` channel (C1/C2), through the real /// LEAK guard for the `api_key` channel (C1/C2), through the real
/// `prepare_sampling_config_for_model` resolution path. /// `prepare_sampling_config_for_model` resolution path.
mod api_key_channel_leak_tests; mod api_key_channel_leak_tests;
mod chokepoint_leak_tests;
/// No load in flight and no session → the wait returns immediately /// No load in flight and no session → the wait returns immediately
/// (the caller then surfaces "unknown session id" exactly as before). /// (the caller then surfaces "unknown session id" exactly as before).
#[tokio::test] #[tokio::test]
@@ -24,12 +24,12 @@ use crate::agent::config::{Config as AgentConfig, EndpointsConfig, EnvKeys, Mode
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig}; use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
use kigi_test_support::EnvGuard; use kigi_test_support::EnvGuard;
const KIMI_TOKEN: &str = "kimi-subscription-token-DO-NOT-LEAK"; pub(super) const KIMI_TOKEN: &str = "kimi-subscription-token-DO-NOT-LEAK";
/// Ambient BYOK env vars unset, so a model with no resolvable credential ends up /// Ambient BYOK env vars unset, so a model with no resolvable credential ends up
/// with `api_key == None` rather than a global-key fallback that could mask the /// with `api_key == None` rather than a global-key fallback that could mask the
/// leak under test. Every test holding these must be `#[serial]`. /// leak under test. Every test holding these must be `#[serial]`.
fn without_ambient_byok_env() -> [EnvGuard; 3] { pub(super) fn without_ambient_byok_env() -> [EnvGuard; 3] {
[ [
EnvGuard::unset(HOUSE_API_KEY_ENV_VAR), EnvGuard::unset(HOUSE_API_KEY_ENV_VAR),
EnvGuard::unset(XAI_API_KEY_ENV_VAR), EnvGuard::unset(XAI_API_KEY_ENV_VAR),
@@ -41,7 +41,7 @@ fn without_ambient_byok_env() -> [EnvGuard; 3] {
/// Kimi subscription bearer — the mainstream configuration in which the leak /// Kimi subscription bearer — the mainstream configuration in which the leak
/// fires. `(tempdir, agent)`; the tempdir is the auth store and is returned so /// fires. `(tempdir, agent)`; the tempdir is the auth store and is returned so
/// the caller keeps it alive. /// the caller keeps it alive.
fn kimi_session_agent() -> (tempfile::TempDir, MvpAgent) { pub(super) fn kimi_session_agent() -> (tempfile::TempDir, MvpAgent) {
let dir = tempfile::tempdir().expect("tempdir"); let dir = tempfile::tempdir().expect("tempdir");
let auth_manager = std::sync::Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); let auth_manager = std::sync::Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
auth_manager.hot_swap(KimiAuth { auth_manager.hot_swap(KimiAuth {
@@ -65,7 +65,7 @@ fn kimi_session_agent() -> (tempfile::TempDir, MvpAgent) {
/// A catalog entry as `resolve_model_list` builds one for a fetched registry /// A catalog entry as `resolve_model_list` builds one for a fetched registry
/// model: managed catalog key, platform base URL, no credential of its own. /// model: managed catalog key, platform base URL, no credential of its own.
fn platform_entry(catalog_key: &str, slug: &str, base_url: &str) -> ModelEntry { pub(super) fn platform_entry(catalog_key: &str, slug: &str, base_url: &str) -> ModelEntry {
let mut entry = ModelEntry::fallback(slug, &EndpointsConfig::default()); let mut entry = ModelEntry::fallback(slug, &EndpointsConfig::default());
entry.info.id = Some(catalog_key.to_string()); entry.info.id = Some(catalog_key.to_string());
entry.info.base_url = base_url.to_string(); entry.info.base_url = base_url.to_string();
@@ -80,8 +80,9 @@ fn platform_entry(catalog_key: &str, slug: &str, base_url: &str) -> ModelEntry {
/// `Authorization: Bearer <Kimi OAuth token>` to `api.moonshot.cn`, which is NOT /// `Authorization: Bearer <Kimi OAuth token>` to `api.moonshot.cn`, which is NOT
/// first-party. /// first-party.
/// ///
/// Revert-to-red: dropping the `platform_takes_session_credential` guard from /// Revert-to-red: make `CredentialAuthority::governing_manager`'s
/// `session_token_for_model` makes every `api_key` below `Some(KIMI_TOKEN)`. /// `Some(platform) => None` arm return `self.primary.clone()` and every `api_key`
/// below becomes `Some(KIMI_TOKEN)`.
#[tokio::test] #[tokio::test]
#[serial_test::serial] #[serial_test::serial]
async fn bundled_default_moonshot_models_never_carry_the_kimi_bearer() { async fn bundled_default_moonshot_models_never_carry_the_kimi_bearer() {
@@ -154,8 +155,8 @@ async fn api_key_platform_models_never_carry_the_kimi_bearer_as_api_key() {
/// unset (or mistyped) `env_key` classifies the model NotByok and the Kimi /// unset (or mistyped) `env_key` classifies the model NotByok and the Kimi
/// bearer went to `api.openai.com` on BOTH channels. /// bearer went to `api.openai.com` on BOTH channels.
/// ///
/// Revert-to-red: making the `None` arm of `platform_takes_session_credential` /// Revert-to-red: make `CredentialAuthority::is_session_coding_endpoint` return
/// return `true` again makes `api_key` here `Some(KIMI_TOKEN)`. /// `true` unconditionally and `api_key` here becomes `Some(KIMI_TOKEN)`.
#[tokio::test] #[tokio::test]
#[serial_test::serial] #[serial_test::serial]
async fn config_model_with_an_unset_env_key_never_carries_the_kimi_bearer() { async fn config_model_with_an_unset_env_key_never_carries_the_kimi_bearer() {
@@ -0,0 +1,589 @@
//! LEAK GUARD (the chokepoint's own call sites) — H1, H2 and the H3 regression.
//!
//! The companion of `api_key_channel_leak_tests`, which drives
//! `MvpAgent::prepare_sampling_config_for_model`. These three pin the OTHER
//! producers of an outgoing credential: `ModelsManager::sampling_config()` (the
//! agent-wide baseline), the shared `MvpAgent::sampling_config` the login/seed
//! paths stamp, and the SESSION's effective coding endpoint as configured from
//! config.toml rather than the environment.
//!
//! Every assertion reads what the real resolution path produced; nothing is
//! hand-stamped.
use super::super::*;
use super::api_key_channel_leak_tests::{
KIMI_TOKEN, kimi_session_agent, platform_entry, without_ambient_byok_env,
};
use crate::agent::auth_method::CACHED_TOKEN_AUTH_METHOD_ID;
use crate::agent::config::{Config as AgentConfig, EndpointsConfig, ModelEntry};
use crate::auth::credential_authority::CredentialClass;
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
use kigi_sampler::BearerResolver;
use kigi_test_support::EnvGuard;
/// Re-seed the shared config the way `MvpAgent::with_models` does: the config
/// AND the platform it was BUILT from, from ONE `ModelsManager::sampling_config()`
/// against whatever `current_model_id` names right now. A test can therefore
/// never set one without the other — which is the whole point of H-a.
fn rebuild_shared_config(agent: &MvpAgent) {
let baseline = agent.models_manager.sampling_config();
agent.sampling_config_platform.set(baseline.platform);
*agent.sampling_config.borrow_mut() = baseline.config;
}
/// H1 — `ModelsManager::sampling_config()`, the OTHER `api_key` producer, and a
/// byte-for-byte repeat of the round-1 defect: it resolved the session bearer
/// itself (`platform.oauth()`, else `auth_manager.current_or_expired()`), so
/// every non-OAuth platform got the primary Kimi token.
///
/// This config is not incidental — it is the `MvpAgent` baseline
/// (`Self::with_models`), which `resolve_sampling_config_for_model` returns
/// verbatim for an unresolved model id and `SubagentSpawnContext` clones as
/// every subagent's baseline, so its `api_key` reaches the wire against its own
/// `base_url`. Zero-config repro: a Kimi subscription + a bundled
/// `moonshot-cn/*` default.
///
/// Revert-to-red (L: this edit COMPILES — the previous wording named a
/// `Option<String>` argument that the `Option<&SessionCredential>` signature
/// rejects, so it could never have been run): in
/// `ModelsManager::sampling_config`, ask the authority about the SESSION's
/// endpoint instead of the current model's own —
/// `.credential_for(None, &config.endpoints.proxy_url())` in place of
/// `.credential_for_model(current_model)`. That is the round-1 defect's shape
/// (the credential decided by something other than the model's own platform +
/// endpoint) and every `assert_ne!` below sees `Some(KIMI_TOKEN)`.
#[tokio::test]
#[serial_test::serial]
async fn models_manager_sampling_config_never_carries_the_kimi_bearer() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
for (catalog_key, slug, base_url) in [
(
"moonshot-cn/kimi-k2-turbo-preview",
"kimi-k2-turbo-preview",
"https://api.moonshot.cn/v1",
),
(
"deepseek/deepseek-chat",
"deepseek-chat",
"https://api.deepseek.com/v1",
),
("openai/gpt-5.2", "gpt-5.2", "https://api.openai.com/v1"),
] {
agent
.models_manager
.insert_test_entry(catalog_key, platform_entry(catalog_key, slug, base_url));
agent
.models_manager
.set_current_model_id(acp::ModelId::new(catalog_key));
let cfg = agent.models_manager.sampling_config().config;
assert_eq!(cfg.base_url, base_url, "{catalog_key}: routed to its own host");
assert_ne!(
cfg.api_key.as_deref(),
Some(KIMI_TOKEN),
"LEAK: ModelsManager::sampling_config sent the Kimi bearer to {base_url}"
);
}
// …and the first-party subscription channel is unchanged, which is what
// proves the Kimi bearer was reachable above.
let kimi_key = "kimi-code/kimi-for-coding";
agent.models_manager.insert_test_entry(
kimi_key,
platform_entry(
kimi_key,
"kimi-for-coding",
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
),
);
agent
.models_manager
.set_current_model_id(acp::ModelId::new(kimi_key));
assert_eq!(
agent
.models_manager
.sampling_config()
.config
.api_key
.as_deref(),
Some(KIMI_TOKEN),
"the kimi-code subscription channel must be byte-identical"
);
}
/// H2 — the SHARED `MvpAgent::sampling_config`. `seed_client_config_auth_if_available`
/// (from `new_session` / `load_session`) and the `cached_token` / `kimi.com/oidc`
/// login handlers all stamped `sampling_config.api_key = Some(<primary bearer>)`
/// with NO platform or endpoint guard — while that very config may point at a
/// third-party model, and is both the subagent baseline and the
/// unresolved-model fallback. `agent_ops`' generic-OAuth login handler already
/// documents the correct rule ("Do NOT stamp this token onto the shared
/// sampling_config"); the Kimi handlers violated it.
///
/// Revert-to-red: make `stamp_session_credential` skip the authority entirely —
/// `self.sampling_config.borrow_mut().api_key =
/// self.auth_manager.current_or_expired().map(|a| a.key); return true;` — and
/// the third-party rows below become `Some(KIMI_TOKEN)`.
#[tokio::test]
#[serial_test::serial]
async fn shared_sampling_config_is_never_stamped_off_the_session_endpoint() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
for (catalog_key, slug, base_url) in [
(
"deepseek/deepseek-chat",
"deepseek-chat",
"https://api.deepseek.com/v1",
),
(
"moonshot-cn/kimi-k2-turbo-preview",
"kimi-k2-turbo-preview",
"https://api.moonshot.cn/v1",
),
] {
agent
.models_manager
.insert_test_entry(catalog_key, platform_entry(catalog_key, slug, base_url));
agent
.models_manager
.set_current_model_id(acp::ModelId::new(catalog_key));
rebuild_shared_config(&agent);
{
let mut shared = agent.sampling_config.borrow_mut();
assert_eq!(shared.model, slug, "{catalog_key}: built from this entry");
assert_eq!(shared.base_url, base_url);
shared.api_key = None;
}
// The `new_session` / `load_session` seed…
agent.seed_client_config_auth_if_available();
assert_eq!(
agent.sampling_config.borrow().api_key, None,
"LEAK: seeding stamped the Kimi bearer onto a config routed at {base_url}"
);
// …and the login handlers, which overwrite rather than seed.
assert!(
!agent.stamp_session_credential(true),
"LEAK: a login handler stamped the Kimi bearer onto a config routed at {base_url}"
);
assert_eq!(agent.sampling_config.borrow().api_key, None);
}
// Byte-identical on the session's own endpoint: both paths still stamp.
let kimi_key = "kimi-code/kimi-for-coding";
agent.models_manager.insert_test_entry(
kimi_key,
platform_entry(
kimi_key,
"kimi-for-coding",
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
),
);
agent
.models_manager
.set_current_model_id(acp::ModelId::new(kimi_key));
rebuild_shared_config(&agent);
agent.sampling_config.borrow_mut().api_key = None;
agent.seed_client_config_auth_if_available();
assert_eq!(
agent.sampling_config.borrow().api_key.as_deref(),
Some(KIMI_TOKEN),
"the subscription endpoint must still be seeded (this is what makes the \
assertions above meaningful)"
);
}
/// C1 — THE CRITICAL. The login stamp used to pair the right QUESTION (*may
/// **a** session credential ride here?*) with the wrong CREDENTIAL (always
/// `auth_manager.current_or_expired().key`, the primary Kimi bearer). For a
/// subscription-OAuth platform at its OWN host the answer is correctly "yes" —
/// the credential that may ride there is that platform's POOLED token
/// ([`CredentialClass::Pooled`]) — so a Claude Pro/Max user whose current model is
/// `claude-pro-max/*` ran `kigi login` and the Kimi subscription bearer landed
/// on a config routed at `api.anthropic.com`. From there it reaches the wire via
/// `resolve_sampling_config_for_model`'s verbatim fallback (offline / stale
/// catalog) and via `SubagentSpawnContext`'s baseline clone.
///
/// All FOUR subscription platforms, each at its own registry host, derived from
/// the registry so the fixture cannot drift.
///
/// The pooled managers are empty here (`pool_home()` is a per-process temp path
/// under `cfg(test)`), so the correct answer is `None` — and `None` is also what
/// proves the primary is not being substituted, because the same agent DOES
/// stamp `KIMI_TOKEN` on its own coding endpoint at the end of the test.
///
/// Revert-to-red (production, compiles): restore the old shape in
/// `MvpAgent::stamp_session_credential` —
/// ```ignore
/// if self.credential_authority().credential_class(platform, &base_url)
/// == CredentialClass::None
/// {
/// return false;
/// }
/// let Some(auth) = self.auth_manager.current_or_expired() else { return false };
/// self.sampling_config.borrow_mut().api_key = Some(auth.key);
/// true
/// ```
/// and every OAuth row below becomes `Some(KIMI_TOKEN)`.
#[tokio::test]
#[serial_test::serial]
async fn oauth_platform_shared_config_never_receives_the_primary_on_login() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
for (platform_id, slug) in [
("claude-pro-max", "claude-opus-4-8"),
("openai-codex", "gpt-5.5-codex"),
("github-copilot", "gpt-4.1"),
("xai-grok", "grok-4.5"),
] {
let platform = kigi_models::PlatformId::parse(platform_id).expect("known platform");
let base_url = platform.base_url();
let catalog_key = format!("{platform_id}/{slug}");
agent
.models_manager
.insert_test_entry(&catalog_key, platform_entry(&catalog_key, slug, &base_url));
agent
.models_manager
.set_current_model_id(acp::ModelId::new(catalog_key.clone()));
rebuild_shared_config(&agent);
{
let mut shared = agent.sampling_config.borrow_mut();
assert_eq!(shared.model, slug, "{catalog_key}: built from this entry");
assert_eq!(shared.base_url, base_url);
shared.api_key = None;
}
// Precondition: this endpoint DOES take a session credential — that is
// exactly why guarding a hand-carried primary with "may a session
// credential ride?" was the bug. The class names WHICH one: the
// platform's POOLED token, never the primary / house key.
assert_eq!(
agent
.credential_authority()
.credential_class(Some(platform), &base_url),
CredentialClass::Pooled,
"{catalog_key}: precondition — its own host takes its POOLED token, \
and never the primary / house credential"
);
// `kigi login` (cached_token and kimi.com/oidc both land here) …
assert!(
!agent.stamp_session_credential(true),
"LEAK: a Kimi login stamped a credential onto a config routed at {base_url}"
);
assert_eq!(
agent.sampling_config.borrow().api_key,
None,
"LEAK: {catalog_key} received a bearer that is not its own pooled token"
);
// … and the `new_session` / `load_session` seed.
agent.seed_client_config_auth_if_available();
assert_ne!(
agent.sampling_config.borrow().api_key.as_deref(),
Some(KIMI_TOKEN),
"LEAK: seeding sent the Kimi subscription bearer to {base_url}"
);
}
// The first-party channel is untouched — this is what makes every
// assertion above meaningful (the Kimi bearer IS live and IS stampable).
let kimi_key = "kimi-code/kimi-for-coding";
agent.models_manager.insert_test_entry(
kimi_key,
platform_entry(
kimi_key,
"kimi-for-coding",
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
),
);
agent
.models_manager
.set_current_model_id(acp::ModelId::new(kimi_key));
rebuild_shared_config(&agent);
agent.sampling_config.borrow_mut().api_key = None;
assert!(agent.stamp_session_credential(true));
assert_eq!(
agent.sampling_config.borrow().api_key.as_deref(),
Some(KIMI_TOKEN),
"the Kimi subscription channel must be byte-identical"
);
}
/// H-a (AVAILABILITY, HIGH) — the stamp guard must resolve the model the shared
/// config was BUILT from, not the one `current_model_id()` names NOW.
///
/// The shared config is built ONCE (`MvpAgent::with_models`) and never rebuilt,
/// while `current_model_id()` moves on every non-Leader model switch
/// (`handlers/model_switch.rs`) and on catalog reselection. Once they drift, a
/// guard re-resolving the config's BARE slug against the live cell fails
/// `entry_for_slug_resolution`'s `entry.info.model == slug` test and falls
/// through to `resolve_catalog_key`'s `.rev()` scan.
///
/// `kimi-code` (subscription, `uses_oauth`) and `kimi-coding` (API-key twin,
/// SAME coding host, `uses_oauth: false`) list the same routing slug, and
/// `PlatformId::ALL` puts `kimi-code` 1st and `kimi-coding` 19th — so that scan
/// answers `kimi-coding`, which takes NO session credential and therefore has NO
/// governing manager. Because `stamp_session_credential` only overwrites ON
/// SUCCESS, a Kimi-subscription user who also has `KIMI_API_KEY` set kept the
/// EXPIRED bearer in the shared config after a successful re-login: every
/// unresolved-model fallback and every subagent baseline turn 401'd until
/// restart.
///
/// The switch happens AFTER the config is built — the previous version of this
/// test left both on the same model, so it could not catch this.
///
/// Revert-to-red (production, compiles): make `MvpAgent::shared_config_platform`
/// re-resolve from the live cell instead of returning the captured value —
/// ```ignore
/// fn shared_config_platform(&self) -> Option<kigi_models::PlatformId> {
/// let model = self.sampling_config.borrow().model.clone();
/// let current_key = self.models_manager.current_model_id();
/// crate::agent::models::platform_for_slug(
/// &self.models_manager.models(),
/// Some(current_key.0.as_ref()),
/// &model,
/// )
/// }
/// ```
/// and the re-login assertion below sees the stale bearer.
#[tokio::test]
#[serial_test::serial]
async fn relogin_restamps_the_shared_config_after_the_model_switched() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
let coding_host = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
let slug = "kimi-for-coding";
// Insertion order mirrors `PlatformId::ALL`: the subscription platform
// first, its API-key twin later — so the `.rev()` scan answers the twin.
for catalog_key in ["kimi-code/kimi-for-coding", "kimi-coding/kimi-for-coding"] {
agent
.models_manager
.insert_test_entry(catalog_key, platform_entry(catalog_key, slug, coding_host));
}
assert_eq!(
crate::agent::models::platform_for_slug(&agent.models_manager.models(), None, slug),
Some(kigi_models::PlatformId::KimiCoding),
"precondition: a `None`-keyed slug scan answers the API-key twin, which takes \
no session credential"
);
// Startup: the picker selected the SUBSCRIPTION entry, and the shared config
// was built from it (config + platform, one call).
agent
.models_manager
.set_current_model_id(acp::ModelId::new("kimi-code/kimi-for-coding"));
rebuild_shared_config(&agent);
assert_eq!(
agent.sampling_config.borrow().model,
slug,
"precondition: the shared config carries the BARE slug, which collides"
);
// The user switches model: `current_model_id` moves to the API-key twin
// while the shared config still represents the subscription entry.
agent
.models_manager
.set_current_model_id(acp::ModelId::new("kimi-coding/kimi-for-coding"));
// The session expires and `kigi login` re-mints it. The handlers overwrite
// (`stamp_session_credential(true)`) AFTER the manager holds the new token.
agent.sampling_config.borrow_mut().api_key = Some("expired-bearer".to_string());
assert!(
agent.stamp_session_credential(true),
"a successful re-login must restamp the shared config"
);
assert_eq!(
agent.sampling_config.borrow().api_key.as_deref(),
Some(KIMI_TOKEN),
"the re-login must REPLACE the expired bearer: resolving the API-key twin \
instead finds no governing manager, leaves the stale token in place, and \
401s every subagent baseline turn until restart"
);
// And the seed path (`new_session` / `load_session`) agrees.
agent.sampling_config.borrow_mut().api_key = None;
agent.seed_client_config_auth_if_available();
assert_eq!(
agent.sampling_config.borrow().api_key.as_deref(),
Some(KIMI_TOKEN),
);
}
/// H3 (REGRESSION) — a MANAGED deployment configures its coding endpoint with
/// `[endpoints] coding_api_base_url` in **config.toml** (what the managed-config
/// sync writes), NOT the `KIGI_CODE_BASE_URL` env var. The previous round's
/// predicate knew only the env var, so `EndpointsConfig::proxy_url()`'s
/// config-key branch was invisible: every model inheriting the managed endpoint
/// classified third-party, lost its api_key AND its resolver, and 401'd on every
/// turn.
///
/// NOTE: no env var is set anywhere in this test — that is the point.
///
/// Revert-to-red: drop the `proxy_url()` arm from
/// `CredentialAuthority::is_session_coding_endpoint` and both `assert_eq!`s
/// below become `None`.
#[tokio::test]
#[serial_test::serial]
async fn managed_config_toml_coding_endpoint_keeps_the_session_bearer() {
let _env = without_ambient_byok_env();
let _no_env_override = EnvGuard::unset("KIGI_CODE_BASE_URL");
let managed = "https://proxy.acme.example/v1";
let dir = tempfile::tempdir().expect("tempdir");
let auth_manager = std::sync::Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
auth_manager.hot_swap(KimiAuth {
key: KIMI_TOKEN.to_string(),
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..KimiAuth::test_default()
});
let cfg = AgentConfig {
endpoints: EndpointsConfig {
coding_api_base_url: Some(managed.to_string()),
..EndpointsConfig::default()
},
..AgentConfig::default()
};
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let agent = MvpAgent::new(GatewaySender::new(tx), &cfg, auth_manager, None)
.expect("valid test config");
agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID));
assert_eq!(
agent.models_manager.endpoints().proxy_url(),
managed,
"precondition: the session's effective coding endpoint is the config.toml key"
);
// A `[model.*]` entry inheriting the managed endpoint …
let mut bare = ModelEntry::fallback("kigi-4.5", &cfg.endpoints);
bare.info.id = None;
bare.info.base_url = managed.to_string();
assert_eq!(
agent
.prepare_sampling_config_for_model(&bare, None)
.api_key
.as_deref(),
Some(KIMI_TOKEN),
"a managed deployment must still receive the session bearer"
);
// … and the kimi-code catalog entry, whose base_url IS `proxy_url()`.
let kimi = platform_entry("kimi-code/kimi-for-coding", "kimi-for-coding", managed);
assert_eq!(
agent
.prepare_sampling_config_for_model(&kimi, None)
.api_key
.as_deref(),
Some(KIMI_TOKEN),
"the managed kimi-code entry must still receive the session bearer"
);
// The spec's requirement was "rides AND still refreshes": an api_key alone
// freezes at login and 401s unrecoverably ~1h in. The managed endpoint must
// also keep a LIVE manager — the primary's, so mid-session refresh and 401
// recovery run against the credential that actually owns that host.
for platform in [None, Some(kigi_models::PlatformId::KimiCode)] {
let manager = agent
.credential_authority()
.manager_for(platform, managed)
.unwrap_or_else(|| {
panic!("{platform:?}: a managed deployment must keep a live manager")
});
assert!(
std::sync::Arc::ptr_eq(&manager, &agent.auth_manager),
"{platform:?}: and it must be the session's OWN primary manager"
);
let resolver = agent
.credential_authority()
.bearer_resolver_for(platform, managed)
.unwrap_or_else(|| panic!("{platform:?}: … exposed as a live bearer_resolver"));
assert_eq!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"{platform:?}: the resolver reads the live primary session bearer"
);
}
// The guard still holds: a third-party host under the SAME managed config
// gets nothing — no key, and no resolver either.
let third_party = platform_entry(
"deepseek/deepseek-chat",
"deepseek-chat",
"https://api.deepseek.com/v1",
);
assert_eq!(
agent
.prepare_sampling_config_for_model(&third_party, None)
.api_key,
None,
"LEAK: a managed deployment must not widen the trust set to third parties"
);
assert!(
agent
.credential_authority()
.bearer_resolver_for(
Some(kigi_models::PlatformId::DeepSeek),
"https://api.deepseek.com/v1"
)
.is_none(),
"LEAK: nor hand it a live resolver over the primary"
);
}
/// M3 (COMPLETION) — the SUMMARY client's `bearer_resolver` must honour the
/// session-token gate, exactly as the session actor's aux path does.
///
/// `build_summary_client` set `cfg.bearer_resolver =
/// authority.bearer_resolver_for(platform_for_slug(…), &cfg.base_url)` with NO
/// gate while `SessionActor::aux_bearer_resolver` had one. `SamplingClient::post`
/// REPLACES the request's auth header from that resolver, so an api-key /
/// house-key session whose `[model.session-summary]` block carries its OWN
/// `env_key` on the session's own coding endpoint had that key overwritten by
/// the primary bearer on every summary request. Both now go through ONE rule,
/// `sampler_turn::aux_bearer_resolver_for`.
///
/// The summary slug is deliberately absent from the catalog and from any
/// config: it classifies `NotByok` definitively, so the only variable left is
/// the ACP auth method — which is the gate term under test.
///
/// Revert-to-red (production, compiles): in `MvpAgent::summary_bearer_resolver`,
/// return `self.credential_authority().bearer_resolver_for(platform, base_url)`
/// directly (the pre-fix shape) and the api-key row below resolves `KIMI_TOKEN`.
#[tokio::test]
#[serial_test::serial]
async fn summary_client_resolver_honours_the_session_gate() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
let coding_host = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
let slug = "kigi-summary-aux-not-in-any-catalog";
let models = agent.models_manager.models();
// A session-based method on the session's OWN endpoint: byte-identical, the
// summary model keeps a LIVE resolver over the primary.
agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID));
let resolver = agent
.summary_bearer_resolver(&models, slug, coding_host)
.expect("the first-party subscription summary channel keeps its resolver");
assert_eq!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"…and it reads the live primary session bearer"
);
// An API-KEY session: the gate is inactive, so the summary model's own key
// must survive to the wire instead of being replaced by the primary.
agent.set_auth_method(acp::AuthMethodId::new(
crate::agent::auth_method::XAI_API_KEY_METHOD_ID,
));
assert!(
agent
.summary_bearer_resolver(&models, slug, coding_host)
.is_none(),
"LEAK: an api-key session's summary model had its own key replaced on the \
wire by the primary bearer"
);
}
@@ -1011,18 +1011,13 @@ fn resolve_model_override_to_config(
// stamp onto the child's api.x.ai / api.moonshot.cn credentials. The // stamp onto the child's api.x.ai / api.moonshot.cn credentials. The
// first-party subscription channel still resolves to the primary // first-party subscription channel still resolves to the primary
// (byte-identical). // (byte-identical).
let session_key = crate::auth::oauth_registry::session_key_for_endpoint( let session_key = crate::auth::credential_authority::CredentialAuthority::new(
entry ctx.models_manager.endpoints(),
.info() Some(ctx.auth_manager.clone()),
.id )
.as_deref() .credential_for_model(&entry);
.and_then(kigi_models::parse_managed_model_key)
.map(|(platform, _)| platform),
&entry.info().base_url,
Some(&ctx.auth_manager),
);
let has_session_key = session_key.is_some(); let has_session_key = session_key.is_some();
let mut credentials = resolve_credentials(&entry, session_key.as_deref()); let mut credentials = resolve_credentials(&entry, session_key.as_ref());
credentials.auth_type = subagent_auth_type(Some(&entry), &ctx.auth_method_id); credentials.auth_type = subagent_auth_type(Some(&entry), &ctx.auth_method_id);
let resolved_auth_type = credentials.auth_type; let resolved_auth_type = credentials.auth_type;
let config = sampling_config_for_model(&entry, credentials, ctx.alpha_test_key.clone()); let config = sampling_config_for_model(&entry, credentials, ctx.alpha_test_key.clone());
@@ -2545,8 +2545,9 @@ async fn subagent_override_first_party_model_still_gets_primary_token() {
/// (`subagent_override_non_oauth_model_still_gets_primary_token`), which encoded /// (`subagent_override_non_oauth_model_still_gets_primary_token`), which encoded
/// the defect. /// the defect.
/// ///
/// Revert-to-red: dropping the `platform_takes_session_credential` term from /// Revert-to-red: make `CredentialAuthority::governing_manager`'s
/// `oauth_registry::session_key_for_endpoint` makes `api_key` `Some("kimi-secret")`. /// `Some(platform) => None` arm return `self.primary.clone()` and `api_key`
/// becomes `Some("kimi-secret")`.
#[tokio::test] #[tokio::test]
async fn subagent_override_api_key_platform_never_gets_the_primary_token() { async fn subagent_override_api_key_platform_never_gets_the_primary_token() {
let (_kd, manager) = kimi_primary_with_token("kimi-secret"); let (_kd, manager) = kimi_primary_with_token("kimi-secret");
@@ -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; pub(crate) mod attribution;
mod config; mod config;
pub(crate) mod credential_authority;
pub mod credential_provider; pub mod credential_provider;
pub(crate) mod device; pub(crate) mod device;
pub mod device_code; pub mod device_code;
@@ -15,8 +15,13 @@
//! on-disk token stays fresh and a 401 recovers via the provider's own manager. //! 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`]), //! 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 //! 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 //! snapshot.
//! (oauth platform) or to the session's primary (everything else). //!
//! 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. //! 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())) POOL.get_or_init(|| Mutex::new(HashMap::new()))
} }
/// The kigi home every OAuth-pool call site resolves from. Single definition so /// The kigi home EVERY OAuth-provider construction in this crate resolves from
/// the pool, the aux/summary token routing and the session's inference manager /// the pool, the catalog fetch's per-platform token resolution, the
/// can never read different homes. /// 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 /// Production: [`crate::util::kigi_home::kigi_home`]. LIB TESTS: a
/// process-lifetime `TempDir`, unconditionally — the pool is process-global and /// per-process path under the system temp dir that is deliberately **never
/// every manager it builds starts a never-cancelled proactive-refresh loop, so /// created**. The pool is process-global and every manager it builds starts a
/// a unit test resolving the real `~/.kigi` would read the developer's stored /// never-cancelled proactive-refresh loop, so a unit test resolving the real
/// OAuth tokens and, 60 s later, fire REAL refresh requests against them. /// `~/.kigi` would read the developer's stored OAuth tokens and, 60 s later,
/// Deliberately not a per-test opt-in that can be forgotten: `kigi_home()` is /// fire REAL refresh requests against them. Deliberately not a per-test opt-in
/// itself a `OnceLock` an earlier test has usually already resolved to the real /// that can be forgotten: `kigi_home()` is itself a `OnceLock` an earlier test
/// home, so setting `KIGI_SHARE_DIR` in a test cannot pin it after the fact. /// 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 { pub(crate) fn pool_home() -> std::path::PathBuf {
#[cfg(test)] #[cfg(test)]
{ {
static TEST_HOME: OnceLock<tempfile::TempDir> = OnceLock::new(); std::env::temp_dir().join(format!("kigi-oauth-pool-test-{}", std::process::id()))
TEST_HOME
.get_or_init(|| tempfile::tempdir().expect("tempdir for the test OAuth pool"))
.path()
.to_path_buf()
} }
#[cfg(not(test))] #[cfg(not(test))]
crate::util::kigi_home::kigi_home() crate::util::kigi_home::kigi_home()
@@ -89,411 +121,73 @@ pub(crate) fn global_manager_for(
manager 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)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::auth::KimiCodeConfig;
use crate::auth::{AuthMode, KimiAuth};
/// A Kimi manager holding a fixed in-memory bearer, standing in for a fn oauth_for(id: &str) -> &'static kigi_models::OAuthConfig {
/// session's primary. The `TempDir` is returned so the caller keeps it kigi_models::PlatformId::parse(id)
/// alive; the token is read from memory (`current_or_expired`), so disk .expect("known platform")
/// 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
.oauth() .oauth()
.expect("xai-grok carries an OAuthConfig") .expect("subscription-OAuth platform carries an OAuthConfig")
} }
fn claude_oauth() -> &'static kigi_models::OAuthConfig { /// Each subscription-OAuth platform gets its OWN process-global pooled
kigi_models::PlatformId::ClaudeProMax /// manager, and no two share one. (Which credential governs a REQUEST is
.oauth() /// not decided here — see
.expect("claude-pro-max carries an OAuthConfig") /// [`crate::auth::credential_authority::CredentialAuthority`].)
}
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.
#[tokio::test] #[tokio::test]
async fn openai_codex_model_resolves_to_its_own_manager_not_kimi() { async fn every_oauth_scope_gets_its_own_pooled_manager() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap(); let home = tempfile::tempdir().unwrap();
let resolved = manager_for_model(home.path(), "openai-codex/gpt-5.5", Some(&kimi)) let ids = [
.expect("openai-codex model resolves to its pooled manager"); "xai-grok",
assert!( "claude-pro-max",
!Arc::ptr_eq(&resolved, &kimi), "github-copilot",
"openai-codex must NOT resolve to the Kimi manager" "openai-codex",
); ];
assert!( let managers: Vec<_> = ids
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), codex_oauth())), .iter()
"openai-codex must resolve to its OWN process-global pooled manager" .map(|id| global_manager_for(home.path(), oauth_for(id)))
); .collect();
assert!( for (i, a) in managers.iter().enumerate() {
!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");
assert!( assert!(
Arc::ptr_eq(&resolved, &kimi), Arc::ptr_eq(a, &global_manager_for(home.path(), oauth_for(ids[i]))),
"{key} must resolve to the primary manager" "{}: 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 /// M8: the test pool home is a per-process path that is never created, so a
/// non-oauth key, never a panic — and without touching the pool. /// test binary leaves nothing behind (and never resolves the developer's
#[test] /// real `~/.kigi`, whose stored OAuth tokens the pool would otherwise read
fn none_primary_is_passed_through_for_non_oauth() { /// and proactively refresh over the network).
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.
/// ///
/// Revert-to-red: dropping the `platform_takes_session_credential` term /// This is also the tripwire for the cleanup claim in [`pool_home`]: a
/// from `session_key_for_endpoint` returns `Some("kimi-tok")` here. /// 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] #[test]
fn session_key_for_a_third_party_endpoint_is_never_the_primary_token() { fn test_pool_home_is_disposable_and_never_the_real_home() {
let (_kd, kimi) = primary_with_token("kimi-tok"); let home = pool_home();
for key in [ assert!(
"moonshot-cn/kimi-k2", home.starts_with(std::env::temp_dir()),
"deepseek/deepseek-chat", "the test pool home must live under the system temp dir, got {home:?}"
"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"
); );
} assert!(
!home.exists(),
/// LEAK guard (aux-model + subagent-override token routing): a grok key with "the test pool home must not be created — nothing to clean up"
/// 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()),
); );
} }
} }
@@ -417,6 +417,21 @@ pub(crate) struct SessionActor {
/// [`SessionActor::model_auth_facts`]. /// [`SessionActor::model_auth_facts`].
pub(crate) model_auth_facts: pub(crate) model_auth_facts:
std::cell::RefCell<Option<(String, crate::agent::config::ModelAuthFacts)>>, std::cell::RefCell<Option<(String, crate::agent::config::ModelAuthFacts)>>,
/// The catalog KEY this session's model was selected by (`{platform}/{model}`
/// for a registry model), owned PER SESSION.
///
/// H4: `SamplingConfig::model` is the bare routing slug, and duplicate slugs
/// across an API-key platform and its subscription-OAuth twin
/// (`xai`/`xai-grok`, `anthropic`/`claude-pro-max`, `openai`/`openai-codex`)
/// are BY DESIGN, so the slug alone cannot name the platform. This used to
/// be read from `ModelsManager::current_model_id()` — a single
/// PROCESS-GLOBAL cell that Leader mode never writes
/// (`agent/handlers/model_switch.rs`) and that is last-writer-wins across
/// concurrent sessions, so both collision directions resolved the wrong
/// platform: the subscription session lost its resolver (unrecoverable 401
/// ~1h in) and the API-key session got the pooled OAuth bearer stamped over
/// its own `sk-…` key. Written at spawn and on every `SetSessionModel`.
pub(crate) selected_catalog_key: std::cell::RefCell<Option<String>>,
/// 401-attribution callback. Joined with the bearer the /// 401-attribution callback. Joined with the bearer the
/// sampler sends on the wire to emit an `auth 401 attribution` /// sampler sends on the wire to emit an `auth 401 attribution`
/// event at each of the six `OaiCompatClient` 401 arms in /// event at each of the six `OaiCompatClient` 401 arms in
@@ -5,12 +5,18 @@ impl SessionActor {
pub(super) async fn handle_set_session_model( pub(super) async fn handle_set_session_model(
&self, &self,
sampling_config: kigi_sampler::SamplerConfig, sampling_config: kigi_sampler::SamplerConfig,
catalog_key: Option<String>,
use_concise: bool, use_concise: bool,
apply_prompt_override: bool, apply_prompt_override: bool,
skip_prompt_rewrite: bool, skip_prompt_rewrite: bool,
auto_compact_threshold_percent: u8, auto_compact_threshold_percent: u8,
) -> Result<acp::ModelId, acp::Error> { ) -> Result<acp::ModelId, acp::Error> {
let model_id = acp::ModelId::new(sampling_config.model.clone()); let model_id = acp::ModelId::new(sampling_config.model.clone());
// H4: record the picker's catalog KEY as this SESSION's own selection.
// `sampling_config.model` is the ambiguous bare slug; the key is what
// disambiguates an API-key platform from its subscription-OAuth twin,
// and it must never come from the process-global `current_model_id()`.
*self.selected_catalog_key.borrow_mut() = catalog_key;
let new_context_window = self.compaction.context_window_override.unwrap_or_else(|| { let new_context_window = self.compaction.context_window_override.unwrap_or_else(|| {
std::num::NonZeroU64::new(sampling_config.context_window).unwrap_or_else(|| { std::num::NonZeroU64::new(sampling_config.context_window).unwrap_or_else(|| {
std::num::NonZeroU64::new(DEFAULT_CONTEXT_WINDOW) std::num::NonZeroU64::new(DEFAULT_CONTEXT_WINDOW)
@@ -64,15 +70,13 @@ impl SessionActor {
// grok model reads the xai-grok token (used only to classify the // grok model reads the xai-grok token (used only to classify the
// credential's auth_type here), never the Kimi one. Kimi / non-oauth // credential's auth_type here), never the Kimi one. Kimi / non-oauth
// models resolve to the primary — byte-identical. // models resolve to the primary — byte-identical.
let session_key = self let session_key = self.session_credential_for_model(&sampling_config.model);
.auth_manager_for_model(&sampling_config.model)
.and_then(|am| am.current_or_expired().map(|a| a.key));
self.chat_state_handle self.chat_state_handle
.update_credentials(kigi_chat_state::Credentials { .update_credentials(kigi_chat_state::Credentials {
api_key: sampling_config.api_key.clone(), api_key: sampling_config.api_key.clone(),
auth_type: crate::agent::config::resolve_chat_state_auth_type( auth_type: crate::agent::config::resolve_chat_state_auth_type(
sampling_config.model.as_str(), sampling_config.model.as_str(),
session_key.as_deref(), session_key.as_ref(),
existing.auth_type, existing.auth_type,
), ),
alpha_test_key: existing.alpha_test_key, alpha_test_key: existing.alpha_test_key,
@@ -624,23 +624,22 @@ impl SessionActor {
let resolved_describe = self let resolved_describe = self
.resolve_aux_sampler_config(&self.image_description_model) .resolve_aux_sampler_config(&self.image_description_model)
.await; .await;
// LEAK 1b: only re-point the aux bearer_resolver when the aux model // LEAK 1b: the aux bearer_resolver is decided at the chokepoint from the
// actually resolved (Some) — the `None` fallback yields the SESSION // IMAGE-DESCRIBE model's own platform + endpoint and passed in
// config, whose Kimi resolver must stay as-is. // explicitly, so an aux model on another provider can never inherit the
let aux_resolved = resolved_describe.is_some(); // session (Kimi) resolver and have its own key overwritten on the aux
let (describe_model, mut sampler_config) = // host. The `None` fallback yields the SESSION config verbatim, whose
// own resolver must stay as-is.
let describe_resolver = resolved_describe
.as_ref()
.map(|cfg| self.aux_bearer_resolver(&self.image_description_model, &cfg.base_url));
let (describe_model, sampler_config) =
crate::agent::config::finalize_image_describe_sampler_config( crate::agent::config::finalize_image_describe_sampler_config(
resolved_describe, resolved_describe,
&active_session_config, &active_session_config,
describe_resolver.flatten(),
Some(self.max_retries), Some(self.max_retries),
); );
// An image-describe model on another provider must not inherit the
// session (Kimi) bearer_resolver stamped by `finalize_*`: re-point it at
// an OAuth model's own manager, or clear it for an API-key-platform /
// third-party endpoint. No-op for the first-party subscription channel.
if aux_resolved {
self.repoint_aux_bearer_resolver(&mut sampler_config, &self.image_description_model);
}
let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| { let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| {
acp::Error::internal_error().data(format!( acp::Error::internal_error().data(format!(
"failed to build image-describe sampling client: {e}" "failed to build image-describe sampling client: {e}"
@@ -300,10 +300,10 @@ pub(super) async fn run_session(
SessionActor::maybe_start_running_task(session.clone(), completion_tx SessionActor::maybe_start_running_task(session.clone(), completion_tx
.clone()). await; } SessionCommand::SessionMode { session_mode, responds_to } .clone()). await; } SessionCommand::SessionMode { session_mode, responds_to }
=> { session.handle_session_mode(session_mode). await; let _ = responds_to => { session.handle_session_mode(session_mode). await; let _ = responds_to
.send(()); } SessionCommand::SetSessionModel { sampling_config, use_concise, .send(()); } SessionCommand::SetSessionModel { sampling_config, catalog_key,
apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent, use_concise, apply_prompt_override, skip_prompt_rewrite,
responds_to } => { let updated_model_id = session auto_compact_threshold_percent, responds_to } => { let updated_model_id =
.handle_set_session_model(sampling_config, use_concise, session.handle_set_session_model(sampling_config, catalog_key, use_concise,
apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent). apply_prompt_override, skip_prompt_rewrite, auto_compact_threshold_percent).
await; let _ = responds_to.send(updated_model_id); } await; let _ = responds_to.send(updated_model_id); }
SessionCommand::RebuildAgentForDefinition { definition, responds_to } => { SessionCommand::RebuildAgentForDefinition { definition, responds_to } => {
@@ -319,11 +319,24 @@ pub(super) async fn run_session(
.signals_handle().set_primary_model(& model_name); cfg.model = model_name .signals_handle().set_primary_model(& model_name); cfg.model = model_name
.clone(); cfg.extra_headers.extend(extra_headers); if let Some(cw) = .clone(); cfg.extra_headers.extend(extra_headers); if let Some(cw) =
context_window && session.compaction.context_window_override.is_none() { cfg context_window && session.compaction.context_window_override.is_none() { cfg
.context_window = cw; } session.chat_state_handle .context_window = cw; } let override_base_url = cfg
.base_url.clone(); session.chat_state_handle
.update_sampling_config(cfg); let existing = session.chat_state_handle .update_sampling_config(cfg); let existing = session.chat_state_handle
.get_credentials(). await; if let Some(r) = crate .get_credentials(). await;
::agent::config::try_resolve_model_credentials(model_name.as_str(), existing // H-c: the rename makes the session's own selected catalog
.api_key.as_deref()) { session.chat_state_handle // key stale unless it still names this model; a stale key is
// exactly what the model→platform rule must not trust.
session.retain_selected_catalog_key_for(& model_name);
// The override model routes to the SAME endpoint the session
// already had; ask the chokepoint whether that endpoint takes
// a session credential rather than re-offering the key
// already in chat state. The platform comes from the session's
// OWN lookup so this and every later turn agree.
let override_session_key = session.credential_authority()
.credential_for(session.model_platform(model_name.as_str()), &
override_base_url); if let Some(r) = crate
::agent::config::try_resolve_model_credentials(model_name.as_str(),
override_session_key.as_ref()) { session.chat_state_handle
.update_credentials(kigi_chat_state::Credentials { api_key : r.api_key, .update_credentials(kigi_chat_state::Credentials { api_key : r.api_key,
auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key, }); } session.model_auth_facts auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key, }); } session.model_auth_facts
.replace(None); } } SessionCommand::GetCurrentModel { responds_to } => { let .replace(None); } } SessionCommand::GetCurrentModel { responds_to } => { let
@@ -29,55 +29,105 @@ pub(super) fn is_auth_tool_error(err: &kigi_tool_runtime::ToolError) -> bool {
/// Gate inputs bundled with the composed decision so the 401-recovery log can /// Gate inputs bundled with the composed decision so the 401-recovery log can
/// report the components. /// report the components.
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
struct SessionTokenAuthGate { pub(crate) struct SessionTokenAuthGate {
is_session_based: bool, is_session_based: bool,
model_byok: crate::agent::auth_method::ModelByok, model_byok: crate::agent::auth_method::ModelByok,
/// Whether the request targets a first-party host. Lets an `Unknown` /// Whether the request targets a first-party host. Lets an `Unknown`
/// BYOK status still refresh against the first-party cli-chat-proxy hosts without /// BYOK status still refresh against the first-party cli-chat-proxy hosts without
/// risking a session-token leak to a third-party BYOK endpoint. /// risking a session-token leak to a third-party BYOK endpoint.
endpoint_is_first_party: bool, endpoint_is_first_party: bool,
/// Whether this model's platform is one whose endpoint accepts a session /// WHICH credential this model's platform/endpoint pair accepts, per the
/// bearer at all (see /// single credential chokepoint
/// [`crate::agent::auth_method::platform_takes_session_credential`]). False /// ([`crate::auth::credential_authority::CredentialAuthority::credential_class`]).
/// for every API-key registry platform, which keeps the primary Kimi bearer /// `None` for every API-key registry platform, which keeps the primary Kimi
/// off `api.deepseek.com` / `api.openai.com` / … . /// bearer off `api.deepseek.com` / `api.openai.com` / … .
endpoint_takes_session_credential: bool, credential_class: crate::auth::credential_authority::CredentialClass,
} }
impl SessionTokenAuthGate { impl SessionTokenAuthGate {
/// Single place `is_session_based` / `endpoint_is_first_party` are derived, /// Single place `is_session_based` / `endpoint_is_first_party` are derived,
/// so all call sites assemble the gate identically. `model_platform` is the /// so all call sites assemble the gate identically. `model_platform` is the
/// registry platform the model routes to (`None` for a bare / `[model.*]` /// registry platform the model routes to (`None` for a bare / `[model.*]`
/// entry) — it MUST be derived from the same lookup /// entry) — it MUST be derived from the same lookup
/// ([`SessionActor::managed_key_for_model`]) that /// ([`SessionActor::model_platform`]) that
/// [`SessionActor::auth_manager_for_model`] uses, so the gate's verdict and /// [`SessionActor::auth_manager_for_endpoint`] uses, so the gate's verdict
/// the manager actually wrapped as the bearer resolver can never disagree. /// and the manager actually wrapped as the bearer resolver can never
fn new( /// disagree. `authority` is that same chokepoint, so the gate cannot answer
/// the endpoint question differently from the manager routing.
pub(crate) fn new(
auth_method_id: Option<&acp::AuthMethodId>, auth_method_id: Option<&acp::AuthMethodId>,
model_byok: crate::agent::auth_method::ModelByok, model_byok: crate::agent::auth_method::ModelByok,
base_url: &str, base_url: &str,
model_platform: Option<kigi_models::PlatformId>, model_platform: Option<kigi_models::PlatformId>,
authority: &crate::auth::credential_authority::CredentialAuthority,
) -> Self { ) -> Self {
Self { Self {
// L13: a model whose OWN credential is a pooled subscription-OAuth
// session is session-based BY ITSELF, whatever the primary ACP
// method is. A user logged in with an API-KEY platform (e.g.
// `deepseek`) who selects a `claude-pro-max/*` model still gets that
// platform's pooled bearer as the request's `api_key` — without this
// term the gate would be inactive, so the config would carry NO
// resolver: the token freezes at selection time and the session dies
// with an unrecoverable 401 once it expires (~1h). The outer
// `credential_class` conjunct keeps this confined to that
// platform's own host.
is_session_based: auth_method_id is_session_based: auth_method_id
.is_some_and(crate::agent::auth_method::is_session_based_method), .is_some_and(crate::agent::auth_method::is_session_based_method)
|| model_platform.is_some_and(|p| p.oauth().is_some()),
model_byok, model_byok,
endpoint_is_first_party: crate::util::is_first_party_url(base_url), endpoint_is_first_party: crate::util::is_first_party_url(base_url),
endpoint_takes_session_credential: credential_class: authority.credential_class(model_platform, base_url),
crate::agent::auth_method::platform_takes_session_credential(
model_platform,
base_url,
),
} }
} }
fn active(self) -> bool { pub(crate) fn active(self) -> bool {
crate::agent::auth_method::session_token_auth_gate( crate::agent::auth_method::session_token_auth_gate(
self.is_session_based, self.is_session_based,
self.model_byok, self.model_byok,
self.endpoint_is_first_party, self.endpoint_is_first_party,
self.endpoint_takes_session_credential, self.credential_class,
) )
} }
} }
/// THE aux / summary `bearer_resolver` rule, stated ONCE.
///
/// `SamplingClient::post` REPLACES the request's auth header from the resolver,
/// so an aux model on a different provider would have its own correctly-resolved
/// key overwritten by the session bearer ON THE AUX HOST. An OAuth aux model
/// gets a live resolver over ITS OWN pooled manager (keeping mid-session
/// refresh); a first-party aux model gets the primary's, but ONLY when the
/// session-token gate is active; everything else gets `None`, so the aux model's
/// own key survives to the wire.
///
/// M3 — the FIRST-PARTY case honours the gate, which is what the old "copy
/// `active_session_config.bearer_resolver`" shape did implicitly: that field is
/// `None` whenever the gate is inactive. Without it, a BYOK / api-key session
/// with a `[model.*]` aux entry carrying its OWN `env_key` on the session's own
/// coding endpoint has that key REPLACED on the wire by the primary bearer on
/// every image-describe / auto-mode-classifier / summary request. A
/// subscription-OAuth aux model is deliberately NOT gated this way: its pooled
/// token IS its credential, and withholding the resolver only costs it
/// mid-session refresh (L13).
///
/// Shared by [`SessionActor::aux_bearer_resolver`] and
/// `MvpAgent::summary_bearer_resolver`: the summary client is built by the
/// AGENT, not the session actor, and its own private copy of this rule is
/// exactly how it stayed ungated after M3 closed the session-actor side.
pub(crate) fn aux_bearer_resolver_for(
authority: &crate::auth::credential_authority::CredentialAuthority,
auth_method_id: Option<&acp::AuthMethodId>,
platform: Option<kigi_models::PlatformId>,
model_byok: crate::agent::auth_method::ModelByok,
base_url: &str,
) -> Option<kigi_sampler::SharedBearerResolver> {
let is_primary_channel = platform.is_none_or(|p| p.oauth().is_none());
if is_primary_channel
&& !SessionTokenAuthGate::new(auth_method_id, model_byok, base_url, platform, authority)
.active()
{
return None;
}
authority.bearer_resolver_for(platform, base_url)
}
/// Run a tool call; on an auth-shaped failure, attempt recovery via /// Run a tool call; on an auth-shaped failure, attempt recovery via
/// `AuthManager` and one retry. When `shared_recovery` is `Some`, concurrent /// `AuthManager` and one retry. When `shared_recovery` is `Some`, concurrent
/// 401s in the same batch deduplicate via `OnceCell::get_or_init`. /// 401s in the same batch deduplicate via `OnceCell::get_or_init`.
@@ -125,7 +175,8 @@ where
/// [`BearerResolver`](kigi_sampler::BearerResolver), resolving the live /// [`BearerResolver`](kigi_sampler::BearerResolver), resolving the live
/// (current-or-expired) bearer at request time. Shared by /// (current-or-expired) bearer at request time. Shared by
/// [`SessionActor::reconstruct_full_config`] (the session model) and the /// [`SessionActor::reconstruct_full_config`] (the session model) and the
/// aux-model bearer repoint ([`SessionActor::repoint_aux_bearer_resolver_for_oauth`]) /// aux-model bearer routing
/// ([`CredentialAuthority::bearer_resolver_for`](crate::auth::credential_authority::CredentialAuthority::bearer_resolver_for))
/// so both wrap ONE definition. SECURITY: the bearer is resolved per request /// so both wrap ONE definition. SECURITY: the bearer is resolved per request
/// and never logged. /// and never logged.
pub(crate) struct AuthManagerBearerResolver(pub(crate) std::sync::Arc<crate::auth::AuthManager>); pub(crate) struct AuthManagerBearerResolver(pub(crate) std::sync::Arc<crate::auth::AuthManager>);
@@ -140,48 +191,11 @@ impl kigi_sampler::BearerResolver for AuthManagerBearerResolver {
} }
} }
/// Wrap `am` as a shared sampler bearer resolver. /// Wrap `am` as a shared sampler bearer resolver.
fn auth_manager_bearer_resolver( pub(crate) fn auth_manager_bearer_resolver(
am: std::sync::Arc<crate::auth::AuthManager>, am: std::sync::Arc<crate::auth::AuthManager>,
) -> kigi_sampler::SharedBearerResolver { ) -> kigi_sampler::SharedBearerResolver {
std::sync::Arc::new(AuthManagerBearerResolver(am)) std::sync::Arc::new(AuthManagerBearerResolver(am))
} }
/// The `bearer_resolver` an AUX / summary `SamplerConfig` may carry, given the
/// SESSION model's resolver (`stamped`) and the AUX model's own platform +
/// endpoint. ONE decision shared by image-describe, the auto-mode classifier
/// (via [`SessionActor::repoint_aux_bearer_resolver`]) and
/// `MvpAgent::build_summary_client`.
///
/// [`crate::agent::config::stamp_session_local_sampler_fields`] copies the
/// session resolver onto every aux config, and `SamplingClient::post` REPLACES
/// the request's auth header from it — so an aux model on a DIFFERENT provider
/// would have its own correctly-resolved key overwritten by the session bearer
/// on the AUX host (H3/H4). Resolve by the aux model instead:
/// - an OAuth platform → a live resolver over ITS OWN pooled manager (so a grok
/// / claude-pro-max / copilot / codex aux model keeps mid-session refresh);
/// - `kimi-code`, or a platform-less model on the session's own coding endpoint
/// → the stamped session resolver, byte-identical;
/// - every API-key registry platform, and any `[model.*]` block pointed at a
/// third-party host → `None`, so the aux model's own key survives to the wire.
///
/// SECURITY: no token is logged.
pub(crate) fn aux_bearer_resolver(
stamped: Option<kigi_sampler::SharedBearerResolver>,
platform: Option<kigi_models::PlatformId>,
base_url: &str,
) -> Option<kigi_sampler::SharedBearerResolver> {
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
return Some(auth_manager_bearer_resolver(
crate::auth::oauth_registry::global_manager_for(
&crate::auth::oauth_registry::pool_home(),
oauth,
),
));
}
if crate::agent::auth_method::platform_takes_session_credential(platform, base_url) {
return stamped;
}
None
}
impl SessionActor { impl SessionActor {
pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) { pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
let mcp_wait_start = std::time::Instant::now(); let mcp_wait_start = std::time::Instant::now();
@@ -225,8 +239,8 @@ impl SessionActor {
let plan_active = self.plan_mode.lock().is_active(); let plan_active = self.plan_mode.lock().is_active();
filter_cursor_tools_by_plan_mode(defs, plan_active) filter_cursor_tools_by_plan_mode(defs, plan_active)
} }
/// Memoized per-model [`ModelAuthFacts`](crate::agent::config::ModelAuthFacts), /// Memoized per-model [`ModelAuthFacts`](crate::agent::config::ModelAuthFacts)
/// keyed by `model_id`. /// for the SESSION's own model, keyed by `model_id`.
/// ///
/// A fresh `Unknown` (config currently unparseable) falls back to the last /// A fresh `Unknown` (config currently unparseable) falls back to the last
/// definite value for the same `model_id` rather than demoting a live session /// definite value for the same `model_id` rather than demoting a live session
@@ -235,6 +249,32 @@ impl SessionActor {
/// `model_id`, keying on `model_id` alone is insufficient — each /// `model_id`, keying on `model_id` alone is insufficient — each
/// model/credential chokepoint must clear this memo (`replace(None)`). /// model/credential chokepoint must clear this memo (`replace(None)`).
pub(super) fn model_auth_facts(&self, model_id: &str) -> crate::agent::config::ModelAuthFacts { pub(super) fn model_auth_facts(&self, model_id: &str) -> crate::agent::config::ModelAuthFacts {
self.resolve_auth_facts(model_id, true)
}
/// [`Self::model_auth_facts`] for a model that is NOT the session's own — an
/// AUX / summary / image-describe slug.
///
/// Identical resolution, but it NEVER WRITES the slot. The memo is a SINGLE
/// slot: when the aux path shared it, one classifier or image-describe call
/// evicted the session model's entry, and (a) the next
/// [`Self::reconstruct_full_config`] paid another `load_effective_config()`
/// + `resolve_model_list()` — the per-turn disk read M7/M9 removed — while
/// (b) a transient `Unknown` for the SESSION model then had no same-`model_id`
/// definite value to fall back to, so it degraded to `endpoint_is_first_party`
/// — `false` for every subscription-OAuth host, costing the session its
/// `bearer_resolver` and 401ing unrecoverably ~1h in (the failure L13
/// prevents). Reading a matching entry is still allowed: it can only hit when
/// the slot already names this same slug.
fn aux_model_auth_facts(&self, model_id: &str) -> crate::agent::config::ModelAuthFacts {
self.resolve_auth_facts(model_id, false)
}
/// Shared body of [`Self::model_auth_facts`] / [`Self::aux_model_auth_facts`].
/// `memoize` is the ONLY difference, so the two can never resolve differently.
fn resolve_auth_facts(
&self,
model_id: &str,
memoize: bool,
) -> crate::agent::config::ModelAuthFacts {
use crate::agent::auth_method::ModelByok; use crate::agent::auth_method::ModelByok;
if let Some((cached_id, facts)) = self.model_auth_facts.borrow().as_ref() if let Some((cached_id, facts)) = self.model_auth_facts.borrow().as_ref()
&& cached_id == model_id && cached_id == model_id
@@ -251,10 +291,12 @@ impl SessionActor {
} }
return fresh; return fresh;
} }
*self.model_auth_facts.borrow_mut() = Some((model_id.to_string(), fresh)); if memoize {
*self.model_auth_facts.borrow_mut() = Some((model_id.to_string(), fresh));
}
fresh fresh
} }
/// Gate inputs for `model_id` routed to `base_url`. See /// Gate inputs for the SESSION model `model_id` routed to `base_url`. See
/// [`crate::agent::auth_method::session_token_auth_gate`] for the rationale /// [`crate::agent::auth_method::session_token_auth_gate`] for the rationale
/// (`base_url` keeps an `Unknown` BYOK status refreshable only /// (`base_url` keeps an `Unknown` BYOK status refreshable only
/// against first-party xAI hosts). /// against first-party xAI hosts).
@@ -266,62 +308,109 @@ impl SessionActor {
byok, byok,
base_url, base_url,
self.model_platform(model_id), self.model_platform(model_id),
&self.credential_authority(),
)
}
/// This session's credential chokepoint: its EFFECTIVE endpoints (so a
/// managed `[endpoints] coding_api_base_url` deployment keeps the session
/// bearer — H3) plus its primary manager, which the authority keeps
/// private. Every inference-auth question this actor asks goes through it.
pub(crate) fn credential_authority(
&self,
) -> crate::auth::credential_authority::CredentialAuthority {
crate::auth::credential_authority::CredentialAuthority::new(
self.models_manager.endpoints(),
self.auth_manager.clone(),
) )
} }
/// The [`AuthManager`](crate::auth::AuthManager) that governs INFERENCE auth /// The [`AuthManager`](crate::auth::AuthManager) that governs INFERENCE auth
/// for the model whose routing slug is `model` (the sampling config's /// for the routing slug `model` against the endpoint the request will
/// `model`). A generic device-code OAuth platform (xai-grok) routes to its /// ACTUALLY be sent to, from the ONE chokepoint
/// OWN scope-keyed manager from the process-global OAuth pool /// ([`crate::auth::credential_authority::CredentialAuthority`]).
/// ([`crate::auth::oauth_registry::manager_for_model`], built on demand from
/// the on-disk token); every other model routes to the primary Kimi
/// `auth_manager`.
/// ///
/// This is the single chokepoint that keeps a grok turn from ever sending /// A subscription-OAuth platform routes to ITS OWN scope-keyed pooled
/// the Kimi bearer (Facet B) and gives it its own proactive-refresh + 401 /// manager; `kimi-code` and a platform-less model on the session's own
/// recovery manager (Facet A). The pool is the single source of truth, so a /// coding endpoint route to the primary; every API-key registry platform,
/// grok login that lands AFTER this session spawned is resolved on the next /// and any endpoint that is neither, routes to `None` — fail fast, never a
/// grok turn (no frozen per-session snapshot). Cheap `Arc` clone. Kimi / /// silent fallback to the primary. `None` also for a BYOK / test session
/// first-party path is byte-identical: a non-oauth model always resolves to /// with no primary.
/// the primary.
/// ///
/// `None` when there is no governing manager: a BYOK / test session with no /// Callers pass the LIVE sampling config's `base_url` so the manager, the
/// primary. A grok model always resolves to its pooled manager (never the /// gate and the wire can never be resolved against three different endpoints
/// Kimi primary); when the user has not logged into grok that manager simply /// (an `OverrideModelName` session keeps its original `base_url` under a
/// holds no token, so no Kimi bearer can leak. /// routing name absent from the catalog). A `model`-only sibling that
pub(super) fn auth_manager_for_model( /// re-derived the endpoint from the CATALOG instead used to exist beside
/// this; it had zero callers and was deleted rather than left as a second,
/// unexercised way to answer the same question (`lib.rs`'s
/// `#![allow(dead_code)]` means such a helper raises no warning).
pub(super) fn auth_manager_for_endpoint(
&self, &self,
model: &str, model: &str,
base_url: &str,
) -> Option<std::sync::Arc<crate::auth::AuthManager>> { ) -> Option<std::sync::Arc<crate::auth::AuthManager>> {
// `model` is the bare routing slug; recover the managed catalog key self.credential_authority()
// (`{platform}/{model}`) so the platform — and thus its OAuth scope — is .manager_for(self.model_platform(model), base_url)
// unambiguous. A bare / config / unlisted model yields no managed key
// and resolves to the primary.
let managed_key = self.managed_key_for_model(model);
crate::auth::oauth_registry::manager_for_model(
&crate::auth::oauth_registry::pool_home(),
managed_key.as_deref().unwrap_or(model),
self.auth_manager.as_ref(),
)
} }
/// Recover the managed catalog key (`{platform}/{model}`) for a routing slug /// The SESSION credential (if any) that may ride a request for the routing
/// from the live catalog. `None` for a bare / config / unlisted model. /// slug `model`. The only producer is the chokepoint.
/// pub(super) fn session_credential_for_model(
/// H5: the catalog KEY the picker selected &self,
/// ([`crate::agent::models::ModelsManager::current_model_id`]) is model: &str,
/// authoritative — `model` is the ambiguous bare slug. See ) -> Option<crate::auth::credential_authority::SessionCredential> {
/// [`crate::agent::models::managed_key_for_slug`]. self.credential_authority()
fn managed_key_for_model(&self, model: &str) -> Option<String> { .credential_for(self.model_platform(model), &self.model_base_url(model))
}
/// The base URL a routing slug actually resolves to in the live catalog.
/// Falls back to the session's own inference endpoint for an unlisted slug,
/// which is exactly where `resolve_aux_model_sampling_config`'s Tier-2
/// fallback entry routes — so the endpoint the rule is applied to is always
/// the endpoint the request is sent to.
fn model_base_url(&self, model: &str) -> String {
let models = self.models_manager.models(); let models = self.models_manager.models();
let current = self.models_manager.current_model_id(); match crate::agent::config::find_model_by_id(&models, model) {
crate::agent::models::managed_key_for_slug(&models, Some(current.0.as_ref()), model) Some(entry) => entry.info().base_url.clone(),
None => self.models_manager.endpoints().resolve_inference_base_url(),
}
}
/// This SESSION's own selected catalog key (H4) — never the process-global
/// `ModelsManager::current_model_id()`, which Leader mode never writes and
/// which is last-writer-wins across concurrent sessions.
pub(super) fn selected_catalog_key(&self) -> Option<String> {
self.selected_catalog_key.borrow().clone()
}
/// Keep the session's own selected catalog key consistent with an
/// `OverrideModelName` rename: KEEP it when it still names `model_name`
/// (same entry, new routing name), otherwise CLEAR it.
///
/// H-c: `OverrideModelName` is the one command that rewrites
/// `SamplingConfig::model` without going through `SetSessionModel`, so it
/// used to leave the field naming a model the session is no longer on.
/// Clearing rather than re-resolving is deliberate: re-resolving would put
/// `resolve_catalog_key`'s `.rev()` guess INTO the field the whole rule
/// treats as the session's deliberate selection, and a cleared field
/// refuses a collided slug instead of guessing its OAuth twin (H-b).
pub(super) fn retain_selected_catalog_key_for(&self, model_name: &str) {
let models = self.models_manager.models();
let still_names_it = self.selected_catalog_key().is_some_and(|key| {
key == model_name
|| models
.get(key.as_str())
.is_some_and(|entry| entry.info.model == model_name)
});
if !still_names_it {
*self.selected_catalog_key.borrow_mut() = None;
}
} }
/// The registry platform the routing slug `model` belongs to, from the SAME /// The registry platform the routing slug `model` belongs to, from the SAME
/// lookup [`Self::auth_manager_for_model`] routes on. `None` for a bare / /// lookup [`Self::auth_manager_for_endpoint`] routes on. `None` for a bare /
/// `[model.*]` / unlisted model. /// `[model.*]` / unlisted model.
fn model_platform(&self, model: &str) -> Option<kigi_models::PlatformId> { pub(super) fn model_platform(&self, model: &str) -> Option<kigi_models::PlatformId> {
let models = self.models_manager.models(); let models = self.models_manager.models();
let current = self.models_manager.current_model_id(); crate::agent::models::platform_for_slug(
crate::agent::models::platform_for_slug(&models, Some(current.0.as_ref()), model) &models,
self.selected_catalog_key().as_deref(),
model,
)
} }
/// Whether `model` routes to the Claude Pro/Max OAuth-Messages platform /// Whether `model` routes to the Claude Pro/Max OAuth-Messages platform
/// (claude-pro-max) — the gate for the sampler's OAuth Messages adaptation /// (claude-pro-max) — the gate for the sampler's OAuth Messages adaptation
@@ -352,20 +441,37 @@ impl SessionActor {
self.model_platform(model) self.model_platform(model)
.is_some_and(kigi_models::PlatformId::sends_codex_responses_headers) .is_some_and(kigi_models::PlatformId::sends_codex_responses_headers)
} }
/// LEAK guard for the stamped aux paths (auto-mode classifier, image /// The `bearer_resolver` an AUX / summary `SamplerConfig` may carry — the
/// describe). Applies [`aux_bearer_resolver`] to the config /// shared [`aux_bearer_resolver_for`] rule applied to the AUX model's own
/// [`crate::agent::config::stamp_session_local_sampler_fields`] just stamped /// platform + endpoint.
/// the SESSION model's `bearer_resolver` onto. ///
pub(super) fn repoint_aux_bearer_resolver( /// The aux config never inherits the session resolver: it is passed this
/// value explicitly (see
/// [`crate::agent::config::stamp_session_local_sampler_fields`]), so
/// "forgot to re-point" is not expressible.
///
/// The BYOK status comes from [`Self::aux_model_auth_facts`], which does NOT
/// write the session model's single-slot memo.
pub(super) fn aux_bearer_resolver(
&self, &self,
cfg: &mut kigi_sampler::SamplerConfig,
slug: &str, slug: &str,
) { base_url: &str,
cfg.bearer_resolver = aux_bearer_resolver( ) -> Option<kigi_sampler::SharedBearerResolver> {
cfg.bearer_resolver.take(), let auth_method = self.auth_method_id.load();
self.model_platform(slug), // An aux slug is NOT the session's selection, so it must not be resolved
&cfg.base_url, // against `selected_catalog_key` — the same rule the aux `api_key` obeys
); // (`credential_for_slug(.., None, ..)`). Keying an aux model on the
// SESSION's selection let a colliding same-vendor slug resolve the OAuth
// twin, whose pooled resolver would then overwrite the user's own key on
// the aux request.
let models = self.models_manager.models();
aux_bearer_resolver_for(
&self.credential_authority(),
auth_method.as_deref(),
crate::agent::models::platform_for_slug(&models, None, slug),
self.aux_model_auth_facts(slug).byok,
base_url,
)
} }
/// Emit a unified-log breadcrumb whenever the session-token refresh gate is /// Emit a unified-log breadcrumb whenever the session-token refresh gate is
/// evaluated with an **`Unknown`** per-model BYOK status on a session-based /// evaluated with an **`Unknown`** per-model BYOK status on a session-based
@@ -385,9 +491,8 @@ impl SessionActor {
let ctx = serde_json::json!( let ctx = serde_json::json!(
{ "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based" { "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based"
: gate.is_session_based, "endpoint_is_first_party" : gate : gate.is_session_based, "endpoint_is_first_party" : gate
.endpoint_is_first_party, "endpoint_takes_session_credential" : gate .endpoint_is_first_party, "credential_class" : gate.credential_class
.endpoint_takes_session_credential, "refresh_active" : refresh_active, .as_str(), "refresh_active" : refresh_active, "base_url" : base_url, }
"base_url" : base_url, }
); );
let sid = Some(self.session_info.id.0.as_ref()); let sid = Some(self.session_info.id.0.as_ref());
if refresh_active { if refresh_active {
@@ -446,6 +551,7 @@ impl SessionActor {
model_facts.byok, model_facts.byok,
&cfg.base_url, &cfg.base_url,
self.model_platform(cfg.model.as_str()), self.model_platform(cfg.model.as_str()),
&self.credential_authority(),
); );
let use_bearer_resolver = gate.active(); let use_bearer_resolver = gate.active();
self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url); self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url);
@@ -455,7 +561,7 @@ impl SessionActor {
// inactive or the oauth provider has no manager (fail-fast, no Kimi // inactive or the oauth provider has no manager (fail-fast, no Kimi
// fallback). // fallback).
let inference_auth_manager = if use_bearer_resolver { let inference_auth_manager = if use_bearer_resolver {
self.auth_manager_for_model(&cfg.model) self.auth_manager_for_endpoint(&cfg.model, &cfg.base_url)
} else { } else {
None None
}; };
@@ -656,17 +762,20 @@ impl SessionActor {
// Kimi session token (which `resolve_credentials` would otherwise stamp // Kimi session token (which `resolve_credentials` would otherwise stamp
// onto an api.x.ai / api.deepseek.com request). The first-party // onto an api.x.ai / api.deepseek.com request). The first-party
// subscription channel still gets the primary (byte-identical). // subscription channel still gets the primary (byte-identical).
let session_key = crate::auth::oauth_registry::session_key_for_catalog_model( // M6: ONE lookup. The platform AND the base URL the rule is applied to
&models, // both come from `credential_for_slug`'s single resolution of `slug`
slug, // against this catalog, so the platform and the endpoint can no longer
self.auth_manager.as_ref(), // disagree (they were previously resolved by two different lookups).
); // Aux slugs are not the session's selection, so no `current_key`.
let session_key = self
.credential_authority()
.credential_for_slug(&models, None, slug);
let endpoints = self.models_manager.endpoints(); let endpoints = self.models_manager.endpoints();
crate::agent::config::resolve_aux_model_sampling_config( crate::agent::config::resolve_aux_model_sampling_config(
slug, slug,
&models, &models,
&endpoints, &endpoints,
session_key.as_deref(), session_key.as_ref(),
creds.alpha_test_key.clone(), creds.alpha_test_key.clone(),
) )
} }
@@ -681,16 +790,17 @@ impl SessionActor {
) -> Option<(kigi_sampler::SamplingClient, String)> { ) -> Option<(kigi_sampler::SamplingClient, String)> {
let active_session_config = self.reconstruct_full_config().await; let active_session_config = self.reconstruct_full_config().await;
let mut cfg = self.resolve_aux_sampler_config(slug).await?; let mut cfg = self.resolve_aux_sampler_config(slug).await?;
// LEAK 1b: the aux classifier must NOT inherit the SESSION model's
// (Kimi) bearer_resolver — the resolver is decided by the AUX model's
// own platform + endpoint at the chokepoint and passed in explicitly,
// so there is no "copy then remember to re-point" step to forget.
let aux_resolver = self.aux_bearer_resolver(slug, &cfg.base_url);
crate::agent::config::stamp_session_local_sampler_fields( crate::agent::config::stamp_session_local_sampler_fields(
&mut cfg, &mut cfg,
&active_session_config, &active_session_config,
aux_resolver,
Some(self.max_retries), Some(self.max_retries),
); );
// LEAK 1b: the aux classifier must not inherit the SESSION model's
// (Kimi) bearer_resolver stamped above — re-point it at an OAuth aux
// model's own manager, or clear it for an API-key-platform / third-party
// aux endpoint. No-op for the first-party subscription channel.
self.repoint_aux_bearer_resolver(&mut cfg, slug);
let model = cfg.model.clone(); let model = cfg.model.clone();
let client = kigi_sampler::SamplingClient::new(cfg) let client = kigi_sampler::SamplingClient::new(cfg)
.map_err(|e| { .map_err(|e| {
@@ -836,8 +946,7 @@ impl SessionActor {
session_id = % self.session_info.id.0, is_session_based = gate session_id = % self.session_info.id.0, is_session_based = gate
.is_session_based, model_byok = gate.model_byok.as_str(), .is_session_based, model_byok = gate.model_byok.as_str(),
endpoint_is_first_party = gate.endpoint_is_first_party, endpoint_is_first_party = gate.endpoint_is_first_party,
endpoint_takes_session_credential = gate credential_class = gate.credential_class.as_str(),
.endpoint_takes_session_credential,
"auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401", "auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401",
); );
kigi_log::unified_log::warn( kigi_log::unified_log::warn(
@@ -848,8 +957,7 @@ impl SessionActor {
.status_code, "is_session_based" : gate.is_session_based, .status_code, "is_session_based" : gate.is_session_based,
"model_byok" : gate.model_byok.as_str(), "model_byok" : gate.model_byok.as_str(),
"endpoint_is_first_party" : gate.endpoint_is_first_party, "endpoint_is_first_party" : gate.endpoint_is_first_party,
"endpoint_takes_session_credential" : gate "credential_class" : gate.credential_class.as_str(), }
.endpoint_takes_session_credential, }
)), )),
); );
} }
@@ -869,13 +977,13 @@ impl SessionActor {
// xai-grok session via the xai-grok manager, never the Kimi one. For a // xai-grok session via the xai-grok manager, never the Kimi one. For a
// Kimi / non-oauth model this resolves to the primary — byte-identical. // Kimi / non-oauth model this resolves to the primary — byte-identical.
if auth_recovery_eligible { if auth_recovery_eligible {
let recovery_model = self let (recovery_model, recovery_base_url) = self
.chat_state_handle .chat_state_handle
.get_sampling_config() .get_sampling_config()
.await .await
.map(|c| c.model) .map(|c| (c.model, c.base_url))
.unwrap_or_default(); .unwrap_or_default();
if let Some(am) = self.auth_manager_for_model(&recovery_model) { if let Some(am) = self.auth_manager_for_endpoint(&recovery_model, &recovery_base_url) {
if am.try_recover_unauthorized().await { if am.try_recover_unauthorized().await {
tracing::info!( tracing::info!(
session_id = % self.session_info.id.0, session_id = % self.session_info.id.0,
@@ -1070,7 +1178,7 @@ impl SessionActor {
// Refresh the ACTIVE model's OWN manager: a grok model refreshes the // Refresh the ACTIVE model's OWN manager: a grok model refreshes the
// xai-grok token via the xai-grok manager, never the Kimi one. For a // xai-grok token via the xai-grok manager, never the Kimi one. For a
// Kimi / non-oauth model this resolves to the primary — byte-identical. // Kimi / non-oauth model this resolves to the primary — byte-identical.
if let Some(am) = self.auth_manager_for_model(&model_id) { if let Some(am) = self.auth_manager_for_endpoint(&model_id, &base_url) {
let creds = self.chat_state_handle.get_credentials().await; let creds = self.chat_state_handle.get_credentials().await;
if self.auth_gate(&model_id, &base_url).active() if self.auth_gate(&model_id, &base_url).active()
&& let Ok(key) = am.get_valid_token().await && let Ok(key) = am.get_valid_token().await
@@ -1102,13 +1210,23 @@ impl SessionActor {
.map(|c| c.model) .map(|c| c.model)
.unwrap_or_default(); .unwrap_or_default();
let Some(ref key) = current_key else { return }; let Some(ref key) = current_key else { return };
// M7: a registry-platform model's key comes from that platform's // M7/M9: a registry-platform model's key normally comes from that
// credential resolved into its catalog entry — it is NEVER a // platform's credential resolved into its catalog entry, so with the
// `[model.*]` block. With the session gate now inactive for every // session gate now inactive for every API-key platform those turns all
// API-key platform, those turns all fell through to here and paid a // fell through to here and paid a `load_effective_config()` disk read
// `load_effective_config()` disk read PER TURN, then logged a // PER TURN, then logged a permanently false "Model not found in
// permanently false "Model not found in config.toml [model.*]" warning. // config.toml [model.*]" warning.
if self.model_platform(&current_model_id).is_some() { //
// But a `[model."deepseek/deepseek-chat"]` override DOES keep the base
// entry's `info.id` (`ConfigModelOverride::apply`), so "has a platform"
// does NOT imply "has no `[model.*]` block" — skipping on the platform
// alone would freeze an on-disk key rotation for the whole session.
// Skip only when the catalog entry carries no own credential at all,
// which is exactly the "key came from the platform, not from config"
// case the disk read cannot improve on.
if self.model_platform(&current_model_id).is_some()
&& !self.model_has_own_credential(&current_model_id)
{
return; return;
} }
let Some(new_key) = self.reload_api_key_from_config(&current_model_id) else { let Some(new_key) = self.reload_api_key_from_config(&current_model_id) else {
@@ -1125,6 +1243,15 @@ impl SessionActor {
creds.api_key = Some(new_key); creds.api_key = Some(new_key);
self.chat_state_handle.update_credentials(creds); self.chat_state_handle.update_credentials(creds);
} }
/// Whether the live catalog entry for `slug` carries its own credential —
/// an `api_key`/`env_key` from a `[model.*]` block, which a config edit can
/// rotate mid-session. A platform entry whose key came from the platform
/// credential has none.
fn model_has_own_credential(&self, slug: &str) -> bool {
let models = self.models_manager.models();
crate::agent::config::find_model_by_id(&models, slug)
.is_some_and(crate::agent::config::ModelEntry::has_own_credentials)
}
fn reload_api_key_from_config(&self, current_model_id: &str) -> Option<String> { fn reload_api_key_from_config(&self, current_model_id: &str) -> Option<String> {
let raw_config = crate::config::load_effective_config() let raw_config = crate::config::load_effective_config()
.map_err(|e| tracing::warn!(error = % e, "Failed to reload config")) .map_err(|e| tracing::warn!(error = % e, "Failed to reload config"))
@@ -1211,8 +1338,8 @@ mod bearer_resolver_tests {
/// of the manager it wraps. So an aux bearer_resolver built over grok's OWN /// of the manager it wraps. So an aux bearer_resolver built over grok's OWN
/// (oauth) pooled manager yields grok's token (or `None`) — NEVER the Kimi /// (oauth) pooled manager yields grok's token (or `None`) — NEVER the Kimi
/// session token that a Kimi-manager resolver would. The /// session token that a Kimi-manager resolver would. The
/// `repoint_aux_bearer_resolver_for_oauth` fix wraps exactly this grok /// [`CredentialAuthority::bearer_resolver_for`](crate::auth::credential_authority::CredentialAuthority::bearer_resolver_for)
/// manager for a grok aux model. /// wraps exactly this grok manager for a grok aux model.
#[tokio::test] #[tokio::test]
async fn resolver_resolves_the_wrapped_manager_never_kimi() { async fn resolver_resolves_the_wrapped_manager_never_kimi() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
@@ -980,10 +980,21 @@ pub(crate) async fn spawn_session_actor(
} }
}; };
let doom_loop_recovery = effective_config.resolve_doom_loop_recovery(); let doom_loop_recovery = effective_config.resolve_doom_loop_recovery();
let session_model_id_for_actor = session_model_id.clone();
let session = Arc::new_cyclic(|weak: &std::sync::Weak<SessionActor>| SessionActor { let session = Arc::new_cyclic(|weak: &std::sync::Weak<SessionActor>| SessionActor {
session_info: session_info.clone(), session_info: session_info.clone(),
auth_method_id, auth_method_id,
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
// H4: seed the session's OWN selected catalog key from the model it was
// spawned with, resolved through the picker's lookup. Never the
// process-global `current_model_id()`. H-c: the rule lives in
// `selected_catalog_key_for_spawn` so it is covered by a test.
selected_catalog_key: std::cell::RefCell::new(
crate::agent::models::selected_catalog_key_for_spawn(
&models_manager.models(),
&session_model_id_for_actor,
),
),
attribution_callback, attribution_callback,
auth_manager, auth_manager,
state, state,
@@ -485,29 +485,44 @@ async fn no_legacy_hint_for_oidc_auth() {
#[test] #[test]
fn session_token_auth_gate_truth_table() { fn session_token_auth_gate_truth_table() {
use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate}; use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate};
use crate::auth::credential_authority::CredentialClass;
// Non-session methods never refresh, regardless of BYOK status or endpoint. // Non-session methods never refresh, regardless of BYOK status or endpoint.
// `Pooled` (an OAuth platform's own pool) and `Primary` (kimi-code, or a
// bare / [model.*] model on the session's own endpoint) behave identically
// here: each names a credential that IS refreshable on that host.
for fp in [false, true] { for fp in [false, true] {
assert!(!gate(false, ModelByok::NotByok, fp, true)); for class in [CredentialClass::Pooled, CredentialClass::Primary] {
assert!(!gate(false, ModelByok::Byok, fp, true)); assert!(!gate(false, ModelByok::NotByok, fp, class));
assert!(!gate(false, ModelByok::Unknown, fp, true)); assert!(!gate(false, ModelByok::Byok, fp, class));
// Session method on an endpoint that DOES take the session credential assert!(!gate(false, ModelByok::Unknown, fp, class));
// (kimi-code, an OAuth platform's own pool, or a bare / [model.*] // Session method on an endpoint that DOES take a session
// model): a definite classification ignores the endpoint — NotByok // credential: a definite classification ignores the endpoint —
// refreshes, a genuine per-model Byok never does. // NotByok refreshes, a genuine per-model Byok never does.
assert!(gate(true, ModelByok::NotByok, fp, true)); assert!(gate(true, ModelByok::NotByok, fp, class));
assert!(!gate(true, ModelByok::Byok, fp, true)); assert!(!gate(true, ModelByok::Byok, fp, class));
// …and an API-key registry platform endpoint is refused on every arm, }
// first-party flag included: that is the leak guard. // …and an API-key registry platform endpoint (`CredentialClass::None`)
assert!(!gate(true, ModelByok::NotByok, fp, false)); // is refused on every arm, first-party flag included: the leak guard.
assert!(!gate(true, ModelByok::Byok, fp, false)); assert!(!gate(true, ModelByok::NotByok, fp, CredentialClass::None));
assert!(!gate(true, ModelByok::Unknown, fp, false)); assert!(!gate(true, ModelByok::Byok, fp, CredentialClass::None));
assert!(!gate(true, ModelByok::Unknown, fp, CredentialClass::None));
} }
// Session method + Unknown BYOK: refresh only against a first-party xAI // Session method + Unknown BYOK: refresh only against a first-party xAI
// host, so a transiently-unclassifiable config can't demote a live session // host, so a transiently-unclassifiable config can't demote a live session
// (the stale-token 401 regression) yet the session token never leaks to a // (the stale-token 401 regression) yet the session token never leaks to a
// third-party BYOK endpoint. This arm was unconditionally `false` pre-fix. // third-party BYOK endpoint. This arm was unconditionally `false` pre-fix.
assert!(gate(true, ModelByok::Unknown, true, true)); assert!(gate(
assert!(!gate(true, ModelByok::Unknown, false, true)); true,
ModelByok::Unknown,
true,
CredentialClass::Primary
));
assert!(!gate(
true,
ModelByok::Unknown,
false,
CredentialClass::Primary
));
} }
/// Pre-fix, the gate read `auth_type` and skipped recovery here, 401'ing every /// Pre-fix, the gate read `auth_type` and skipped recovery here, 401'ing every
@@ -875,7 +890,7 @@ async fn set_session_model_invalidates_byok_memo_for_same_model_id() {
header_injector: None, header_injector: None,
}; };
let _ = actor let _ = actor
.handle_set_session_model(cfg, false, false, true, 85) .handle_set_session_model(cfg, None, false, false, true, 85)
.await; .await;
assert!( assert!(
@@ -109,6 +109,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
session_info, session_info,
auth_method_id: test_auth_method_id("test-auth"), auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
attribution_callback: None, attribution_callback: None,
auth_manager: None, auth_manager: None,
state: TokioMutex::new(State { state: TokioMutex::new(State {
@@ -562,6 +563,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
session_info: session_info.clone(), session_info: session_info.clone(),
auth_method_id: test_auth_method_id("test-auth"), auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
attribution_callback: None, attribution_callback: None,
auth_manager: None, auth_manager: None,
state: TokioMutex::new(State { state: TokioMutex::new(State {
@@ -825,6 +827,7 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
}, },
auth_method_id: test_auth_method_id("test-auth"), auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
attribution_callback: None, attribution_callback: None,
auth_manager: None, auth_manager: None,
state, state,
@@ -1817,6 +1820,7 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
}, },
auth_method_id: test_auth_method_id("test-auth"), auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
attribution_callback: None, attribution_callback: None,
auth_manager: None, auth_manager: None,
state, state,
@@ -128,6 +128,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
attribution_callback: None, attribution_callback: None,
auth_method_id: test_auth_method_id("cached_token"), auth_method_id: test_auth_method_id("cached_token"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
auth_manager: { auth_manager: {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new( let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
@@ -72,6 +72,7 @@ async fn create_test_actor(
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(), rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
auth_method_id: test_auth_method_id("test-auth"), auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
attribution_callback: None, attribution_callback: None,
auth_manager: None, auth_manager: None,
state, state,
@@ -511,6 +512,7 @@ async fn create_test_actor_with_memory(
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(), rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
auth_method_id: test_auth_method_id("test-auth"), auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
attribution_callback: None, attribution_callback: None,
auth_manager: None, auth_manager: None,
state, state,
@@ -1266,6 +1268,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(), rebuild_spec: crate::session::agent_rebuild::test_rebuild_spec_default(),
auth_method_id: test_auth_method_id("cached_token"), auth_method_id: test_auth_method_id("cached_token"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
auth_manager: { auth_manager: {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new( let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
@@ -127,6 +127,7 @@ async fn create_test_actor_with_memory(
}, },
auth_method_id: test_auth_method_id("test-auth"), auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
attribution_callback: None, attribution_callback: None,
auth_manager: None, auth_manager: None,
state, state,
@@ -78,6 +78,7 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
}, },
auth_method_id: test_auth_method_id("test-auth"), auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
attribution_callback: None, attribution_callback: None,
auth_manager: None, auth_manager: None,
state, state,
@@ -8,8 +8,38 @@ use super::session_bearer_leak_tests::{
}; };
use super::*; use super::*;
use kigi_sampler::BearerResolver; use kigi_sampler::BearerResolver;
use kigi_test_support::EnvGuard;
use std::sync::Arc; use std::sync::Arc;
/// The host BOTH halves of the `anthropic` / `claude-pro-max` collision route
/// to, derived from the registry (as the sibling at
/// [`oauth_platform_models_keep_a_live_resolver_from_their_own_pool`] does) so
/// the fixture cannot drift, with the twin agreement asserted rather than
/// assumed — the collision is only a collision because both platforms serve the
/// same host.
fn anthropic_collision_host() -> String {
let oauth_host = kigi_models::PlatformId::ClaudeProMax.base_url();
assert_eq!(
kigi_models::PlatformId::Anthropic.base_url(),
oauth_host,
"the API-key platform and its subscription-OAuth twin must serve the same host, \
or this fixture is not testing the dual-credential collision"
);
oauth_host
}
/// Ambient BYOK env unset. `resolve_model_auth_facts` probes `std::env::var` at
/// call time, so a developer (or CI) holding `ANTHROPIC_API_KEY` flips the
/// fixture to `Byok` and switches off the session-token gate for a reason that
/// has nothing to do with the platform lookup under test. Every holder must be
/// `#[serial]`.
fn anthropic_collision_env_guard() -> [EnvGuard; 2] {
[
EnvGuard::unset("ANTHROPIC_API_KEY"),
EnvGuard::unset("KIGI_CODE_BASE_URL"),
]
}
#[tokio::test(flavor = "current_thread")] #[tokio::test(flavor = "current_thread")]
async fn dual_credential_slug_collision_resolves_the_selected_oauth_platform() { async fn dual_credential_slug_collision_resolves_the_selected_oauth_platform() {
let local = tokio::task::LocalSet::new(); let local = tokio::task::LocalSet::new();
@@ -127,28 +157,21 @@ async fn oauth_platform_models_keep_a_live_resolver_from_their_own_pool() {
let local = tokio::task::LocalSet::new(); let local = tokio::task::LocalSet::new();
local local
.run_until(async { .run_until(async {
for (catalog_key, slug, base_url) in [ // The base URL is the platform's OWN registry host, exactly as
( // `models_fetch::platform_fetch_base` builds every fetched entry —
"claude-pro-max/claude-opus-4-8", // derived here rather than hard-coded so the fixture cannot drift
"claude-opus-4-8", // from the registry (L10 compares against precisely this).
"https://api.anthropic.com/v1", for (catalog_key, slug) in [
), ("claude-pro-max/claude-opus-4-8", "claude-opus-4-8"),
( ("github-copilot/gpt-4.1", "gpt-4.1"),
"github-copilot/gpt-4.1", ("xai-grok/grok-4-latest", "grok-4-latest"),
"gpt-4.1", ("openai-codex/gpt-5.5", "gpt-5.5"),
"https://api.githubcopilot.com",
),
(
"xai-grok/grok-4-latest",
"grok-4-latest",
"https://api.x.ai/v1",
),
(
"openai-codex/gpt-5.5",
"gpt-5.5",
"https://chatgpt.com/backend-api/codex",
),
] { ] {
let base_url = kigi_models::parse_managed_model_key(catalog_key)
.expect("managed key")
.0
.base_url();
let base_url = base_url.as_str();
let (_dir, actor, _rx) = let (_dir, actor, _rx) =
actor_on_managed_model(catalog_key, slug, base_url, "unused").await; actor_on_managed_model(catalog_key, slug, base_url, "unused").await;
let cfg = actor.reconstruct_full_config().await; let cfg = actor.reconstruct_full_config().await;
@@ -165,12 +188,13 @@ async fn oauth_platform_models_keep_a_live_resolver_from_their_own_pool() {
// token rotated inside the pool is observed by the // token rotated inside the pool is observed by the
// already-built resolver (this is what mid-session refresh // already-built resolver (this is what mid-session refresh
// does). The pool is read here, never mutated. // does). The pool is read here, never mutated.
let pooled = crate::auth::oauth_registry::manager_for_model( let pooled = actor
&crate::auth::oauth_registry::pool_home(), .credential_authority()
catalog_key, .manager_for(
actor.auth_manager.as_ref(), kigi_models::parse_managed_model_key(catalog_key).map(|(p, _)| p),
) base_url,
.expect("an OAuth platform always resolves a manager"); )
.expect("an OAuth platform on its own host always resolves a manager");
assert!( assert!(
!Arc::ptr_eq( !Arc::ptr_eq(
&pooled, &pooled,
@@ -189,25 +213,30 @@ async fn oauth_platform_models_keep_a_live_resolver_from_their_own_pool() {
} }
/// H3/H4 — the stamped AUX paths (image-describe, the auto-mode classifier and /// H3/H4 — the stamped AUX paths (image-describe, the auto-mode classifier and
/// the session-summary client all funnel through `aux_bearer_resolver`). /// the session-summary client all funnel through
/// `stamp_session_local_sampler_fields` copies the SESSION model's resolver onto /// `CredentialAuthority::bearer_resolver_for`). `SamplingClient::post` REPLACES
/// every aux config and `SamplingClient::post` REPLACES the request's auth /// the request's auth header from the resolver, so an API-key-platform aux model
/// header from it, so an API-key-platform aux model would have its own key /// would have its own key overwritten by the Kimi bearer ON THE AUX HOST.
/// overwritten by the Kimi bearer ON THE AUX HOST.
/// ///
/// Revert-to-red: returning `stamped` unconditionally (the pre-fix /// Revert-to-red: make `CredentialAuthority::governing_manager`'s
/// "re-point only when the aux model is OAuth" shape) makes the deepseek / /// `Some(platform) => None` arm return `self.primary.clone()` and the deepseek /
/// openai / `[model.*]`-on-openai.com rows resolve `KIMI_TOKEN`. /// openai / moonshot rows resolve `KIMI_TOKEN`.
#[test] #[test]
fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host() { fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host() {
#[derive(Debug)] let dir = tempfile::tempdir().expect("tempdir");
struct Fixed(&'static str); let primary = Arc::new(crate::auth::AuthManager::new(
impl BearerResolver for Fixed { dir.path(),
fn current_bearer(&self) -> Option<String> { crate::auth::KimiCodeConfig::default(),
Some(self.0.to_string()) ));
} primary.hot_swap(crate::auth::KimiAuth {
} key: KIMI_TOKEN.to_string(),
let stamped: kigi_sampler::SharedBearerResolver = Arc::new(Fixed(KIMI_TOKEN)); auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
let authority = crate::auth::credential_authority::CredentialAuthority::new(
crate::agent::config::EndpointsConfig::default(),
Some(primary),
);
let platform = let platform =
|key: &str| kigi_models::parse_managed_model_key(key).map(|(platform, _)| platform); |key: &str| kigi_models::parse_managed_model_key(key).map(|(platform, _)| platform);
@@ -222,22 +251,16 @@ fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host()
), ),
] { ] {
assert!( assert!(
crate::session::acp_session::sampler_turn::aux_bearer_resolver( authority
Some(stamped.clone()), .bearer_resolver_for(platform(key), base_url)
platform(key), .is_none(),
base_url,
)
.is_none(),
"LEAK: an aux model on {base_url} must not inherit the session bearer resolver" "LEAK: an aux model on {base_url} must not inherit the session bearer resolver"
); );
} }
assert!( assert!(
crate::session::acp_session::sampler_turn::aux_bearer_resolver( authority
Some(stamped.clone()), .bearer_resolver_for(None, "https://api.openai.com/v1")
None, .is_none(),
"https://api.openai.com/v1",
)
.is_none(),
"LEAK: a [model.*] aux model on a third-party host must not inherit it either" "LEAK: a [model.*] aux model on a third-party host must not inherit it either"
); );
@@ -251,17 +274,15 @@ fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host()
(None, kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url), (None, kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url),
(None, "http://127.0.0.1:4141/v1"), (None, "http://127.0.0.1:4141/v1"),
] { ] {
let resolved = crate::session::acp_session::sampler_turn::aux_bearer_resolver( let resolved = authority
Some(stamped.clone()), .bearer_resolver_for(key.and_then(platform), base_url)
key.and_then(platform), .unwrap_or_else(|| panic!("{key:?} @ {base_url} must keep the session resolver"));
base_url,
)
.unwrap_or_else(|| panic!("{key:?} @ {base_url} must keep the session resolver"));
assert_eq!(resolved.current_bearer(), Some(KIMI_TOKEN.to_string())); assert_eq!(resolved.current_bearer(), Some(KIMI_TOKEN.to_string()));
} }
// Re-pointed: an OAuth aux model gets a LIVE resolver over its OWN pool // Re-pointed: an OAuth aux model on ITS OWN host gets a LIVE resolver over
// (empty here), never the stamped Kimi one. // its own pool (empty here), never the Kimi primary. L10: the same model
// redirected to a third-party host gets NOTHING.
let rt = tokio::runtime::Builder::new_current_thread() let rt = tokio::runtime::Builder::new_current_thread()
.enable_all() .enable_all()
.build() .build()
@@ -273,17 +294,697 @@ fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host()
"github-copilot/gpt-4.1", "github-copilot/gpt-4.1",
"openai-codex/gpt-5.5", "openai-codex/gpt-5.5",
] { ] {
let resolved = crate::session::acp_session::sampler_turn::aux_bearer_resolver( let p = platform(key).expect("managed key");
Some(stamped.clone()), let resolved = authority
platform(key), .bearer_resolver_for(Some(p), &p.base_url())
"https://example.invalid/v1", .unwrap_or_else(|| panic!("{key} must keep a live resolver from its own pool"));
)
.unwrap_or_else(|| panic!("{key} must keep a live resolver from its own pool"));
assert_ne!( assert_ne!(
resolved.current_bearer(), resolved.current_bearer(),
Some(KIMI_TOKEN.to_string()), Some(KIMI_TOKEN.to_string()),
"{key}: the aux resolver must never resolve the Kimi session bearer" "{key}: the aux resolver must never resolve the Kimi session bearer"
); );
assert!(
authority
.bearer_resolver_for(Some(p), "https://example.invalid/v1")
.is_none(),
"LEAK ({key}): an OAuth aux model redirected off its own host gets nothing"
);
} }
}); });
} }
/// H4 — TWO CONCURRENT SESSIONS on a colliding slug. The model→platform lookup
/// used to key on `ModelsManager::current_model_id()`, a single PROCESS-GLOBAL
/// `RwLock<acp::ModelId>` written by whichever session switched last. With one
/// session on `xai-grok/grok-4.5` and another on `xai/grok-4.5` — same routing
/// slug, by design — the loser resolved the OTHER session's platform:
/// the subscription session lost its live resolver (unrecoverable 401 ~1h in)
/// and the API-key session got the pooled OAuth bearer stamped over its own
/// `sk-…` key, which the provider rejects.
///
/// Here the global cell is deliberately set to the API-key twin for BOTH
/// sessions (last writer wins, and it was the API-key one). Each session must
/// still resolve ITS OWN selection.
///
/// Revert-to-red: replace `self.selected_catalog_key()` in
/// `SessionActor::model_platform` with
/// `Some(self.models_manager.current_model_id().0.as_ref())`. Under THIS
/// fixture both of the subscription session's assertions fail. (L: the fixture
/// is what makes that true — `managed_entry` carries `api_key: None` and
/// `actor_with_catalog` pins `NotByok`, which is exactly what a FETCHED registry
/// entry resolves to. A user who additionally sets `ANTHROPIC_API_KEY` /
/// `[model.*] env_key` classifies `Byok`, the gate is inactive for that reason
/// alone, and only the `anthropic_oauth` assertion would still catch the
/// mis-resolution — hence the env guard below.)
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn concurrent_sessions_on_a_colliding_slug_each_resolve_their_own_platform() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let api_key_twin = "anthropic/claude-opus-4-8";
let oauth_twin = "claude-pro-max/claude-opus-4-8";
let slug = "claude-opus-4-8";
let host = &anthropic_collision_host();
let catalog = || {
vec![
// API-key platform FIRST, exactly as `PlatformId::ALL` orders them.
managed_entry(api_key_twin, slug, host),
managed_entry(oauth_twin, slug, host),
]
};
let (_d1, subscription, _r1) =
actor_with_catalog(catalog(), oauth_twin, "unused").await;
let (_d2, api_key, _r2) =
actor_with_catalog(catalog(), api_key_twin, "sk-ant-user").await;
// The other session switched last: the process-global cell now names
// the API-key twin for BOTH.
for actor in [&subscription, &api_key] {
actor
.models_manager
.set_current_model_id(acp::ModelId::new(api_key_twin.to_string()));
}
let sub_cfg = subscription.reconstruct_full_config().await;
let resolver = sub_cfg.bearer_resolver.as_ref().expect(
"the subscription session must keep its own live bearer_resolver even when \
another session switched the process-global model last",
);
assert_ne!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"it must read the claude-pro-max pool, never the Kimi primary"
);
assert!(
sub_cfg.anthropic_oauth,
"the Claude OAuth Messages adaptation must follow the SUBSCRIPTION session"
);
let api_cfg = api_key.reconstruct_full_config().await;
assert!(
api_cfg.bearer_resolver.is_none(),
"LEAK: the API-key session must get no session bearer_resolver"
);
assert!(
!api_cfg.anthropic_oauth,
"the API-key session must not get the OAuth Messages adaptation"
);
assert_eq!(
api_cfg.api_key.as_deref(),
Some("sk-ant-user"),
"the API-key session keeps its own provider key"
);
})
.await;
}
/// H4 in LEADER mode, where `agent/handlers/model_switch.rs` skips
/// `set_current_model_id` ENTIRELY, so the process-global cell is frozen at the
/// startup default for the whole process lifetime. `platform_for_slug` then
/// fell through to the `.rev()` scan, which returns the LAST match — the OAuth
/// twin — so a Leader-mode session on the API-KEY twin was handed the pooled
/// OAuth bearer plus the Messages adaptation, and Anthropic rejects both.
///
/// Revert-to-red: same edit as above; with the global cell naming the startup
/// default (not in this catalog) the `.rev()` fallback resolves
/// `claude-pro-max/*` and both assertions fail.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn leader_mode_session_resolves_its_own_platform_without_the_global_cell() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let slug = "claude-opus-4-8";
let host = &anthropic_collision_host();
let (_dir, actor, _rx) = actor_with_catalog(
vec![
managed_entry("anthropic/claude-opus-4-8", slug, host),
managed_entry("claude-pro-max/claude-opus-4-8", slug, host),
],
"anthropic/claude-opus-4-8",
"sk-ant-user",
)
.await;
// Leader mode never writes the global cell: it still names the
// startup default, which is not in this catalog at all.
assert!(
!actor
.models_manager
.models()
.contains_key(actor.models_manager.current_model_id().0.as_ref()),
"precondition: the process-global model id is stale (Leader mode)"
);
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_none(),
"LEAK: a Leader-mode API-key session must get no session bearer_resolver"
);
assert!(
!cfg.anthropic_oauth,
"a Leader-mode API-key session must not get the OAuth Messages adaptation"
);
})
.await;
}
/// L13 — a user whose ACP auth method is an API-KEY registry platform (e.g.
/// `deepseek`) can still SELECT a subscription-OAuth model, and the chokepoint
/// hands it that platform's pooled bearer as the request's `api_key`. The gate
/// keys on the primary method, which is not session-based, so the config used
/// to carry NO `bearer_resolver`: the pooled token froze at selection time and
/// the session died with an unrecoverable 401 once it expired (~1h).
///
/// The model's own credential now makes the gate session-based, confined to
/// that platform's own host by the gate's `credential_class` conjunct.
#[tokio::test(flavor = "current_thread")]
async fn oauth_model_under_an_api_key_auth_method_keeps_its_pooled_resolver() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let base_url = kigi_models::PlatformId::ClaudeProMax.base_url();
let (_dir, actor, _rx) = actor_with_catalog(
vec![managed_entry(
"claude-pro-max/claude-opus-4-8",
"claude-opus-4-8",
&base_url,
)],
"claude-pro-max/claude-opus-4-8",
"unused",
)
.await;
// The PRIMARY ACP method is an API-key registry platform login.
actor
.auth_method_id
.store(Some(std::sync::Arc::new(acp::AuthMethodId::new(
"deepseek",
))));
let cfg = actor.reconstruct_full_config().await;
let resolver = cfg.bearer_resolver.as_ref().expect(
"a subscription-OAuth model keeps a live resolver whatever the primary \
ACP auth method is, or it cannot refresh mid-session",
);
assert_ne!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"and it reads the claude-pro-max pool, never the primary"
);
// An API-key-platform model under the same method stays resolver-free.
let (_d2, deepseek, _r2) = actor_with_catalog(
vec![managed_entry(
"deepseek/deepseek-chat",
"deepseek-chat",
"https://api.deepseek.com/v1",
)],
"deepseek/deepseek-chat",
"sk-deepseek",
)
.await;
deepseek
.auth_method_id
.store(Some(std::sync::Arc::new(acp::AuthMethodId::new(
"deepseek",
))));
assert!(
deepseek
.reconstruct_full_config()
.await
.bearer_resolver
.is_none(),
"LEAK: an API-key-platform model must never get a session resolver"
);
})
.await;
}
/// H-b — a `None` or STALE per-session catalog key must REFUSE, not degrade to
/// the subscription-OAuth twin.
///
/// `model_platform` falls through to `resolve_catalog_key`'s `.rev()` scan when
/// the session's own key does not name the slug, and that scan returns the LAST
/// match — the OAuth twin, because `PlatformId::ALL` orders every API-key
/// platform first. Combined with the L13 disjunct (a model whose own credential
/// is a pooled OAuth session is session-based BY ITSELF), an API-KEY session on
/// `anthropic/claude-opus-4-8` with no per-session key got
/// `is_session_based = true`, `credential_class = Pooled` (same
/// host) and, at `NotByok`, an ACTIVE gate — so `manager_for` handed it the
/// Claude POOLED manager, whose `bearer_resolver` REPLACES the user's own
/// `sk-ant-…` on the wire, plus the OAuth Messages adaptation. Anthropic rejects
/// both. This is exactly what H4 prevents, reached through the `None` path.
///
/// A key that is absent or names a different model is not evidence for either
/// twin: resolve to NO platform, which the chokepoint then decides purely by the
/// ENDPOINT (the OAuth host is not this session's coding endpoint ⇒ nothing
/// rides).
///
/// Revert-to-red (production, compiles): delete the
/// `platform.oauth().is_some() && !disambiguated && slug_collides_across_platforms(..)`
/// refusal from `crate::agent::models::platform_for_slug` and every
/// `bearer_resolver` / `anthropic_oauth` assertion below fails.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn a_missing_or_stale_session_key_refuses_instead_of_guessing_the_oauth_twin() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let api_key_twin = "anthropic/claude-opus-4-8";
let slug = "claude-opus-4-8";
let host = anthropic_collision_host();
let catalog = || {
vec![
// API-key platform FIRST, exactly as `PlatformId::ALL` orders them.
managed_entry(api_key_twin, slug, &host),
managed_entry("claude-pro-max/claude-opus-4-8", slug, &host),
]
};
for (case, stale_key) in [
// No key at all: a session spawned on a model that left the
// catalog, or one an older build never seeded.
("absent", None),
// Stale: an `OverrideModelName` rename, or a key naming a model
// this session is no longer on.
("stale", Some("claude-pro-max/some-other-model".to_string())),
] {
let (_dir, actor, _rx) =
actor_with_catalog(catalog(), api_key_twin, "sk-ant-user").await;
*actor.selected_catalog_key.borrow_mut() = stale_key;
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_none(),
"{case}: LEAK — an unresolvable selection must get NO session bearer \
resolver; the pooled OAuth bearer would REPLACE the user's own key"
);
assert!(
!cfg.anthropic_oauth,
"{case}: nor the Claude OAuth Messages adaptation"
);
assert_eq!(
cfg.api_key.as_deref(),
Some("sk-ant-user"),
"{case}: the user's own provider key must survive untouched"
);
assert!(
actor
.credential_authority()
.manager_for(
crate::agent::models::platform_for_slug(
&actor.models_manager.models(),
actor.selected_catalog_key().as_deref(),
slug,
),
&host,
)
.is_none(),
"{case}: and no manager either — refuse, never guess"
);
}
// …while a session that DID select the OAuth twin still gets its
// pooled resolver: the refusal is about the guess, not the platform.
let (_dir, selected, _rx) =
actor_with_catalog(catalog(), "claude-pro-max/claude-opus-4-8", "unused").await;
let cfg = selected.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_some() && cfg.anthropic_oauth,
"a DELIBERATE subscription selection keeps its pooled resolver and \
adaptation (this is what makes the refusals above meaningful)"
);
})
.await;
}
/// H-c — coverage for the FIRST of the two production writers of
/// `selected_catalog_key`: the spawn seed
/// (`crate::agent::models::selected_catalog_key_for_spawn`, called from
/// `spawn.rs`). Every other test in this module sets the field by hand, so a
/// wrong seed was silent.
///
/// Both spawn shapes are covered: a FRESH session, spawned on the catalog key
/// the picker resolved, and a RESUME/LOAD, which spawns with the RAW persisted
/// `summary.current_model_id` — a BARE routing slug after any `SetSessionModel`,
/// since `handle_set_session_model` persists `sampling_config.model`. This seed
/// is where that slug becomes a key. The assertion is end-to-end: the seeded key
/// is fed to the very function the auth layer keys on.
#[test]
fn spawn_seeds_the_session_key_the_auth_layer_keys_on() {
let slug = "claude-opus-4-8";
let host = anthropic_collision_host();
let api_key_twin = "anthropic/claude-opus-4-8";
let oauth_twin = "claude-pro-max/claude-opus-4-8";
let models: indexmap::IndexMap<String, crate::agent::config::ModelEntry> = [
managed_entry(api_key_twin, slug, &host),
managed_entry(oauth_twin, slug, &host),
]
.into_iter()
.collect();
// FRESH: spawned on the catalog key. Idempotent, and it disambiguates.
for selected in [api_key_twin, oauth_twin] {
let seeded = crate::agent::models::selected_catalog_key_for_spawn(
&models,
&acp::ModelId::new(selected.to_string()),
);
assert_eq!(
seeded.as_deref(),
Some(selected),
"a fresh session must record the catalog key it was spawned with"
);
assert_eq!(
crate::agent::models::platform_for_slug(&models, seeded.as_deref(), slug),
kigi_models::parse_managed_model_key(selected).map(|(p, _)| p),
"…and that key must resolve THIS session's own platform for the bare slug"
);
}
// RESUME/LOAD: `acp_agent::load_session` spawns with the RAW persisted id,
// which after any model switch is the bare routing slug. THIS seed is what
// turns it into a key — the picker's `.rev()` answer, which is the resume
// default for a collided slug.
let resumed = crate::agent::models::selected_catalog_key_for_spawn(
&models,
&acp::ModelId::new(slug.to_string()),
);
assert_eq!(
resumed.as_deref(),
Some(oauth_twin),
"a bare persisted slug resolves through the picker's own lookup"
);
// A model that is no longer in the catalog seeds NOTHING, which (H-b) then
// refuses rather than guessing a twin.
assert_eq!(
crate::agent::models::selected_catalog_key_for_spawn(
&models,
&acp::ModelId::new("gone/model".to_string()),
),
None,
"a model that left the catalog must not seed a key"
);
}
/// H-c — coverage for the SECOND production writer: `SetSessionModel`
/// (`handle_set_session_model`), the picker's own path. The existing test
/// through this handler passes `None`, so a handler that dropped the key on the
/// floor stayed green.
///
/// End-to-end: after the switch the session's per-turn config must carry the
/// SELECTED twin's pooled resolver and adaptation, even though the bare slug in
/// the config is ambiguous.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn set_session_model_records_the_key_the_next_turn_resolves_on() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let slug = "claude-opus-4-8";
let host = anthropic_collision_host();
let api_key_twin = "anthropic/claude-opus-4-8";
let oauth_twin = "claude-pro-max/claude-opus-4-8";
let (_dir, actor, _rx) = actor_with_catalog(
vec![
managed_entry(api_key_twin, slug, &host),
managed_entry(oauth_twin, slug, &host),
],
api_key_twin,
"sk-ant-user",
)
.await;
// Switch to the SUBSCRIPTION twin, exactly as
// `agent/handlers/model_switch.rs` does: the ambiguous slug in the
// sampler config plus the catalog KEY the picker resolved.
let models = actor.models_manager.models();
let entry = models.get(oauth_twin).expect("catalog entry");
let sampler = crate::agent::config::sampling_config_for_model(
entry,
crate::agent::config::resolve_credentials(entry, None),
None,
);
actor
.handle_set_session_model(
sampler,
Some(oauth_twin.to_string()),
false,
false,
true,
85,
)
.await
.expect("model switch");
assert_eq!(
actor.selected_catalog_key().as_deref(),
Some(oauth_twin),
"SetSessionModel must record the picker's catalog key"
);
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_some(),
"the switched-to subscription model must keep a live pooled resolver"
);
assert!(
cfg.anthropic_oauth,
"…and the Claude OAuth Messages adaptation"
);
// And back: switching to the API-key twin must UNDO both.
let entry = models.get(api_key_twin).expect("catalog entry");
let sampler = crate::agent::config::sampling_config_for_model(
entry,
crate::agent::config::resolve_credentials(entry, None),
None,
);
actor
.handle_set_session_model(
sampler,
Some(api_key_twin.to_string()),
false,
false,
true,
85,
)
.await
.expect("model switch");
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_none() && !cfg.anthropic_oauth,
"LEAK: switching back to the API-key twin must drop the pooled resolver \
and the OAuth adaptation"
);
})
.await;
}
/// H-c — `OverrideModelName` is the one command that rewrites
/// `SamplingConfig::model` WITHOUT going through `SetSessionModel`, so it used
/// to leave `selected_catalog_key` naming a model the session is no longer on.
/// It must keep the field consistent: KEEP when the key still names the new
/// routing name, CLEAR otherwise — never re-resolve, which would put the
/// `.rev()` guess into the field the rule treats as a deliberate selection.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn override_model_name_keeps_the_session_key_consistent() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let slug = "claude-opus-4-8";
let host = anthropic_collision_host();
let oauth_twin = "claude-pro-max/claude-opus-4-8";
let (_dir, actor, _rx) = actor_with_catalog(
vec![
managed_entry("anthropic/claude-opus-4-8", slug, &host),
managed_entry(oauth_twin, slug, &host),
],
oauth_twin,
"unused",
)
.await;
// A rename to the SAME model's routing slug (or to its catalog key)
// keeps the selection.
for same in [slug, oauth_twin] {
actor.retain_selected_catalog_key_for(same);
assert_eq!(
actor.selected_catalog_key().as_deref(),
Some(oauth_twin),
"{same}: still names the selected entry — keep it"
);
}
// A rename to a DIFFERENT name makes the key stale: clear it, so the
// collided slug refuses (H-b) instead of resolving the old model.
actor.retain_selected_catalog_key_for("some-harness-model-name");
assert_eq!(
actor.selected_catalog_key(),
None,
"a stale key must be cleared, not carried into the next turn's \
platform lookup"
);
})
.await;
}
/// M3 — the FIRST-PARTY aux case must honour the session gate, which is what
/// the old shape did implicitly.
///
/// `stamp_session_local_sampler_fields` used to copy
/// `active_session_config.bearer_resolver`, and that field is `None` whenever
/// the gate is inactive. Re-pointing the aux resolver at the chokepoint (the
/// LEAK 1b fix) made it `Some(primary)` for the session's own coding endpoint
/// REGARDLESS of the gate — so a BYOK / api-key session with a `[model.*]` aux
/// entry carrying its own key on that endpoint had that key REPLACED by the
/// primary bearer on every image-describe / auto-mode-classifier / summary
/// request (`SamplingClient::post` overrides the auth header from the resolver).
///
/// Revert-to-red (production, compiles): delete the
/// `if is_primary_channel && !SessionTokenAuthGate::new(…).active()` early
/// return from `sampler_turn::aux_bearer_resolver_for` and the first two rows
/// below resolve `KIMI_TOKEN`.
#[tokio::test(flavor = "current_thread")]
async fn first_party_aux_resolver_honours_the_session_gate() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let coding_host = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
let aux_slug = "kigi-aux";
let mut info = crate::agent::config::ModelInfo::fallback(aux_slug);
info.id = None; // a `[model.kigi-aux]` block, not a registry entry
info.base_url = coding_host.to_string();
let aux_entry = crate::agent::config::ModelEntry {
info,
api_key: None,
env_key: None,
api_base_url: None,
};
// (case, ACP auth method, the aux model's own BYOK status, expected)
for (case, auth_method, byok, expect_resolver) in [
(
"an API-key session: the aux model's own key must survive",
"deepseek",
crate::agent::auth_method::ModelByok::NotByok,
false,
),
(
"a BYOK aux entry under a session method: its env_key wins",
"cached_token",
crate::agent::auth_method::ModelByok::Byok,
false,
),
(
"the first-party subscription aux channel: byte-identical",
"cached_token",
crate::agent::auth_method::ModelByok::NotByok,
true,
),
] {
let (_dir, actor, _rx) = actor_with_catalog(
vec![(aux_slug.to_string(), aux_entry.clone())],
aux_slug,
"",
)
.await;
actor
.auth_method_id
.store(Some(Arc::new(acp::AuthMethodId::new(auth_method))));
actor.model_auth_facts.replace(Some((
aux_slug.to_string(),
crate::agent::config::ModelAuthFacts {
byok,
auth_scheme: Default::default(),
},
)));
let resolved = actor.aux_bearer_resolver(aux_slug, coding_host);
assert_eq!(
resolved.is_some(),
expect_resolver,
"{case}: aux resolver presence on the session's own endpoint"
);
if let Some(resolver) = resolved {
assert_eq!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"{case}: and when it IS kept it is the primary's, live"
);
}
}
})
.await;
}
/// M-aux (REGRESSION this remediation introduced) — an aux call must NOT evict
/// the SESSION model's memoized auth facts.
///
/// `SessionActor::model_auth_facts` is a SINGLE slot. When `aux_bearer_resolver`
/// began asking it about the AUX slug, a definite result overwrote the session
/// model's entry, and:
/// (a) the next `reconstruct_full_config` re-paid `load_effective_config()` +
/// `resolve_model_list()` — the per-turn disk read M7/M9 removed — on top
/// of the one the aux call itself paid; and
/// (b) the memo's documented purpose (a transient `Unknown` falling back to
/// the last DEFINITE value FOR THE SAME model_id) was defeated: with the
/// aux slug in the slot, the session model's `Unknown` degrades to
/// `endpoint_is_first_party`, which is `false` for every
/// subscription-OAuth host — the session loses its `bearer_resolver` and
/// 401s unrecoverably ~1h in, the failure L13 exists to prevent.
///
/// Round 3's deleted `repoint_aux_bearer_resolver` never touched the memo.
///
/// Revert-to-red (production, compiles): make `SessionActor::aux_bearer_resolver`
/// call `self.model_auth_facts(slug)` instead of `self.aux_model_auth_facts(slug)`
/// — the slot then names the aux slug and both assertions below fail.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn an_aux_call_does_not_evict_the_session_models_auth_facts() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let session_slug = "claude-opus-4-8";
let host = anthropic_collision_host();
let oauth_twin = "claude-pro-max/claude-opus-4-8";
let (_dir, actor, _rx) = actor_with_catalog(
vec![managed_entry(oauth_twin, session_slug, &host)],
oauth_twin,
"unused",
)
.await;
// The session model's DEFINITE facts, as a turn would have memoized
// them.
actor.model_auth_facts.replace(Some((
session_slug.to_string(),
crate::agent::config::ModelAuthFacts {
byok: crate::agent::auth_method::ModelByok::NotByok,
auth_scheme: Default::default(),
},
)));
// An aux turn: the auto-mode classifier / image-describe slug, which
// is NOT the session's model.
let _ = actor.aux_bearer_resolver("kigi-aux-classifier", &host);
let memo = actor.model_auth_facts.borrow();
let (cached_id, facts) = memo
.as_ref()
.expect("the session model's memo must survive an aux call");
assert_eq!(
cached_id, session_slug,
"an aux call evicted the SESSION model's memo: the next turn re-reads \
config from disk, and a transient Unknown loses its definite fallback"
);
assert_eq!(facts.byok, crate::agent::auth_method::ModelByok::NotByok);
})
.await;
}
@@ -5,7 +5,7 @@
//! `kimi-code` / any OAuth platform) + a selected API-key-platform model //! `kimi-code` / any OAuth platform) + a selected API-key-platform model
//! classifies `ModelByok::NotByok` (the model carries no `[model.*]` key), the //! classifies `ModelByok::NotByok` (the model carries no `[model.*]` key), the
//! pre-fix `session_token_auth_gate` returned `true` unconditionally on that //! pre-fix `session_token_auth_gate` returned `true` unconditionally on that
//! arm, `auth_manager_for_model` fell through to the primary Kimi manager for a //! arm, the manager lookup fell through to the primary Kimi manager for a
//! non-OAuth platform, and `SamplingClient::post` then REPLACED the correctly //! non-OAuth platform, and `SamplingClient::post` then REPLACED the correctly
//! resolved provider key with the Kimi bearer on the wire. //! resolved provider key with the Kimi bearer on the wire.
//! //!
@@ -21,11 +21,12 @@
//! `bearer_resolver` drawn from their OWN pooled `AuthManager`, or they lose //! `bearer_resolver` drawn from their OWN pooled `AuthManager`, or they lose
//! mid-session token refresh. //! mid-session token refresh.
//! //!
//! STORAGE DISCIPLINE (H6): nothing here touches the developer's real `~/.kigi` //! STORAGE DISCIPLINE (H6/M8): nothing here touches the developer's real
//! and nothing hot-swaps the process-global OAuth pool. Under `cfg(test)` //! `~/.kigi` and nothing hot-swaps the process-global OAuth pool. Under
//! `oauth_registry::pool_home()` is a process-lifetime `TempDir`, so every //! `cfg(test)` `oauth_registry::pool_home()` is a per-process temp path that is
//! pooled manager is empty — which is exactly what the assertions need (a live //! never created, so every pooled manager is empty — exactly what the
//! resolver that is provably NOT the Kimi one). //! assertions need (a live resolver that is provably NOT the Kimi one) — and
//! the binary leaves nothing behind.
use super::support::*; use super::support::*;
use super::*; use super::*;
@@ -106,9 +107,14 @@ pub(super) async fn actor_with_catalog(
actor.models_manager.insert_test_entry(key, entry); actor.models_manager.insert_test_entry(key, entry);
} }
let selected_entry = selected_entry.expect("the selected key must be in the catalog"); let selected_entry = selected_entry.expect("the selected key must be in the catalog");
actor // H4: the SESSION owns its selection. The process-global
.models_manager // `ModelsManager::current_model_id()` is deliberately left UNSET (it still
.set_current_model_id(acp::ModelId::new(selected.to_string())); // names the startup default, which is not in this catalog) — exactly what
// Leader mode produces, since `agent/handlers/model_switch.rs` never calls
// `set_current_model_id` there, and what a second concurrent session on a
// colliding slug produces (last writer wins). Every assertion below
// therefore rides the per-session key, not the global cell.
*actor.selected_catalog_key.borrow_mut() = Some(selected.to_string());
let slug = selected_entry.info().model.clone(); let slug = selected_entry.info().model.clone();
actor actor
@@ -166,7 +172,7 @@ pub(super) async fn actor_on_managed_model(
/// (`cached_token`) method with a live Kimi primary must send DeepSeek's own key /// (`cached_token`) method with a live Kimi primary must send DeepSeek's own key
/// — the Kimi subscription bearer must not appear anywhere in the request. /// — the Kimi subscription bearer must not appear anywhere in the request.
/// ///
/// Revert-to-red: dropping `endpoint_takes_session_credential` from /// Revert-to-red: dropping the `credential_class` conjunct from
/// `session_token_auth_gate` puts `Bearer <KIMI_TOKEN>` on this request. /// `session_token_auth_gate` puts `Bearer <KIMI_TOKEN>` on this request.
#[tokio::test(flavor = "multi_thread")] #[tokio::test(flavor = "multi_thread")]
async fn deepseek_turn_under_a_kimi_session_sends_no_kimi_bearer_on_the_wire() { async fn deepseek_turn_under_a_kimi_session_sends_no_kimi_bearer_on_the_wire() {
@@ -282,11 +288,12 @@ async fn api_key_platform_models_get_no_session_bearer_resolver() {
/// C2 at the resolver channel: a `[model.*]` entry has NO platform /// C2 at the resolver channel: a `[model.*]` entry has NO platform
/// (`info.id == None`), which used to be a blanket allow. Pointed at a /// (`info.id == None`), which used to be a blanket allow. Pointed at a
/// third-party host it must get no session resolver; pointed at the session's /// third-party host it must get no session resolver; pointed at the session's
/// own coding endpoint (a `KIGI_CODE_BASE_URL` deployment or a local dev proxy) /// own coding endpoint (a config.toml `[endpoints] coding_api_base_url`
/// it must keep one — that is why the predicate is not `is_first_party_url`. /// deployment, a `KIGI_CODE_BASE_URL` override, or a local dev proxy) it must
/// keep one — that is why the predicate is not `is_first_party_url`.
/// ///
/// Revert-to-red: making the `None` arm of `platform_takes_session_credential` /// Revert-to-red: make `CredentialAuthority::is_session_coding_endpoint` return
/// return `true` again puts a Kimi resolver on the openai.com config. /// `true` unconditionally and a Kimi resolver lands on the openai.com config.
#[tokio::test(flavor = "current_thread")] #[tokio::test(flavor = "current_thread")]
async fn config_model_entry_takes_a_session_resolver_only_on_its_own_endpoint() { async fn config_model_entry_takes_a_session_resolver_only_on_its_own_endpoint() {
let local = tokio::task::LocalSet::new(); let local = tokio::task::LocalSet::new();
@@ -343,21 +350,15 @@ async fn config_model_entry_takes_a_session_resolver_only_on_its_own_endpoint()
.await; .await;
} }
/// H5 — the slug collision. A user holding BOTH an xAI API key and a Grok /// The Kimi / first-party subscription channel must be BYTE-IDENTICAL: the
/// subscription has `xai/grok-4.5` AND `xai-grok/grok-4.5` in one catalog, in /// session model keeps its live bearer_resolver AND the pre-flight refresh
/// `PlatformId::ALL` order (`Xai`(15) before `XaiGrok`(25)) and with the SAME /// still heals a stale buffered key. This is also what proves the Kimi bearer is
/// routing slug. `cfg.model` is that bare slug, so the auth layer used to /// live in every LEAK assertion in this module — it WOULD leak if the guard
/// first-match the API-key entry: no bearer_resolver, no live refresh (the /// were missing.
/// session dies ~1h in with an unrecoverable 401), and — for the Anthropic and
/// Codex twins — the OAuth Messages adaptation and the Codex identity headers
/// silently dropped.
/// ///
/// The catalog KEY the picker selected is now authoritative. /// (L12: the slug-collision commentary that used to sit here belongs to the
/// /// collision tests in `session_bearer_leak_platform_tests`, which is where its
/// Revert-to-red: resolving the platform from `find_model_by_id(models, slug)` /// revert-to-red actually reproduces; on this first-party test it never could.)
/// instead of `current_model_id` resolves `xai/grok-4.5` /
/// `anthropic/claude-opus-4-8` / `openai/gpt-5.5-codex` and every assertion
/// below fails.
#[tokio::test(flavor = "current_thread")] #[tokio::test(flavor = "current_thread")]
async fn kimi_first_party_model_still_rides_the_primary_session_bearer() { async fn kimi_first_party_model_still_rides_the_primary_session_bearer() {
let local = tokio::task::LocalSet::new(); let local = tokio::task::LocalSet::new();
@@ -188,6 +188,7 @@ pub(crate) async fn create_test_actor_ex(
}, },
auth_method_id: test_auth_method_id("test-auth"), auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
attribution_callback: None, attribution_callback: None,
auth_manager: None, auth_manager: None,
state, state,
@@ -145,6 +145,13 @@ pub enum SessionCommand {
}, },
SetSessionModel { SetSessionModel {
sampling_config: kigi_sampler::SamplerConfig, sampling_config: kigi_sampler::SamplerConfig,
/// The catalog KEY the picker resolved (`{platform}/{model}` for a
/// registry model), which `sampling_config.model` — the bare routing
/// slug — cannot express when an API-key platform and its
/// subscription-OAuth twin list the same id. The session stores it as
/// its OWN selection instead of reading the process-global
/// `ModelsManager::current_model_id()` (H4).
catalog_key: Option<String>,
use_concise: bool, use_concise: bool,
/// When `false`, skip the system prompt rewrite (concise/default swap). /// When `false`, skip the system prompt rewrite (concise/default swap).
/// Set to `false` for forked sessions so mid-session model switches /// Set to `false` for forked sessions so mid-session model switches
@@ -2168,6 +2168,7 @@ mod inline_auto_compact_flow_tests {
}, },
auth_method_id: test_auth_method_id("test-auth"), auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None), model_auth_facts: std::cell::RefCell::new(None),
selected_catalog_key: std::cell::RefCell::new(None),
attribution_callback: None, attribution_callback: None,
auth_manager: None, auth_manager: None,
state, state,