This commit is contained in:
2026-07-22 11:08:54 -04:00
parent 422e241e13
commit 2d00a4e6e6
19 changed files with 1955 additions and 296 deletions
@@ -302,14 +302,26 @@ impl ModelByok {
/// buffered token on every turn and 401s with `bad-credentials` until restart.
/// It refreshes when `endpoint_is_first_party` — the request targets the
/// first-party API, where sending the session token cannot leak to a
/// third-party BYOK endpoint. A definite `NotByok` always refreshes (it only
/// ever routes to the session endpoint); a definite `Byok` never does.
/// third-party BYOK endpoint. A definite `NotByok` refreshes (within the
/// session's own providers); a definite `Byok` never does.
///
/// `endpoint_takes_session_credential` ([`platform_takes_session_credential`])
/// is the outer, non-negotiable guard and the reason the `NotByok` arm is safe:
/// it is `false` for every API-key registry platform (deepseek, openai,
/// anthropic, moonshot-*, …), whose models classify `NotByok` (they carry no
/// `[model.*]` key) yet route to a THIRD-PARTY inference host while
/// `oauth_registry::manager_for_model` falls through to the primary (Kimi)
/// manager. Without this term the mainstream "session auth + API-key-platform
/// model" configuration stamps the user's Kimi subscription bearer on every
/// request to that host.
pub fn session_token_auth_gate(
is_session_based_method: bool,
model_byok: ModelByok,
endpoint_is_first_party: bool,
endpoint_takes_session_credential: bool,
) -> bool {
is_session_based_method
&& endpoint_takes_session_credential
&& match model_byok {
ModelByok::NotByok => true,
ModelByok::Byok => false,
@@ -317,6 +329,42 @@ 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 =
"Session expired. Run `kigi login` to re-authenticate.";
@@ -683,15 +731,117 @@ mod tests {
#[test]
fn session_token_auth_gate_matrix() {
// Session method + NotByok → refresh.
assert!(session_token_auth_gate(true, ModelByok::NotByok, false));
// Session method + NotByok → refresh (endpoint takes the session cred).
assert!(session_token_auth_gate(
true,
ModelByok::NotByok,
false,
true
));
// Session method + Byok → never.
assert!(!session_token_auth_gate(true, ModelByok::Byok, true));
assert!(!session_token_auth_gate(true, ModelByok::Byok, true, true));
// Session method + Unknown → only on first-party endpoints.
assert!(session_token_auth_gate(true, ModelByok::Unknown, true));
assert!(!session_token_auth_gate(true, ModelByok::Unknown, false));
assert!(session_token_auth_gate(
true,
ModelByok::Unknown,
true,
true
));
assert!(!session_token_auth_gate(
true,
ModelByok::Unknown,
false,
true
));
// Non-session method → never.
assert!(!session_token_auth_gate(false, ModelByok::NotByok, true));
assert!(!session_token_auth_gate(
false,
ModelByok::NotByok,
true,
true
));
// An endpoint that does not take the session credential (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 first_party in [false, true] {
assert!(
!session_token_auth_gate(true, byok, first_party, false),
"byok={byok:?} first_party={first_party}: an API-key-platform \
endpoint must never receive a session bearer"
);
}
}
}
/// 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).
+58 -22
View File
@@ -2779,13 +2779,20 @@ pub fn default_model_entries(endpoints: &EndpointsConfig) -> IndexMap<String, Mo
}
/// Resolve a model against the available model map.
/// Checks the map key (id) first, then falls back to a slug scan.
///
/// The slug scan takes the LAST match, exactly like the picker's
/// [`crate::agent::models::resolve_catalog_key`]. Duplicate slugs across
/// platforms are by design (`moonshot-cn` and `moonshot-ai` list the same ids;
/// so do every API-key platform and its subscription-OAuth twin), and a
/// first-match scan here made the auth layer resolve a DIFFERENT entry than the
/// one the picker selected — the H5 slug collision. One direction, one answer.
pub fn find_model_by_id<'a>(
models: &'a IndexMap<String, ModelEntry>,
model_id: &str,
) -> Option<&'a ModelEntry> {
models
.get(model_id)
.or_else(|| models.values().find(|m| m.model == model_id))
.or_else(|| models.values().rev().find(|m| m.model == model_id))
}
/// Whether the EFFECTIVE Auto-mode classifier model supports reasoning effort:
/// the model actually routed to (`aux_model` when the aux sampler resolved) else
@@ -4643,14 +4650,22 @@ reasoning_effort = "low"
});
(dir, manager)
}
/// LEAK 1a (aux/summary model): the aux `session_key` is resolved by the aux
/// model's OWN platform (as `build_summary_client` / `resolve_aux_sampler_config`
/// now do). A grok (oauth-platform) aux model's resolved sampler `api_key` is
/// therefore grok's own token or `None` — NEVER the primary Kimi key — while a
/// first-party / non-oauth aux model still gets the primary (byte-identical).
/// LEAK 1a (aux/summary model `api_key` channel): the aux `session_key` is
/// resolved by the aux model's OWN platform AND endpoint (as
/// `build_summary_client` / `resolve_aux_sampler_config` now do), so it is
/// never the primary Kimi key on a host that does not own it.
///
/// - a grok (oauth-platform) aux model → grok's own pooled token or `None`;
/// - a `moonshot-cn` (API-key platform) aux model → NO session key at all
/// (pre-fix it received `kimi-tok`, which `resolve_credentials` then
/// stamped as the `api_key` on an `api.moonshot.cn` request);
/// - a `kimi-code` aux model → the primary, byte-identical.
///
/// Revert-to-red: dropping the `platform_takes_session_credential` term
/// from `session_key_for_endpoint` makes the moonshot assertion see
/// `Some("kimi-tok")`.
#[tokio::test]
async fn aux_model_session_key_is_platform_scoped_never_leaking_kimi() {
let home = tempfile::tempdir().unwrap();
let (_kd, kimi) = kimi_primary("kimi-tok");
let endpoints = EndpointsConfig::default();
// A grok aux catalog entry (managed id → oauth platform), no own key.
@@ -4658,9 +4673,9 @@ reasoning_effort = "low"
grok.info.id = Some("xai-grok/grok-4-latest".to_string());
let mut grok_catalog = IndexMap::new();
grok_catalog.insert("grok".to_string(), grok);
let grok_key = crate::auth::oauth_registry::session_key_for_model(
home.path(),
"xai-grok/grok-4-latest",
let grok_key = crate::auth::oauth_registry::session_key_for_catalog_model(
&grok_catalog,
"grok",
Some(&kimi),
);
let grok_cfg = resolve_aux_model_sampling_config(
@@ -4675,10 +4690,11 @@ reasoning_effort = "low"
Some("kimi-tok"),
"a grok aux model must never receive the primary Kimi session token",
);
// A non-oauth aux catalog entry still resolves to the primary token.
// An API-key registry platform gets NO session key — the pre-fix
// behaviour handed it the primary Kimi token on api.moonshot.cn.
let mut k2 = test_model_entry(
"kimi-k2-0905-preview",
"https://vendor/v1",
"https://api.moonshot.cn/v1",
None,
None,
None,
@@ -4686,23 +4702,43 @@ reasoning_effort = "low"
k2.info.id = Some("moonshot-cn/kimi-k2".to_string());
let mut k2_catalog = IndexMap::new();
k2_catalog.insert("k2".to_string(), k2);
let k2_key = crate::auth::oauth_registry::session_key_for_model(
home.path(),
"moonshot-cn/kimi-k2",
assert_eq!(
crate::auth::oauth_registry::session_key_for_catalog_model(
&k2_catalog,
"k2",
Some(&kimi),
),
None,
"LEAK: an API-key-platform aux model must receive no session token",
);
// The first-party subscription channel is byte-identical.
let mut kimi_code = test_model_entry(
"kimi-for-coding",
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
None,
None,
None,
);
kimi_code.info.id = Some("kimi-code/kimi-for-coding".to_string());
let mut kimi_catalog = IndexMap::new();
kimi_catalog.insert("kimi-for-coding".to_string(), kimi_code);
let kimi_key = crate::auth::oauth_registry::session_key_for_catalog_model(
&kimi_catalog,
"kimi-for-coding",
Some(&kimi),
);
let k2_cfg = resolve_aux_model_sampling_config(
"k2",
&k2_catalog,
let kimi_cfg = resolve_aux_model_sampling_config(
"kimi-for-coding",
&kimi_catalog,
&endpoints,
k2_key.as_deref(),
kimi_key.as_deref(),
None,
)
.expect("non-oauth aux resolves via the primary session token");
.expect("a kimi-code aux model resolves via the primary session token");
assert_eq!(
k2_cfg.api_key.as_deref(),
kimi_cfg.api_key.as_deref(),
Some("kimi-tok"),
"a non-oauth aux model must still receive the primary session token",
"the first-party subscription aux path must be byte-identical",
);
}
#[test]
@@ -1917,6 +1917,51 @@ pub(crate) fn resolve_catalog_key(
.map(|(key, _)| acp::ModelId::new(key.clone()))
}
/// The managed catalog key (`{platform}/{model}`) a routing slug belongs to.
///
/// H5: `SamplingConfig::model` is the BARE routing slug, never the catalog key,
/// and duplicate slugs across platforms are BY DESIGN — the registry guarantees
/// an API-key platform and its subscription-OAuth twin list the SAME ids
/// (`xai`/`xai-grok`, `anthropic`/`claude-pro-max`, `openai`/`openai-codex`),
/// and `PlatformId::ALL` orders every API-key platform FIRST. A slug scan
/// therefore resolves the WRONG platform for a user holding both credentials:
/// the OAuth twin loses its live `bearer_resolver` (no mid-session refresh → an
/// unrecoverable 401 ~1h in), its Messages adaptation and its Copilot/Codex
/// identity headers.
///
/// `current_key` — [`ModelsManager::current_model_id`], the catalog key the
/// picker actually selected — is therefore authoritative whenever it names this
/// slug. 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.
pub(crate) fn managed_key_for_slug(
models: &IndexMap<String, ModelEntry>,
current_key: Option<&str>,
slug: &str,
) -> Option<String> {
if let Some(entry) = current_key.and_then(|key| models.get(key))
&& (entry.info.model == slug || current_key == Some(slug))
{
return entry.info.id.clone();
}
let key = resolve_catalog_key(models, &acp::ModelId::new(slug.to_string()))?;
models.get(key.0.as_ref())?.info.id.clone()
}
/// The registry platform a routing slug resolves to, via [`managed_key_for_slug`].
/// `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 gate, the manager and the wire adaptations can never disagree.
pub(crate) fn platform_for_slug(
models: &IndexMap<String, ModelEntry>,
current_key: Option<&str>,
slug: &str,
) -> Option<kigi_models::PlatformId> {
let key = managed_key_for_slug(models, current_key, slug);
kigi_models::parse_managed_model_key(key.as_deref().unwrap_or(slug))
.map(|(platform, _)| platform)
}
/// Catalog key for a persisted session model id, restricted to **selectable**
/// entries. A selectable exact-key match wins (as in [`resolve_catalog_key`]);
/// otherwise the last selectable entry whose routing slug matches `id`, so a
@@ -3668,6 +3713,52 @@ mod tests {
assert!(!info.visible_for_auth(false));
}
/// SHIP-BLOCKER regression: a user who signed in with ONLY a Claude Pro/Max
/// subscription has no PRIMARY (Kimi) session, so `is_session_auth()` is
/// false. Stamping `supported_in_api = !uses_oauth()` therefore hid every
/// one of their models — `available()` returned an empty picker and the
/// whole subscription-OAuth feature was dead for its target user. The
/// claude-pro-max entry must be visible with no primary session at all.
#[test]
fn claude_pro_max_only_user_sees_their_models_in_the_picker() {
let wire: kigi_models::WireModel =
serde_json::from_value(serde_json::json!({ "id": "claude-opus-4-8" }))
.expect("wire model fixture");
let entry_config = crate::agent::models_fetch::platform_wire_model_to_entry(
kigi_models::PlatformId::ClaudeProMax,
wire,
"https://api.anthropic.com/v1",
);
let entry = ModelEntry::from_config_entry(&entry_config);
let key = entry
.info
.id
.clone()
.expect("platform entries carry a managed catalog key");
assert_eq!(key, "claude-pro-max/claude-opus-4-8");
let mut catalog = IndexMap::new();
catalog.insert(key.clone(), entry);
// Empty home ⇒ the primary AuthManager holds no credential at all.
let home = tempfile::tempdir().expect("tempdir");
let mgr = ModelsManager::new(
None,
catalog,
acp::ModelId::new(Arc::from(key.clone())),
Arc::new(AuthManager::new(home.path(), KimiCodeConfig::default())),
config::Config::default(),
);
assert!(
!mgr.is_session_auth(),
"a claude-pro-max-only user has no PRIMARY (Kimi) OAuth session"
);
assert!(
mgr.available()
.contains_key(&acp::ModelId::new(Arc::from(key.clone()))),
"the claude-pro-max model must reach the picker without a primary session"
);
}
// ── duplicate model slug re-keying (A/B experiment "auto" alias) ──
fn make_entry_config(model: &str, name: Option<&str>) -> config::ModelEntryConfig {
@@ -645,9 +645,15 @@ pub(crate) fn platform_wire_model_to_entry(
inference_idle_timeout_secs: None,
max_retries: None,
hidden: false,
// Subscription models require the OAuth session; open-platform
// models are usable by API-key users.
supported_in_api: !platform.uses_oauth(),
// `supported_in_api: false` hides a model unless the PRIMARY session is
// an OAuth session (`ModelInfo::visible_for_auth`). Only `kimi-code`
// rides that primary session, so only it may be gated on it. Every
// other OAuth platform (claude-pro-max, openai-codex, github-copilot,
// xai-grok) carries its OWN pooled credential, and its models only
// enter the catalog once THAT provider is signed in — gating them on
// the Kimi session would hide every model from a user who signed in
// with only a Claude/ChatGPT/Copilot/Grok subscription.
supported_in_api: platform != kigi_models::PlatformId::KimiCode,
supports_backend_search: false,
compactions_remaining: None,
compaction_at_tokens: None,
@@ -898,6 +904,16 @@ fn get_string_map(
mod tests {
use super::*;
/// Whether a freshly-fetched platform entry shows up in the picker for a
/// user whose PRIMARY session is NOT an OAuth session (`is_session_auth ==
/// false`): an API-key user, or — the case that made this a ship-blocker —
/// someone who signed in with ONLY a Claude Pro/Max, ChatGPT, Copilot, or
/// Grok subscription. `ModelInfo::visible_for_auth` is the picker's real
/// predicate (`agent/models.rs` → `available()`).
fn visible_to_non_primary_session_user(entry: &crate::agent::config::ModelEntryConfig) -> bool {
crate::agent::config::ModelEntry::from_config_entry(entry).visible_for_auth(false)
}
/// OpenAI-cycle e2e (mock wire): a polluted bare-id `/models` listing +
/// a models.dev refresh produce a catalog with ONLY chat models, enriched
/// context windows / efforts, and the Responses backend — the full
@@ -1277,8 +1293,14 @@ mod tests {
"enrichment fills the context window from models.dev anthropic"
);
assert!(
!opus.supported_in_api,
"subscription (uses_oauth) models require the OAuth session"
opus.supported_in_api,
"claude-pro-max carries its OWN pooled credential — it must NOT be \
gated on the primary (Kimi) session"
);
assert!(
visible_to_non_primary_session_user(opus),
"a Claude-Pro/Max-only user has no primary OAuth session; their \
models must still appear in the picker"
);
}
@@ -1409,8 +1431,14 @@ mod tests {
"context window comes from models.dev github-copilot enrichment"
);
assert!(
!entry.supported_in_api,
"subscription (uses_oauth) models require the OAuth session"
entry.supported_in_api,
"github-copilot carries its OWN pooled credential — it must NOT be \
gated on the primary (Kimi) session"
);
assert!(
visible_to_non_primary_session_user(entry),
"a Copilot-only user has no primary OAuth session; their models \
must still appear in the picker"
);
}
@@ -1484,8 +1512,14 @@ mod tests {
assert_eq!(sol.context_window.get(), 272_000);
assert_eq!(sol.name.as_deref(), Some("GPT-5.6-Sol"));
assert!(
!sol.supported_in_api,
"subscription (uses_oauth) models require the OAuth session"
sol.supported_in_api,
"openai-codex carries its OWN pooled credential — it must NOT be \
gated on the primary (Kimi) session"
);
assert!(
visible_to_non_primary_session_user(sol),
"a ChatGPT/Codex-only user has no primary OAuth session; their \
models must still appear in the picker"
);
assert!(sol.supports_reasoning_effort);
assert_eq!(
@@ -3563,8 +3597,14 @@ mod tests {
"an OAuth channel carries no api-key env"
);
assert!(
!entry.supported_in_api,
"subscription models require the OAuth session (not the public API)"
entry.supported_in_api,
"xai-grok carries its OWN pooled credential — it must NOT be gated \
on the primary (Kimi) session"
);
assert!(
visible_to_non_primary_session_user(entry),
"a Grok-only user has no primary OAuth session; their models must \
still appear in the picker"
);
// Passthrough dialect (identical to the API-key xai wire).
let model_entry = crate::agent::config::ModelEntry::from_config_entry(entry);
@@ -25,17 +25,19 @@ impl MvpAgent {
primary: &SamplingConfig,
) -> Result<(OaiCompatClient, String), acp::Error> {
let slug = self.resolve_session_summary_model();
// Resolve the aux token by the summary model's OWN platform: a grok
// (oauth-platform) summary model draws its pooled grok token or `None`
// — NEVER the primary Kimi session token (which `resolve_credentials`
// would otherwise stamp onto an api.x.ai request). A first-party /
// non-oauth summary model still gets the primary (byte-identical).
let session_key = crate::auth::oauth_registry::session_key_for_model(
&crate::util::kigi_home::kigi_home(),
let models = self.models_manager.models();
// Resolve the aux token by the summary model's OWN platform AND
// endpoint: a grok (oauth-platform) summary model draws its pooled grok
// token or `None`, and an API-key registry platform draws NOTHING —
// NEVER the primary Kimi session token (which `resolve_credentials`
// would otherwise stamp onto an api.x.ai / api.deepseek.com request).
// The first-party subscription channel still gets the primary
// (byte-identical).
let session_key = crate::auth::oauth_registry::session_key_for_catalog_model(
&models,
&slug,
Some(&self.auth_manager),
);
let models = self.models_manager.models();
let endpoints = self.models_manager.endpoints();
let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone();
let config = match crate::agent::config::resolve_aux_model_sampling_config(
@@ -47,7 +49,17 @@ impl MvpAgent {
) {
Some(mut cfg) => {
cfg.attribution_callback = primary.attribution_callback.clone();
cfg.bearer_resolver = primary.bearer_resolver.clone();
// H4: the SESSION model's bearer_resolver must not ride to a
// summary model on a different provider —
// `SamplingClient::post` REPLACES the request's auth header
// from it, overwriting the summary model's own resolved key on
// ITS host. Route through the one aux-resolver decision.
cfg.bearer_resolver =
crate::session::acp_session::sampler_turn::aux_bearer_resolver(
primary.bearer_resolver.clone(),
crate::agent::models::platform_for_slug(&models, None, &slug),
&cfg.base_url,
);
cfg.max_retries = primary.max_retries;
cfg
}
@@ -669,31 +681,51 @@ impl MvpAgent {
);
Ok(entry.clone())
}
/// Resolve the SESSION token for `model` by the model's OWN platform — the
/// single guard against the api_key-channel token leak.
/// Resolve the SESSION token for `model` by the model's OWN platform AND
/// endpoint — the single guard against the api_key-channel token leak.
///
/// An oauth-platform model (xai-grok) draws its session token from ITS OWN
/// process-global pool manager (build-on-demand from the on-disk grok token,
/// proactively refreshed), INDEPENDENT of the primary `auth_method`; when
/// that provider has no stored session the token is `None` — NEVER the
/// primary Kimi key. Every other model (first-party / Kimi) 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.
/// - an oauth-platform model (xai-grok, claude-pro-max, github-copilot,
/// openai-codex) draws its session token from ITS OWN process-global pool
/// manager (built on demand from the on-disk token, proactively
/// refreshed), INDEPENDENT of the primary `auth_method`; when that
/// provider has no stored session the token is `None` — NEVER the primary
/// Kimi key;
/// - an endpoint that does not take a session credential at all
/// ([`crate::agent::auth_method::platform_takes_session_credential`] —
/// every API-key registry platform, and any `[model.*]` block pointed at a
/// third-party host) gets `None`. This is C1: `resolve_credentials` takes
/// the `else if let Some(key) = session_key` arm and sets
/// `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.
fn session_token_for_model(&self, model: &ModelEntry) -> Option<crate::auth::KimiAuth> {
if let Some(oauth) = model
.info()
let info = model.info();
let platform = info
.id
.as_deref()
.and_then(kigi_models::parse_managed_model_key)
.and_then(|(platform, _)| platform.oauth())
{
.map(|(platform, _)| platform);
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
return crate::auth::oauth_registry::global_manager_for(
&crate::util::kigi_home::kigi_home(),
&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;
}
if self.is_session_based_auth() {
self.auth_manager.current_or_expired()
} else {
@@ -1197,18 +1197,26 @@ async fn prepare_sampling_config_never_stamps_kimi_key_on_grok_model() {
let endpoints = EndpointsConfig::default();
// First-party Kimi model (non-oauth platform): the primary session key IS
// its api_key — the byte-identical primary path, and proof the Kimi token is
// First-party SUBSCRIPTION model (kimi-code): the primary session key IS its
// api_key — the byte-identical primary path, and proof the Kimi token is
// live (so it WOULD leak if mis-routed onto a grok request). This assertion
// also confirms the session-based primary path is active.
let mut kimi_model = ModelEntry::fallback("kimi-k2-0905-preview", &endpoints);
kimi_model.info.id = Some("moonshot-cn/kimi-k2-0905-preview".to_string());
//
// This used to use `moonshot-cn/kimi-k2-0905-preview` and assert the SAME
// thing, which encoded the C1 defect: moonshot-cn is an API-key registry
// platform on `api.moonshot.cn`, NOT first-party, so "must carry the primary
// session key" was asserting the leak. `api_key_channel_leak_tests` now pins
// the opposite for every moonshot entry.
let mut kimi_model = ModelEntry::fallback("kimi-for-coding", &endpoints);
kimi_model.info.id = Some("kimi-code/kimi-for-coding".to_string());
kimi_model.info.base_url = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url.to_string();
assert!(!kimi_model.has_own_credentials());
let kimi_cfg = agent.prepare_sampling_config_for_model(&kimi_model, None);
assert_eq!(
kimi_cfg.api_key.as_deref(),
Some(KIMI_KEY),
"a first-party Kimi model must carry the primary session key (primary path unchanged)"
"the first-party subscription model must carry the primary session key \
(primary path unchanged)"
);
// xai-grok model (oauth platform): the session token resolves from its OWN
@@ -1280,6 +1288,9 @@ async fn ensure_plugin_registry_lazily_populates_snapshot() {
);
}
mod subagent_spawn_context_tests;
/// LEAK guard for the `api_key` channel (C1/C2), through the real
/// `prepare_sampling_config_for_model` resolution path.
mod api_key_channel_leak_tests;
/// No load in flight and no session → the wait returns immediately
/// (the caller then surfaces "unknown session id" exactly as before).
#[tokio::test]
@@ -0,0 +1,313 @@
//! LEAK GUARD (`api_key` channel) — C1/C2, driven through the REAL resolution
//! path, `MvpAgent::prepare_sampling_config_for_model`.
//!
//! This is the channel the `bearer_resolver` guard does NOT close, and the one
//! the first round of leak tests assumed away by hand-stamping a provider key
//! into chat state. The chain: `session_token_for_model` used to fall through to
//! `self.auth_manager.current_or_expired()` (the primary Kimi bearer) for every
//! non-OAuth model, `resolve_credentials` then took its
//! `else if let Some(key) = session_key` arm and set `api_key = <Kimi token>`
//! with the THIRD-PARTY `base_url`, and `SamplingClient` builds
//! `Authorization: Bearer <api_key>` straight into `default_headers` — which
//! `post()` only overrides when a resolver exists, so `bearer_resolver: None`
//! does not save it.
//!
//! Nothing here stamps a credential by hand: every assertion reads what the
//! resolution path actually produced.
use super::super::*;
use crate::agent::auth_method::{
CACHED_TOKEN_AUTH_METHOD_ID, HOUSE_API_KEY_ENV_VAR, LEGACY_XAI_API_KEY_ENV_VAR,
XAI_API_KEY_ENV_VAR,
};
use crate::agent::config::{Config as AgentConfig, EndpointsConfig, EnvKeys, ModelEntry};
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
use kigi_test_support::EnvGuard;
const KIMI_TOKEN: &str = "kimi-subscription-token-DO-NOT-LEAK";
/// 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
/// leak under test. Every test holding these must be `#[serial]`.
fn without_ambient_byok_env() -> [EnvGuard; 3] {
[
EnvGuard::unset(HOUSE_API_KEY_ENV_VAR),
EnvGuard::unset(XAI_API_KEY_ENV_VAR),
EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR),
]
}
/// An `MvpAgent` on a session-based (`cached_token`) ACP method holding a live
/// Kimi subscription bearer — the mainstream configuration in which the leak
/// fires. `(tempdir, agent)`; the tempdir is the auth store and is returned so
/// the caller keeps it alive.
fn kimi_session_agent() -> (tempfile::TempDir, MvpAgent) {
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 (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let agent = MvpAgent::new(
GatewaySender::new(tx),
&AgentConfig::default(),
auth_manager,
None,
)
.expect("valid test config");
agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID));
(dir, agent)
}
/// 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.
fn platform_entry(catalog_key: &str, slug: &str, base_url: &str) -> ModelEntry {
let mut entry = ModelEntry::fallback(slug, &EndpointsConfig::default());
entry.info.id = Some(catalog_key.to_string());
entry.info.base_url = base_url.to_string();
entry
}
/// C1, the ZERO-CONFIGURATION repro. `default_models.json` bundles
/// `moonshot-cn/*` and `moonshot-ai/*` entries with `api_key: None`, and
/// `resolve_model_list` keeps the bundled defaults whenever no catalog fetch has
/// succeeded — so on first launch / offline a Kimi-subscription user sees them
/// in the picker with no configuration whatsoever. Selecting one used to send
/// `Authorization: Bearer <Kimi OAuth token>` to `api.moonshot.cn`, which is NOT
/// first-party.
///
/// Revert-to-red: dropping the `platform_takes_session_credential` guard from
/// `session_token_for_model` makes every `api_key` below `Some(KIMI_TOKEN)`.
#[tokio::test]
#[serial_test::serial]
async fn bundled_default_moonshot_models_never_carry_the_kimi_bearer() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
let bundled = crate::agent::config::default_model_entries(&EndpointsConfig::default());
let moonshot: Vec<_> = bundled
.iter()
.filter(|(key, _)| key.starts_with("moonshot-cn/") || key.starts_with("moonshot-ai/"))
.collect();
assert_eq!(
moonshot.len(),
4,
"default_models.json still bundles the four moonshot open-platform entries"
);
for (key, entry) in moonshot {
assert!(
!entry.has_own_credentials(),
"{key}: the bundled entry carries no credential of its own"
);
assert!(
!crate::util::is_effective_coding_endpoint_url(&entry.info().base_url),
"{key}: routes to a third-party host ({})",
entry.info().base_url
);
let cfg = agent.prepare_sampling_config_for_model(entry, None);
assert_ne!(
cfg.api_key.as_deref(),
Some(KIMI_TOKEN),
"LEAK: selecting the bundled {key} sent the Kimi subscription bearer to {}",
entry.info().base_url
);
assert_eq!(
cfg.base_url,
entry.info().base_url,
"{key}: still routes to its own host (the fix must not reroute traffic)"
);
}
}
/// C1 across the API-key registry platform shapes a fetched catalog produces.
#[tokio::test]
#[serial_test::serial]
async fn api_key_platform_models_never_carry_the_kimi_bearer_as_api_key() {
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"),
("openai/gpt-5.2", "gpt-5.2", "https://api.openai.com/v1"),
("anthropic/claude-opus-4-8", "claude-opus-4-8", "https://api.anthropic.com/v1"),
("groq/llama-4", "llama-4", "https://api.groq.com/openai/v1"),
("xai/grok-4.5", "grok-4.5", "https://api.x.ai/v1"),
] {
let entry = platform_entry(catalog_key, slug, base_url);
let cfg = agent.prepare_sampling_config_for_model(&entry, None);
assert_ne!(
cfg.api_key.as_deref(),
Some(KIMI_TOKEN),
"LEAK: {catalog_key} carried the primary Kimi session bearer to {base_url}"
);
}
}
/// C2, the `[model.*]` repro. A `[model.gpt-4o]` block has `info.id == None`, so
/// it has no platform at all — which used to be a blanket allow. BYOK is
/// `has_own_credentials()`, which probes `std::env::var` AT CALL TIME, so an
/// unset (or mistyped) `env_key` classifies the model NotByok and the Kimi
/// bearer went to `api.openai.com` on BOTH channels.
///
/// Revert-to-red: making the `None` arm of `platform_takes_session_credential`
/// return `true` again makes `api_key` here `Some(KIMI_TOKEN)`.
#[tokio::test]
#[serial_test::serial]
async fn config_model_with_an_unset_env_key_never_carries_the_kimi_bearer() {
let _env = without_ambient_byok_env();
let _typo = EnvGuard::unset("OPENAI_API_KEY_TYPO");
let (_dir, agent) = kimi_session_agent();
let mut entry = ModelEntry::fallback("gpt-4o", &EndpointsConfig::default());
entry.info.id = None; // a `[model.gpt-4o]` config block
entry.info.base_url = "https://api.openai.com/v1".to_string();
entry.env_key = Some(EnvKeys::single("OPENAI_API_KEY_TYPO"));
assert!(
!entry.has_own_credentials(),
"the env var is unset, so this classifies NotByok — the precondition of the defect"
);
let cfg = agent.prepare_sampling_config_for_model(&entry, None);
assert_ne!(
cfg.api_key.as_deref(),
Some(KIMI_TOKEN),
"LEAK: a [model.*] block with an unset env_key sent the Kimi bearer to api.openai.com"
);
assert_eq!(cfg.api_key, None, "no credential resolves — fail fast");
}
/// The first-party subscription channel must stay BYTE-IDENTICAL: `kimi-code/*`
/// (and a `[model.*]` block on the session's own coding endpoint, including a
/// `KIGI_CODE_BASE_URL` deployment / a local dev proxy) still carries the
/// primary session key. This assertion is also what proves the Kimi token is
/// live in the tests above — it WOULD leak if the guard were missing.
#[tokio::test]
#[serial_test::serial]
async fn the_first_party_subscription_channel_still_carries_the_session_key() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
let kimi = platform_entry(
"kimi-code/kimi-for-coding",
"kimi-for-coding",
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
);
assert_eq!(
agent
.prepare_sampling_config_for_model(&kimi, None)
.api_key
.as_deref(),
Some(KIMI_TOKEN),
"the kimi-code subscription channel must be unchanged"
);
for base_url in [
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
"http://127.0.0.1:4141/v1",
"http://localhost:8080/v1",
] {
let mut bare = ModelEntry::fallback("kigi-4.5", &EndpointsConfig::default());
bare.info.id = None;
bare.info.base_url = base_url.to_string();
assert_eq!(
agent
.prepare_sampling_config_for_model(&bare, None)
.api_key
.as_deref(),
Some(KIMI_TOKEN),
"{base_url}: a custom deployment / local proxy keeps the session key"
);
}
}
/// A subscription-OAuth model draws its `api_key` from ITS OWN pooled manager,
/// never the Kimi primary — and never falls back to it when that provider has no
/// stored session (the pool home is an empty TempDir under `cfg(test)`).
#[tokio::test]
#[serial_test::serial]
async fn oauth_platform_models_never_carry_the_kimi_bearer_as_api_key() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
for (catalog_key, slug, base_url) in [
("xai-grok/grok-4-latest", "grok-4-latest", "https://api.x.ai/v1"),
(
"claude-pro-max/claude-opus-4-8",
"claude-opus-4-8",
"https://api.anthropic.com/v1",
),
("github-copilot/gpt-4.1", "gpt-4.1", "https://api.githubcopilot.com"),
(
"openai-codex/gpt-5.5",
"gpt-5.5",
"https://chatgpt.com/backend-api/codex",
),
] {
let entry = platform_entry(catalog_key, slug, base_url);
let cfg = agent.prepare_sampling_config_for_model(&entry, None);
assert_ne!(
cfg.api_key.as_deref(),
Some(KIMI_TOKEN),
"LEAK: {catalog_key} carried the primary Kimi session bearer to {base_url}"
);
}
}
/// H5 at the api_key channel: `resolve_model_id` (the picker's own lookup) must
/// hand `prepare_sampling_config_for_model` the entry the user SELECTED, even
/// when an API-key platform and its subscription-OAuth twin list the same
/// routing slug in `PlatformId::ALL` order. Selecting the OAuth twin by catalog
/// key must not resolve the API-key twin — and vice versa.
#[tokio::test]
#[serial_test::serial]
async fn dual_credential_slug_collision_resolves_the_selected_catalog_key() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
// API-key platform FIRST, exactly as `PlatformId::ALL` orders them.
for key in ["xai/grok-4.5", "xai-grok/grok-4.5"] {
agent.models_manager.insert_test_entry(
key,
platform_entry(key, "grok-4.5", "https://api.x.ai/v1"),
);
}
for key in ["xai/grok-4.5", "xai-grok/grok-4.5"] {
let resolved = agent
.resolve_model_id(&acp::ModelId::new(key))
.expect("both twins resolve");
assert_eq!(
resolved.info().id.as_deref(),
Some(key),
"selecting {key} must resolve THAT catalog entry, not its slug twin"
);
}
// And the bare slug resolves the same entry the picker's `resolve_catalog_key`
// does — one direction, one answer (the auth layer used to first-match).
let by_slug = agent
.resolve_model_id(&acp::ModelId::new("grok-4.5"))
.expect("the bare slug resolves");
let models = agent.models_manager.models();
let picker_key = crate::agent::models::resolve_catalog_key(
&models,
&acp::ModelId::new("grok-4.5"),
)
.expect("the picker resolves the bare slug");
assert_eq!(
by_slug.info().id.as_deref(),
Some(picker_key.0.as_ref()),
"the auth layer and the picker must resolve the SAME entry for one slug"
);
assert_eq!(
crate::agent::config::find_model_by_id(&models, "grok-4.5")
.and_then(|e| e.info().id.as_deref()),
Some(picker_key.0.as_ref()),
"find_model_by_id must agree with resolve_catalog_key by construction"
);
}
@@ -1003,16 +1003,22 @@ fn resolve_model_override_to_config(
} else {
acp::ModelId::new(entry.info().model.clone())
};
// Resolve the child's session token by the OVERRIDE model's OWN platform,
// not the parent's primary auth: a grok (oauth-platform) override draws its
// pooled grok token or `None` — NEVER the primary Kimi session token (which
// `resolve_credentials` would otherwise stamp onto the child's api.x.ai
// credentials, leaking it in the logout-mid-session edge). A first-party /
// non-oauth override still resolves to the primary (byte-identical).
let managed_key = entry.info().id.as_deref().unwrap_or(model_id);
let session_key = crate::auth::oauth_registry::session_key_for_model(
&crate::util::kigi_home::kigi_home(),
managed_key,
// Resolve the child's session token by the OVERRIDE model's OWN platform
// AND endpoint, not the parent's primary auth: a grok (oauth-platform)
// override draws its pooled grok token or `None`, and an API-key registry
// platform / a third-party `[model.*]` host draws NOTHING — NEVER the
// primary Kimi session token, which `resolve_credentials` would otherwise
// stamp onto the child's api.x.ai / api.moonshot.cn credentials. The
// first-party subscription channel still resolves to the primary
// (byte-identical).
let session_key = crate::auth::oauth_registry::session_key_for_endpoint(
entry
.info()
.id
.as_deref()
.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();
@@ -3131,19 +3131,25 @@ fn fresh_tool_model_rejects_unavailable_exact_key_over_visible_slug_collision()
"validation must inspect the unavailable exact-key entry selected by execution"
);
}
/// Validation must inspect the SAME slug-collision entry execution selects.
/// Both go through `find_model_by_id`, whose slug scan takes the LAST match —
/// aligned with the picker's `resolve_catalog_key` so the auth layer and the
/// picker can never resolve different platforms for one slug (the H5 collision).
/// So a blocked LAST entry must be rejected even though an available earlier one
/// shares the slug.
#[test]
fn fresh_tool_model_rejects_unavailable_first_slug_collision() {
fn fresh_tool_model_rejects_unavailable_last_slug_collision() {
let mut models = indexmap::IndexMap::new();
let mut unavailable_first = test_model_entry("shared-routing-slug");
unavailable_first.info.user_selectable = false;
models.insert("blocked-first".to_string(), unavailable_first);
models.insert("visible-second".to_string(), test_model_entry("shared-routing-slug"));
models.insert("visible-first".to_string(), test_model_entry("shared-routing-slug"));
let mut unavailable_last = test_model_entry("shared-routing-slug");
unavailable_last.info.user_selectable = false;
models.insert("blocked-last".to_string(), unavailable_last);
assert_eq!(
super::handle_request::task_model_override_error(Some("shared-routing-slug"),
ModelOverrideProvenance::Tool, false, & models, false,).as_deref(),
Some("Unknown Task.model slug 'shared-routing-slug'. Valid model slugs: \
visible-second. Omit `model` to inherit the parent model."),
"validation must inspect the first routing-slug entry selected by execution"
visible-first. Omit `model` to inherit the parent model."),
"validation must inspect the last routing-slug entry selected by execution"
);
}
#[test]
@@ -2511,14 +2511,48 @@ async fn subagent_override_grok_model_never_leaks_kimi_session_token() {
"a grok override must never receive the primary Kimi session token",
);
}
/// Byte-identical guard: a non-oauth override with a Kimi primary still resolves
/// to the primary session token — passes both before and after the fix (the
/// non-oauth path is unchanged).
/// Byte-identical guard: an override on the SESSION's own first-party endpoint
/// (the kimi-code subscription channel) still resolves to the primary session
/// token — the primary path is unchanged.
#[tokio::test]
async fn subagent_override_non_oauth_model_still_gets_primary_token() {
async fn subagent_override_first_party_model_still_gets_primary_token() {
let (_kd, manager) = kimi_primary_with_token("kimi-secret");
let mut entry = test_model_entry("kimi-for-coding");
entry.info.id = Some("kimi-code/kimi-for-coding".to_string());
entry.info.base_url = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url.to_string();
let mut models = indexmap::IndexMap::new();
models.insert("kfc".to_string(), entry);
let mut ctx = ctx_with_toggle(HashMap::new());
ctx.available_models = models;
ctx.auth = Some(crate::auth::KimiAuth {
key: "kimi-secret".to_string(),
auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
ctx.auth_manager = manager;
let (config, _model_id) = resolve_model_override_to_config("kfc", &ctx)
.expect("first-party override resolves to a config");
assert_eq!(
config.api_key.as_deref(),
Some("kimi-secret"),
"a first-party override must still receive the primary session token",
);
}
/// LEAK guard (C1, subagent-override `api_key` channel): an API-key registry
/// platform override must NOT receive the parent's primary Kimi session token —
/// `resolve_credentials` would stamp it as the child's `api_key` on
/// `api.moonshot.cn`. This test previously asserted the opposite
/// (`subagent_override_non_oauth_model_still_gets_primary_token`), which encoded
/// the defect.
///
/// Revert-to-red: dropping the `platform_takes_session_credential` term from
/// `oauth_registry::session_key_for_endpoint` makes `api_key` `Some("kimi-secret")`.
#[tokio::test]
async fn subagent_override_api_key_platform_never_gets_the_primary_token() {
let (_kd, manager) = kimi_primary_with_token("kimi-secret");
let mut entry = test_model_entry("kimi-k2-0905-preview");
entry.info.id = Some("moonshot-cn/kimi-k2".to_string());
entry.info.base_url = "https://api.moonshot.cn/v1".to_string();
let mut models = indexmap::IndexMap::new();
models.insert("k2".to_string(), entry);
let mut ctx = ctx_with_toggle(HashMap::new());
@@ -2529,12 +2563,12 @@ async fn subagent_override_non_oauth_model_still_gets_primary_token() {
..crate::auth::KimiAuth::test_default()
});
ctx.auth_manager = manager;
let (config, _model_id) =
resolve_model_override_to_config("k2", &ctx).expect("non-oauth override resolves to a config");
assert_eq!(
let (config, _model_id) = resolve_model_override_to_config("k2", &ctx)
.expect("an API-key-platform override still resolves to a config");
assert_ne!(
config.api_key.as_deref(),
Some("kimi-secret"),
"a non-oauth override must still receive the primary session token",
"LEAK: an API-key-platform override must never receive the primary Kimi session token",
);
}
/// An unresolvable `AgentDefinition.model` pin (model absent from