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