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
@@ -13,11 +13,10 @@
//! generic-oauth scope, each wired with the SAME lifecycle as the primary Kimi
//! manager (`configure_refresher()` + `start_proactive_refresh()`) so the
//! on-disk token stays fresh and a 401 recovers via the provider's own manager.
//! Managers are built ON DEMAND: the first grok turn (or model switch) reads the
//! on-disk token via [`global_manager_for`], so a login that lands AFTER a
//! session spawned self-heals — there is no frozen per-session snapshot to go
//! stale. [`manager_for_model`] routes a managed catalog key to the pool (oauth
//! platform) or to the session's primary (everything else).
//! Managers are built ON DEMAND from the on-disk token ([`global_manager_for`]),
//! so a login landing AFTER a session spawned self-heals — no frozen per-session
//! snapshot. [`manager_for_model`] routes a managed catalog key to the pool
//! (oauth platform) or to the session's primary (everything else).
//!
//! SECURITY: access/refresh tokens and resolved bearers are NEVER logged here.
@@ -39,6 +38,31 @@ fn oauth_manager_pool() -> &'static Mutex<HashMap<&'static str, Arc<AuthManager>
POOL.get_or_init(|| Mutex::new(HashMap::new()))
}
/// The kigi home every OAuth-pool call site resolves from. Single definition so
/// the pool, the aux/summary token routing and the session's inference manager
/// can never read different homes.
///
/// Production: [`crate::util::kigi_home::kigi_home`]. LIB TESTS: a
/// process-lifetime `TempDir`, unconditionally — the pool is process-global and
/// every manager it builds starts a never-cancelled proactive-refresh loop, so
/// a unit test resolving the real `~/.kigi` would read the developer's stored
/// OAuth tokens and, 60 s later, fire REAL refresh requests against them.
/// Deliberately not a per-test opt-in that can be forgotten: `kigi_home()` is
/// itself a `OnceLock` an earlier test has usually already resolved to the real
/// home, so setting `KIGI_SHARE_DIR` in a test cannot pin it after the fact.
pub(crate) fn pool_home() -> std::path::PathBuf {
#[cfg(test)]
{
static TEST_HOME: OnceLock<tempfile::TempDir> = OnceLock::new();
TEST_HOME
.get_or_init(|| tempfile::tempdir().expect("tempdir for the test OAuth pool"))
.path()
.to_path_buf()
}
#[cfg(not(test))]
crate::util::kigi_home::kigi_home()
}
/// Get-or-create the process-global manager for `oauth`, wiring the same
/// refresher + proactive-refresh lifecycle as the primary Kimi manager the
/// FIRST time a scope is seen. The manager reads the on-disk token at
@@ -93,27 +117,70 @@ pub(crate) fn manager_for_model(
primary.cloned()
}
/// The SESSION token (the raw bearer/key string) that governs INFERENCE auth
/// for `managed_key`, resolved by the model's OWN platform. Thin wrapper over
/// [`manager_for_model`] used by the aux-model and subagent-override wire paths
/// so a `{platform}/{model}` key never receives the primary token of a
/// DIFFERENT provider.
/// 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) draws its token from ITS OWN
/// pooled manager; when that provider has no stored session the result is
/// `None` — NEVER the primary Kimi key. Every other key routes to `primary` and
/// yields the primary's current-or-expired token, byte-identical to reading it
/// directly. SECURITY: the resolved token is never logged.
pub(crate) fn session_key_for_model(
kigi_home: &Path,
managed_key: &str,
/// - 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> {
manager_for_model(kigi_home, managed_key, primary)
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
return global_manager_for(&pool_home(), oauth)
.current_or_expired()
.map(|a| a.key);
}
if !crate::agent::auth_method::platform_takes_session_credential(platform, base_url) {
return None;
}
primary
.and_then(|am| am.current_or_expired())
.map(|a| a.key)
}
/// [`session_key_for_endpoint`] for the catalog model whose routing slug (or
/// catalog key) is `slug`.
///
/// A slug absent from the catalog keeps the pre-registry behaviour: the aux
/// resolver's Tier-2 fallback builds its entry against
/// `EndpointsConfig::resolve_inference_base_url` (first-party), so the primary
/// still governs.
pub(crate) fn session_key_for_catalog_model(
models: &indexmap::IndexMap<String, crate::agent::config::ModelEntry>,
slug: &str,
primary: Option<&Arc<AuthManager>>,
) -> Option<String> {
let Some(entry) = crate::agent::config::find_model_by_id(models, slug) else {
return primary
.and_then(|am| am.current_or_expired())
.map(|a| a.key);
};
let info = entry.info();
session_key_for_endpoint(
info.id
.as_deref()
.and_then(kigi_models::parse_managed_model_key)
.map(|(platform, _)| platform),
&info.base_url,
primary,
)
}
#[cfg(test)]
mod tests {
use super::*;
@@ -159,6 +226,20 @@ mod tests {
.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
@@ -187,7 +268,7 @@ mod tests {
"openai-codex and claude-pro-max must not share a pooled manager"
);
assert_ne!(
session_key_for_model(home.path(), "openai-codex/gpt-5.5", Some(&kimi)),
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"
);
@@ -218,7 +299,7 @@ mod tests {
// 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_model(home.path(), "github-copilot/gpt-4.1", Some(&kimi)),
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"
);
@@ -256,9 +337,8 @@ mod tests {
#[tokio::test]
async fn session_key_for_claude_pro_max_is_never_the_kimi_primary() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
assert_ne!(
session_key_for_model(home.path(), "claude-pro-max/claude-opus-4-8", Some(&kimi)),
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"
);
@@ -342,22 +422,62 @@ mod tests {
);
}
/// `session_key_for_model`: a non-oauth / bare key yields the primary Kimi
/// token exactly as reading it directly would — byte-identical to the
/// pre-fix aux/override wire path (no runtime / pool touched).
/// `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_non_oauth_is_the_primary_token() {
fn session_key_for_the_sessions_own_endpoint_is_the_primary_token() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
for key in ["moonshot-cn/kimi-k2", "kimi-k2-0905-preview"] {
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_model(home.path(), key, Some(&kimi)),
session_key_for_endpoint(None, url, Some(&kimi)),
Some("kimi-tok".to_string()),
"{key} (non-oauth) must yield the primary token unchanged"
"{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
/// from `session_key_for_endpoint` returns `Some("kimi-tok")` here.
#[test]
fn session_key_for_a_third_party_endpoint_is_never_the_primary_token() {
let (_kd, kimi) = primary_with_token("kimi-tok");
for key in [
"moonshot-cn/kimi-k2",
"deepseek/deepseek-chat",
"openai/gpt-5",
] {
assert_eq!(
session_key_for_key(key, Some(&kimi)),
None,
"LEAK: {key} is an API-key platform — no session token may ride there"
);
}
assert_eq!(
session_key_for_endpoint(None, "https://api.openai.com/v1", Some(&kimi)),
None,
"LEAK: a [model.*] block on a third-party host gets no session token"
);
}
/// LEAK guard (aux-model + subagent-override token routing): a grok key with
/// a Kimi primary NEVER yields the primary Kimi token — it draws from the
/// pooled xai manager (its own token, or `None`). This is the exact source
@@ -365,15 +485,14 @@ mod tests {
#[tokio::test]
async fn session_key_for_grok_is_never_the_kimi_primary() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
assert_ne!(
session_key_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi)),
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_model(home.path(), "xai-grok/grok-4-fast", None),
session_key_for_key("xai-grok/grok-4-fast", None),
Some("kimi-tok".to_string()),
);
}
@@ -135,7 +135,7 @@ use prompt_build::*;
mod session_mode;
use session_mode::*;
#[path = "acp_session_impl/sampler_turn.rs"]
mod sampler_turn;
pub(crate) mod sampler_turn;
use sampler_turn::*;
#[path = "acp_session_impl/tool_dispatch.rs"]
mod tool_dispatch;
@@ -1256,6 +1256,17 @@ mod rewind_synthetic_turn_tests;
#[cfg(test)]
#[path = "acp_session_tests/rewrite_zero_turn_prefix_tests.rs"]
mod rewrite_zero_turn_prefix_tests;
/// The same guard for the model→platform lookup (the dual-credential slug
/// collision) and for the stamped aux/summary configs.
#[cfg(test)]
#[path = "acp_session_tests/session_bearer_leak_platform_tests.rs"]
mod session_bearer_leak_platform_tests;
/// LEAK guard: the primary Kimi subscription bearer must never ride a request
/// to an API-key registry platform's host, while the subscription-OAuth
/// platforms keep a live resolver from their OWN pooled manager.
#[cfg(test)]
#[path = "acp_session_tests/session_bearer_leak_tests.rs"]
mod session_bearer_leak_tests;
/// Pins the `SubagentFinished` usage-fold attribution gate.
#[cfg(test)]
#[path = "acp_session_tests/subagent_usage_fold_tests.rs"]
@@ -634,14 +634,12 @@ impl SessionActor {
&active_session_config,
Some(self.max_retries),
);
// A grok (oauth-platform) image-describe model must not inherit the
// session (Kimi) bearer_resolver stamped by `finalize_*`; re-point it at
// grok's own manager. No-op for a first-party / non-oauth model.
// 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_for_oauth(
&mut sampler_config,
&self.image_description_model,
);
self.repoint_aux_bearer_resolver(&mut sampler_config, &self.image_description_model);
}
let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| {
acp::Error::internal_error().data(format!(
@@ -36,20 +36,37 @@ struct SessionTokenAuthGate {
/// BYOK status still refresh against the first-party cli-chat-proxy hosts without
/// risking a session-token leak to a third-party BYOK endpoint.
endpoint_is_first_party: bool,
/// Whether this model's platform is one whose endpoint accepts a session
/// bearer at all (see
/// [`crate::agent::auth_method::platform_takes_session_credential`]). False
/// for every API-key registry platform, which keeps the primary Kimi bearer
/// off `api.deepseek.com` / `api.openai.com` / … .
endpoint_takes_session_credential: bool,
}
impl SessionTokenAuthGate {
/// Single place `is_session_based` / `endpoint_is_first_party` are derived,
/// so all call sites assemble the gate identically.
/// so all call sites assemble the gate identically. `model_platform` is the
/// registry platform the model routes to (`None` for a bare / `[model.*]`
/// entry) — it MUST be derived from the same lookup
/// ([`SessionActor::managed_key_for_model`]) that
/// [`SessionActor::auth_manager_for_model`] uses, so the gate's verdict and
/// the manager actually wrapped as the bearer resolver can never disagree.
fn new(
auth_method_id: Option<&acp::AuthMethodId>,
model_byok: crate::agent::auth_method::ModelByok,
base_url: &str,
model_platform: Option<kigi_models::PlatformId>,
) -> Self {
Self {
is_session_based: auth_method_id
.is_some_and(crate::agent::auth_method::is_session_based_method),
model_byok,
endpoint_is_first_party: crate::util::is_first_party_url(base_url),
endpoint_takes_session_credential:
crate::agent::auth_method::platform_takes_session_credential(
model_platform,
base_url,
),
}
}
fn active(self) -> bool {
@@ -57,6 +74,7 @@ impl SessionTokenAuthGate {
self.is_session_based,
self.model_byok,
self.endpoint_is_first_party,
self.endpoint_takes_session_credential,
)
}
}
@@ -127,6 +145,43 @@ fn auth_manager_bearer_resolver(
) -> kigi_sampler::SharedBearerResolver {
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 {
pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
let mcp_wait_start = std::time::Instant::now();
@@ -206,7 +261,12 @@ impl SessionActor {
fn auth_gate(&self, model_id: &str, base_url: &str) -> SessionTokenAuthGate {
let byok = self.model_auth_facts(model_id).byok;
let auth_method = self.auth_method_id.load();
SessionTokenAuthGate::new(auth_method.as_deref(), byok, base_url)
SessionTokenAuthGate::new(
auth_method.as_deref(),
byok,
base_url,
self.model_platform(model_id),
)
}
/// The [`AuthManager`](crate::auth::AuthManager) that governs INFERENCE auth
/// for the model whose routing slug is `model` (the sampling config's
@@ -238,26 +298,30 @@ impl SessionActor {
// and resolves to the primary.
let managed_key = self.managed_key_for_model(model);
crate::auth::oauth_registry::manager_for_model(
&crate::util::kigi_home::kigi_home(),
&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
/// from the live catalog. `None` for a bare / config / unlisted model.
///
/// H5: the catalog KEY the picker selected
/// ([`crate::agent::models::ModelsManager::current_model_id`]) is
/// authoritative — `model` is the ambiguous bare slug. See
/// [`crate::agent::models::managed_key_for_slug`].
fn managed_key_for_model(&self, model: &str) -> Option<String> {
let models = self.models_manager.models();
crate::agent::config::find_model_by_id(&models, model).and_then(|e| e.info().id.clone())
let current = self.models_manager.current_model_id();
crate::agent::models::managed_key_for_slug(&models, Some(current.0.as_ref()), model)
}
/// Whether the aux/session model `model` routes to a generic device-code
/// OAuth platform (xai-grok). The gate for re-pointing an aux model's
/// bearer_resolver away from the session (Kimi) resolver — a first-party /
/// non-oauth model returns `false` and keeps the stamped session resolver.
fn model_is_oauth_platform(&self, model: &str) -> bool {
let managed_key = self.managed_key_for_model(model);
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model))
.and_then(|(platform, _)| platform.oauth())
.is_some()
/// The registry platform the routing slug `model` belongs to, from the SAME
/// lookup [`Self::auth_manager_for_model`] routes on. `None` for a bare /
/// `[model.*]` / unlisted model.
fn model_platform(&self, model: &str) -> Option<kigi_models::PlatformId> {
let models = self.models_manager.models();
let current = self.models_manager.current_model_id();
crate::agent::models::platform_for_slug(&models, Some(current.0.as_ref()), model)
}
/// Whether `model` routes to the Claude Pro/Max OAuth-Messages platform
/// (claude-pro-max) — the gate for the sampler's OAuth Messages adaptation
@@ -266,22 +330,18 @@ impl SessionActor {
/// which is ChatCompletions) returns `false`, keeping the API-key Anthropic
/// / MiniMax Messages requests byte-identical.
fn model_is_anthropic_oauth(&self, model: &str) -> bool {
let managed_key = self.managed_key_for_model(model);
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model)).is_some_and(
|(platform, _)| {
platform.oauth().is_some()
&& platform.wire_api() == kigi_models::PlatformWireApi::Messages
},
)
self.model_platform(model).is_some_and(|platform| {
platform.oauth().is_some()
&& platform.wire_api() == kigi_models::PlatformWireApi::Messages
})
}
/// Whether `model` routes to the GitHub Copilot ChatCompletions platform
/// (github-copilot) — the gate for the sampler's editor-identity headers +
/// `X-Initiator`. Every other model returns `false`, keeping the other
/// ChatCompletions providers byte-identical.
fn model_is_github_copilot(&self, model: &str) -> bool {
let managed_key = self.managed_key_for_model(model);
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model))
.is_some_and(|(platform, _)| platform.sends_copilot_editor_headers())
self.model_platform(model)
.is_some_and(kigi_models::PlatformId::sends_copilot_editor_headers)
}
/// Whether `model` routes to the ChatGPT/Codex Responses platform
/// (openai-codex) — the gate for the sampler's Codex identity headers
@@ -289,28 +349,23 @@ impl SessionActor {
/// returns `false`, keeping the API-key `openai` Responses request
/// byte-identical.
fn model_is_openai_codex(&self, model: &str) -> bool {
let managed_key = self.managed_key_for_model(model);
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model))
.is_some_and(|(platform, _)| platform.sends_codex_responses_headers())
self.model_platform(model)
.is_some_and(kigi_models::PlatformId::sends_codex_responses_headers)
}
/// LEAK guard for the stamped aux paths (auto-mode classifier, image
/// describe). After [`crate::agent::config::stamp_session_local_sampler_fields`]
/// has copied the SESSION model's `bearer_resolver` onto an aux
/// `SamplerConfig`, re-point it at the AUX model's OWN platform manager when
/// the aux model is an oauth platform (xai-grok) — so a grok aux model never
/// inherits the live Kimi session bearer (→ api.x.ai). No-op for a
/// first-party / non-oauth aux model (keeps the stamped session resolver →
/// byte-identical). SECURITY: no token is logged.
pub(super) fn repoint_aux_bearer_resolver_for_oauth(
/// describe). Applies [`aux_bearer_resolver`] to the config
/// [`crate::agent::config::stamp_session_local_sampler_fields`] just stamped
/// the SESSION model's `bearer_resolver` onto.
pub(super) fn repoint_aux_bearer_resolver(
&self,
cfg: &mut kigi_sampler::SamplerConfig,
slug: &str,
) {
if self.model_is_oauth_platform(slug)
&& let Some(manager) = self.auth_manager_for_model(slug)
{
cfg.bearer_resolver = Some(auth_manager_bearer_resolver(manager));
}
cfg.bearer_resolver = aux_bearer_resolver(
cfg.bearer_resolver.take(),
self.model_platform(slug),
&cfg.base_url,
);
}
/// Emit a unified-log breadcrumb whenever the session-token refresh gate is
/// evaluated with an **`Unknown`** per-model BYOK status on a session-based
@@ -330,8 +385,9 @@ impl SessionActor {
let ctx = serde_json::json!(
{ "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based"
: gate.is_session_based, "endpoint_is_first_party" : gate
.endpoint_is_first_party, "refresh_active" : refresh_active, "base_url" :
base_url, }
.endpoint_is_first_party, "endpoint_takes_session_credential" : gate
.endpoint_takes_session_credential, "refresh_active" : refresh_active,
"base_url" : base_url, }
);
let sid = Some(self.session_info.id.0.as_ref());
if refresh_active {
@@ -385,8 +441,12 @@ impl SessionActor {
let creds = self.chat_state_handle.get_credentials().await;
let model_facts = self.model_auth_facts(cfg.model.as_str());
let auth_method = self.auth_method_id.load();
let gate =
SessionTokenAuthGate::new(auth_method.as_deref(), model_facts.byok, &cfg.base_url);
let gate = SessionTokenAuthGate::new(
auth_method.as_deref(),
model_facts.byok,
&cfg.base_url,
self.model_platform(cfg.model.as_str()),
);
let use_bearer_resolver = gate.active();
self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url);
// Resolve the bearer from the ACTIVE model's OWN manager: a grok model
@@ -589,17 +649,18 @@ impl SessionActor {
slug: &str,
) -> Option<kigi_sampler::SamplerConfig> {
let creds = self.chat_state_handle.get_credentials().await;
// Resolve the aux token by the aux model's OWN platform: a grok
// (oauth-platform) aux 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 aux 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 aux model's OWN platform AND endpoint: a
// grok (oauth-platform) aux 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,
self.auth_manager.as_ref(),
);
let models = self.models_manager.models();
let endpoints = self.models_manager.endpoints();
crate::agent::config::resolve_aux_model_sampling_config(
slug,
@@ -625,10 +686,11 @@ impl SessionActor {
&active_session_config,
Some(self.max_retries),
);
// LEAK 1b: a grok aux classifier must not inherit the SESSION model's
// (Kimi) bearer_resolver stamped above; re-point it at grok's own
// manager (its pooled token, or None). No-op for a non-oauth aux.
self.repoint_aux_bearer_resolver_for_oauth(&mut cfg, slug);
// 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 client = kigi_sampler::SamplingClient::new(cfg)
.map_err(|e| {
@@ -774,6 +836,8 @@ impl SessionActor {
session_id = % self.session_info.id.0, is_session_based = gate
.is_session_based, model_byok = gate.model_byok.as_str(),
endpoint_is_first_party = gate.endpoint_is_first_party,
endpoint_takes_session_credential = gate
.endpoint_takes_session_credential,
"auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401",
);
kigi_log::unified_log::warn(
@@ -783,7 +847,9 @@ impl SessionActor {
{ "kind" : error.kind.as_str(), "status_code" : error
.status_code, "is_session_based" : gate.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
.endpoint_takes_session_credential, }
)),
);
}
@@ -1036,6 +1102,15 @@ impl SessionActor {
.map(|c| c.model)
.unwrap_or_default();
let Some(ref key) = current_key else { return };
// M7: a registry-platform model's key comes from that platform's
// credential resolved into its catalog entry — it is NEVER a
// `[model.*]` block. With the session gate now inactive for every
// API-key platform, those turns all fell through to here and paid a
// `load_effective_config()` disk read PER TURN, then logged a
// permanently false "Model not found in config.toml [model.*]" warning.
if self.model_platform(&current_model_id).is_some() {
return;
}
let Some(new_key) = self.reload_api_key_from_config(&current_model_id) else {
return;
};
@@ -1161,7 +1236,7 @@ mod bearer_resolver_tests {
.oauth()
.expect("xai-grok carries an OAuthConfig");
let grok = crate::auth::oauth_registry::global_manager_for(
&crate::util::kigi_home::kigi_home(),
&crate::auth::oauth_registry::pool_home(),
oauth,
);
assert_ne!(
@@ -487,21 +487,27 @@ fn session_token_auth_gate_truth_table() {
use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate};
// Non-session methods never refresh, regardless of BYOK status or endpoint.
for fp in [false, true] {
assert!(!gate(false, ModelByok::NotByok, fp));
assert!(!gate(false, ModelByok::Byok, fp));
assert!(!gate(false, ModelByok::Unknown, fp));
// Session method: a definite classification ignores the endpoint —
// NotByok always refreshes (only ever routes to the session endpoint),
// a genuine per-model Byok never does.
assert!(gate(true, ModelByok::NotByok, fp));
assert!(!gate(true, ModelByok::Byok, fp));
assert!(!gate(false, ModelByok::NotByok, fp, true));
assert!(!gate(false, ModelByok::Byok, fp, true));
assert!(!gate(false, ModelByok::Unknown, fp, true));
// Session method on an endpoint that DOES take the session credential
// (kimi-code, an OAuth platform's own pool, or a bare / [model.*]
// model): a definite classification ignores the endpoint — NotByok
// refreshes, a genuine per-model Byok never does.
assert!(gate(true, ModelByok::NotByok, fp, true));
assert!(!gate(true, ModelByok::Byok, fp, true));
// …and an API-key registry platform endpoint is refused on every arm,
// first-party flag included: that is the leak guard.
assert!(!gate(true, ModelByok::NotByok, fp, false));
assert!(!gate(true, ModelByok::Byok, fp, false));
assert!(!gate(true, ModelByok::Unknown, fp, false));
}
// Session method + Unknown BYOK: refresh only against a first-party xAI
// 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
// third-party BYOK endpoint. This arm was unconditionally `false` pre-fix.
assert!(gate(true, ModelByok::Unknown, true));
assert!(!gate(true, ModelByok::Unknown, false));
assert!(gate(true, ModelByok::Unknown, true, true));
assert!(!gate(true, ModelByok::Unknown, false, true));
}
/// Pre-fix, the gate read `auth_type` and skipped recovery here, 401'ing every
@@ -0,0 +1,289 @@
//! LEAK GUARD, part 2: the model→platform lookup (H5) and the AUX resolver
//! decision (H3/H4). Shares the fixtures in
//! [`super::session_bearer_leak_tests`]; see that module's header for the chain
//! and the storage-discipline contract.
use super::session_bearer_leak_tests::{
KIMI_TOKEN, actor_on_managed_model, actor_with_catalog, managed_entry,
};
use super::*;
use kigi_sampler::BearerResolver;
use std::sync::Arc;
#[tokio::test(flavor = "current_thread")]
async fn dual_credential_slug_collision_resolves_the_selected_oauth_platform() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
// (api-key twin, oauth twin, shared slug, host)
for (api_key_twin, oauth_twin, slug, base_url) in [
(
"xai/grok-4.5",
"xai-grok/grok-4.5",
"grok-4.5",
"https://api.x.ai/v1",
),
(
"anthropic/claude-opus-4-8",
"claude-pro-max/claude-opus-4-8",
"claude-opus-4-8",
"https://api.anthropic.com/v1",
),
(
"openai/gpt-5.5-codex",
"openai-codex/gpt-5.5-codex",
"gpt-5.5-codex",
"https://chatgpt.com/backend-api/codex",
),
] {
// API-key platform FIRST, exactly as `PlatformId::ALL` orders them.
let catalog = vec![
managed_entry(api_key_twin, slug, base_url),
managed_entry(oauth_twin, slug, base_url),
];
let (_dir, actor, _rx) = actor_with_catalog(catalog, oauth_twin, "unused").await;
let cfg = actor.reconstruct_full_config().await;
let resolver = cfg.bearer_resolver.as_ref().unwrap_or_else(|| {
panic!(
"{oauth_twin}: selecting the OAuth twin must keep a LIVE bearer_resolver \
(mid-session refresh); resolving {api_key_twin} instead drops it"
)
});
assert_ne!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"{oauth_twin}: the resolver must read its OWN pool, never the Kimi primary"
);
let platform = kigi_models::parse_managed_model_key(oauth_twin)
.expect("managed key")
.0;
assert_eq!(
cfg.anthropic_oauth,
platform.wire_api() == kigi_models::PlatformWireApi::Messages,
"{oauth_twin}: the Claude OAuth Messages adaptation must follow the \
SELECTED platform"
);
assert_eq!(
cfg.openai_codex,
platform.sends_codex_responses_headers(),
"{oauth_twin}: the Codex identity headers must follow the SELECTED platform"
);
assert_eq!(
cfg.github_copilot,
platform.sends_copilot_editor_headers(),
"{oauth_twin}: the Copilot editor headers must follow the SELECTED platform"
);
}
})
.await;
}
/// The other half of H5: selecting the API-KEY twin of a colliding slug must
/// still resolve the API-key platform — no bearer_resolver, no adaptations. The
/// unified lookup must not simply prefer OAuth.
#[tokio::test(flavor = "current_thread")]
async fn dual_credential_slug_collision_resolves_the_selected_api_key_platform() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let catalog = vec![
managed_entry(
"anthropic/claude-opus-4-8",
"claude-opus-4-8",
"https://api.anthropic.com/v1",
),
managed_entry(
"claude-pro-max/claude-opus-4-8",
"claude-opus-4-8",
"https://api.anthropic.com/v1",
),
];
let (_dir, actor, _rx) =
actor_with_catalog(catalog, "anthropic/claude-opus-4-8", "sk-ant-byok").await;
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_none(),
"selecting the API-key twin must get NO session bearer resolver"
);
assert!(
!cfg.anthropic_oauth,
"the API-key Anthropic Messages request must stay byte-identical"
);
assert_eq!(cfg.api_key.as_deref(), Some("sk-ant-byok"));
})
.await;
}
/// MANDATORY counterpart: the subscription-OAuth platforms have NON-first-party
/// base URLs, so the fix must not disable their resolver. Each must still get a
/// LIVE `bearer_resolver` — and it must read THAT platform's own pooled
/// `AuthManager`, never the Kimi primary. The pooled managers are empty here (a
/// TempDir pool home), which is what makes `current_bearer() == None` a proof
/// that the Kimi bearer cannot be what they resolve.
#[tokio::test(flavor = "current_thread")]
async fn oauth_platform_models_keep_a_live_resolver_from_their_own_pool() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
for (catalog_key, slug, base_url) in [
(
"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",
),
(
"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 (_dir, actor, _rx) =
actor_on_managed_model(catalog_key, slug, base_url, "unused").await;
let cfg = actor.reconstruct_full_config().await;
let resolver = cfg.bearer_resolver.as_ref().unwrap_or_else(|| {
panic!("{catalog_key}: must keep a live bearer_resolver for refresh")
});
assert_ne!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"{catalog_key}: the Kimi bearer must never be what it resolves"
);
// The resolver is LIVE over that platform's pooled manager: a
// token rotated inside the pool is observed by the
// already-built resolver (this is what mid-session refresh
// does). The pool is read here, never mutated.
let pooled = crate::auth::oauth_registry::manager_for_model(
&crate::auth::oauth_registry::pool_home(),
catalog_key,
actor.auth_manager.as_ref(),
)
.expect("an OAuth platform always resolves a manager");
assert!(
!Arc::ptr_eq(
&pooled,
actor.auth_manager.as_ref().expect("primary is present")
),
"{catalog_key}: must route to its OWN pooled manager, not the Kimi primary"
);
assert_eq!(
resolver.current_bearer(),
pooled.current_or_expired().map(|a| a.key),
"{catalog_key}: the resolver must read THIS platform's pooled manager"
);
}
})
.await;
}
/// H3/H4 — the stamped AUX paths (image-describe, the auto-mode classifier and
/// the session-summary client all funnel through `aux_bearer_resolver`).
/// `stamp_session_local_sampler_fields` copies the SESSION model's resolver onto
/// every aux config and `SamplingClient::post` REPLACES the request's auth
/// header from it, so an API-key-platform aux model would have its own key
/// overwritten by the Kimi bearer ON THE AUX HOST.
///
/// Revert-to-red: returning `stamped` unconditionally (the pre-fix
/// "re-point only when the aux model is OAuth" shape) makes the deepseek /
/// openai / `[model.*]`-on-openai.com rows resolve `KIMI_TOKEN`.
#[test]
fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host() {
#[derive(Debug)]
struct Fixed(&'static str);
impl BearerResolver for Fixed {
fn current_bearer(&self) -> Option<String> {
Some(self.0.to_string())
}
}
let stamped: kigi_sampler::SharedBearerResolver = Arc::new(Fixed(KIMI_TOKEN));
let platform =
|key: &str| kigi_models::parse_managed_model_key(key).map(|(platform, _)| platform);
// Cleared: every API-key registry platform, and a `[model.*]` aux model
// pointed at a third-party host.
for (key, base_url) in [
("deepseek/deepseek-chat", "https://api.deepseek.com/v1"),
("openai/gpt-5-mini", "https://api.openai.com/v1"),
(
"moonshot-cn/kimi-k2-turbo-preview",
"https://api.moonshot.cn/v1",
),
] {
assert!(
crate::session::acp_session::sampler_turn::aux_bearer_resolver(
Some(stamped.clone()),
platform(key),
base_url,
)
.is_none(),
"LEAK: an aux model on {base_url} must not inherit the session bearer resolver"
);
}
assert!(
crate::session::acp_session::sampler_turn::aux_bearer_resolver(
Some(stamped.clone()),
None,
"https://api.openai.com/v1",
)
.is_none(),
"LEAK: a [model.*] aux model on a third-party host must not inherit it either"
);
// Kept (byte-identical): the first-party subscription channel and a
// platform-less aux model on the session's own endpoint.
for (key, base_url) in [
(
Some("kimi-code/kimi-for-coding"),
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"),
] {
let resolved = crate::session::acp_session::sampler_turn::aux_bearer_resolver(
Some(stamped.clone()),
key.and_then(platform),
base_url,
)
.unwrap_or_else(|| panic!("{key:?} @ {base_url} must keep the session resolver"));
assert_eq!(resolved.current_bearer(), Some(KIMI_TOKEN.to_string()));
}
// Re-pointed: an OAuth aux model gets a LIVE resolver over its OWN pool
// (empty here), never the stamped Kimi one.
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime for the pooled manager's refresh task");
rt.block_on(async {
for key in [
"xai-grok/grok-4-latest",
"claude-pro-max/claude-opus-4-8",
"github-copilot/gpt-4.1",
"openai-codex/gpt-5.5",
] {
let resolved = crate::session::acp_session::sampler_turn::aux_bearer_resolver(
Some(stamped.clone()),
platform(key),
"https://example.invalid/v1",
)
.unwrap_or_else(|| panic!("{key} must keep a live resolver from its own pool"));
assert_ne!(
resolved.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"{key}: the aux resolver must never resolve the Kimi session bearer"
);
}
});
}
@@ -0,0 +1,436 @@
//! LEAK GUARD (bearer_resolver channel): the primary (Kimi) subscription bearer
//! must never be stamped on a request to a host that does not own it.
//!
//! Chain the guard closes: a session-based ACP method (`cached_token` /
//! `kimi-code` / any OAuth platform) + a selected API-key-platform model
//! classifies `ModelByok::NotByok` (the model carries no `[model.*]` key), the
//! 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
//! non-OAuth platform, and `SamplingClient::post` then REPLACED the correctly
//! resolved provider key with the Kimi bearer on the wire.
//!
//! The `api_key` half of the same defect (the config never even gets the
//! provider key, because `resolve_credentials` stamps the session token) is
//! pinned in `agent/mvp_agent/tests/api_key_channel_leak_tests.rs`, which drives
//! the real `prepare_sampling_config_for_model` resolution path. These tests
//! deliberately do NOT hand-stamp a provider key except where the assertion is
//! about the resolver overwriting one that already resolved correctly.
//!
//! The counterpart contract these tests also pin: the four subscription-OAuth
//! platforms have non-first-party base URLs but MUST keep a live
//! `bearer_resolver` drawn from their OWN pooled `AuthManager`, or they lose
//! mid-session token refresh.
//!
//! STORAGE DISCIPLINE (H6): nothing here touches the developer's real `~/.kigi`
//! and nothing hot-swaps the process-global OAuth pool. Under `cfg(test)`
//! `oauth_registry::pool_home()` is a process-lifetime `TempDir`, so every
//! pooled manager is empty — which is exactly what the assertions need (a live
//! resolver that is provably NOT the Kimi one).
use super::support::*;
use super::*;
use crate::agent::auth_method::ModelByok;
use crate::agent::config::{ModelAuthFacts, ModelEntry, ModelInfo};
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
use kigi_sampler::BearerResolver;
use std::sync::Arc;
use tokio::sync::mpsc;
/// The primary session bearer. Any occurrence of this string in an outgoing
/// request to a third-party host is the defect.
pub(super) const KIMI_TOKEN: &str = "kimi-subscription-token-DO-NOT-LEAK";
/// `(tempdir, manager)` standing in for the session's primary Kimi
/// `AuthManager`, holding a live (unexpired) OAuth session bearer.
fn kimi_primary() -> (tempfile::TempDir, Arc<AuthManager>) {
let dir = tempfile::tempdir().expect("tempdir");
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
am.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()
});
(dir, am)
}
/// One catalog entry: catalog key `catalog_key`, routing slug `slug`, routed at
/// `base_url`, carrying no credential of its own (the shape every fetched
/// registry model has).
pub(super) fn managed_entry(catalog_key: &str, slug: &str, base_url: &str) -> (String, ModelEntry) {
let mut info = ModelInfo::fallback(slug);
info.id = Some(catalog_key.to_string());
info.base_url = base_url.to_string();
(
catalog_key.to_string(),
ModelEntry {
info,
api_key: None,
env_key: None,
api_base_url: None,
},
)
}
/// A `SessionActor` on a session-based ACP method with a live Kimi primary,
/// whose live catalog holds `catalog` and whose SELECTED model is the catalog
/// key `selected` (the picker's own notion of "current"). `wire_key` is the
/// already-correctly-resolved provider credential sitting in chat state.
///
/// The per-model BYOK memo is pinned to `NotByok` on purpose: that is what a
/// fetched registry model actually resolves to (`resolve_model_auth_facts` only
/// ever sees `default_models.json` + `[model.*]`), and pinning it keeps the test
/// independent of the developer's on-disk `~/.kigi/config.toml`.
pub(super) async fn actor_with_catalog(
catalog: Vec<(String, ModelEntry)>,
selected: &str,
wire_key: &str,
) -> (
tempfile::TempDir,
Arc<SessionActor>,
mpsc::UnboundedReceiver<PersistenceMsg>,
) {
let (dir, am) = kimi_primary();
let (gateway_tx, _gateway_rx) = mpsc::unbounded_channel();
let (persistence_tx, persistence_rx) = mpsc::unbounded_channel();
let mut actor = create_test_actor(50_000, 200_000, 85, gateway_tx, persistence_tx).await;
actor.auth_manager = Some(am);
actor.auth_method_id = test_auth_method_id("cached_token");
let mut selected_entry = None;
for (key, entry) in catalog {
if key == selected {
selected_entry = Some(entry.clone());
}
actor.models_manager.insert_test_entry(key, entry);
}
let selected_entry = selected_entry.expect("the selected key must be in the catalog");
actor
.models_manager
.set_current_model_id(acp::ModelId::new(selected.to_string()));
let slug = selected_entry.info().model.clone();
actor
.chat_state_handle
.update_sampling_config(kigi_sampling_types::SamplingConfig {
base_url: selected_entry.info().base_url.clone(),
model: slug.clone(),
max_completion_tokens: None,
temperature: None,
top_p: None,
api_backend: Default::default(),
chat_compat: Default::default(),
extra_headers: Default::default(),
context_window: std::num::NonZeroU64::new(200_000).unwrap(),
reasoning_effort: None,
stream_tool_calls: None,
});
actor
.chat_state_handle
.update_credentials(kigi_chat_state::Credentials {
api_key: Some(wire_key.to_string()),
auth_type: kigi_chat_state::AuthType::SessionToken,
..Default::default()
});
actor.model_auth_facts.replace(Some((
slug,
ModelAuthFacts {
byok: ModelByok::NotByok,
auth_scheme: Default::default(),
},
)));
(dir, Arc::new(actor), persistence_rx)
}
/// Single-entry convenience over [`actor_with_catalog`].
pub(super) async fn actor_on_managed_model(
catalog_key: &str,
slug: &str,
base_url: &str,
wire_key: &str,
) -> (
tempfile::TempDir,
Arc<SessionActor>,
mpsc::UnboundedReceiver<PersistenceMsg>,
) {
actor_with_catalog(
vec![managed_entry(catalog_key, slug, base_url)],
catalog_key,
wire_key,
)
.await
}
/// THE leak test, at the wire. A `deepseek/deepseek-chat` turn on a session
/// (`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.
///
/// Revert-to-red: dropping `endpoint_takes_session_credential` from
/// `session_token_auth_gate` puts `Bearer <KIMI_TOKEN>` on this request.
#[tokio::test(flavor = "multi_thread")]
async fn deepseek_turn_under_a_kimi_session_sends_no_kimi_bearer_on_the_wire() {
let server = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("POST"))
.and(wiremock::matchers::path("/chat/completions"))
.respond_with(
wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
"id": "cmpl-1",
"object": "chat.completion",
"created": 0,
"model": "deepseek-chat",
"choices": [{
"index": 0,
"message": { "role": "assistant", "content": "ok" },
"finish_reason": "stop"
}]
})),
)
.mount(&server)
.await;
let uri = server.uri();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (_dir, actor, _rx) = actor_on_managed_model(
"deepseek/deepseek-chat",
"deepseek-chat",
&uri,
"sk-deepseek-provider-key",
)
.await;
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_none(),
"an API-key platform model must get NO session bearer resolver"
);
let client =
kigi_sampler::SamplingClient::new(cfg).expect("sampling client must construct");
let _ = client
.chat_completion(kigi_sampling_types::ChatCompletionRequest::new(
"deepseek-chat",
vec![kigi_sampling_types::ChatRequestMessage::user("hi")],
))
.await;
})
.await;
let requests = server
.received_requests()
.await
.expect("wiremock records requests");
assert_eq!(requests.len(), 1, "exactly one inference request was sent");
let auth = requests[0]
.headers
.get("authorization")
.and_then(|v| v.to_str().ok())
.expect("the request must carry an Authorization header")
.to_string();
assert!(
!auth.contains(KIMI_TOKEN),
"the Kimi subscription bearer must never reach a third-party inference host"
);
assert_eq!(
auth, "Bearer sk-deepseek-provider-key",
"the correctly-resolved provider key must survive to the wire"
);
}
/// The same guard for every other API-key registry platform shape: OpenAI
/// (Responses), Anthropic (x-api-key/Messages), Groq, Together and Z.AI CN — all
/// classify `NotByok`, all route to a non-first-party host, none may receive a
/// session bearer resolver.
#[tokio::test(flavor = "current_thread")]
async fn api_key_platform_models_get_no_session_bearer_resolver() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
for (catalog_key, slug, base_url) in [
("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"),
("together/qwen-3", "qwen-3", "https://api.together.xyz/v1"),
(
"zai-coding-cn/glm-5",
"glm-5",
"https://open.bigmodel.cn/api/paas/v4",
),
] {
let (_dir, actor, _rx) =
actor_on_managed_model(catalog_key, slug, base_url, "sk-provider-key").await;
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_none(),
"{catalog_key}: an API-key platform must get no session bearer resolver"
);
assert_eq!(
cfg.api_key.as_deref(),
Some("sk-provider-key"),
"{catalog_key}: the provider key must stay on the config"
);
}
})
.await;
}
/// C2 at the resolver channel: a `[model.*]` entry has NO platform
/// (`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
/// own coding endpoint (a `KIGI_CODE_BASE_URL` deployment 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`
/// return `true` again puts a Kimi resolver on the openai.com config.
#[tokio::test(flavor = "current_thread")]
async fn config_model_entry_takes_a_session_resolver_only_on_its_own_endpoint() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
for base_url in ["https://api.openai.com/v1", "https://api.deepseek.com/v1"] {
let mut info = ModelInfo::fallback("gpt-4o");
info.id = None; // a `[model.gpt-4o]` block
info.base_url = base_url.to_string();
let entry = ModelEntry {
info,
api_key: None,
// An env_key that is NOT set: `has_own_credentials()` probes
// `std::env::var` at call time, so this classifies NotByok.
env_key: None,
api_base_url: None,
};
let (_dir, actor, _rx) =
actor_with_catalog(vec![("gpt-4o".to_string(), entry)], "gpt-4o", "").await;
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_none(),
"LEAK: a [model.*] block at {base_url} must get no session bearer resolver"
);
}
for base_url in [
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
"http://127.0.0.1:4141/v1",
] {
let mut info = ModelInfo::fallback("kigi-4.5");
info.id = None;
info.base_url = base_url.to_string();
let entry = ModelEntry {
info,
api_key: None,
env_key: None,
api_base_url: None,
};
let (_dir, actor, _rx) =
actor_with_catalog(vec![("kigi-4.5".to_string(), entry)], "kigi-4.5", "").await;
let resolver = actor
.reconstruct_full_config()
.await
.bearer_resolver
.expect("the session's own endpoint keeps the session resolver");
assert_eq!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"{base_url}: a custom deployment / dev proxy is unchanged"
);
}
})
.await;
}
/// H5 — the slug collision. A user holding BOTH an xAI API key and a Grok
/// subscription has `xai/grok-4.5` AND `xai-grok/grok-4.5` in one catalog, in
/// `PlatformId::ALL` order (`Xai`(15) before `XaiGrok`(25)) and with the SAME
/// routing slug. `cfg.model` is that bare slug, so the auth layer used to
/// first-match the API-key entry: no bearer_resolver, no live refresh (the
/// 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.
///
/// Revert-to-red: resolving the platform from `find_model_by_id(models, slug)`
/// 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")]
async fn kimi_first_party_model_still_rides_the_primary_session_bearer() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (_dir, actor, _rx) = actor_on_managed_model(
"kimi-code/kimi-for-coding",
"kimi-for-coding",
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
"stale-buffered-token",
)
.await;
let cfg = actor.reconstruct_full_config().await;
let resolver = cfg
.bearer_resolver
.as_ref()
.expect("the subscription model must keep the live session resolver");
assert_eq!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"the first-party model resolves the primary session bearer"
);
actor.refresh_token_if_expired().await;
assert_eq!(
actor
.chat_state_handle
.get_credentials()
.await
.api_key
.as_deref(),
Some(KIMI_TOKEN),
"the first-party pre-flight refresh must still heal the stale key"
);
})
.await;
}
/// The persistence half of the defect: `refresh_token_if_expired` used to write
/// the Kimi session token into `chat_state` `creds.api_key` for ANY
/// session-method turn, from where it propagated to subagents and aux configs.
/// A deepseek turn must leave the provider key untouched.
///
/// M7 rides along: a registry-platform model must not fall into
/// `reload_api_key_from_config` at all (a `load_effective_config()` disk read
/// per turn plus a permanently false "not found in config.toml" warning), so
/// the key is left exactly as resolved.
#[tokio::test(flavor = "current_thread")]
async fn preflight_refresh_never_writes_the_kimi_token_into_a_platform_credential() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (_dir, actor, _rx) = actor_on_managed_model(
"deepseek/deepseek-chat",
"deepseek-chat",
"https://api.deepseek.com/v1",
"sk-deepseek-provider-key",
)
.await;
actor.refresh_token_if_expired().await;
assert_eq!(
actor
.chat_state_handle
.get_credentials()
.await
.api_key
.as_deref(),
Some("sk-deepseek-provider-key"),
"the Kimi session token must never overwrite a platform credential"
);
})
.await;
}