ZacharyZhang-NY/Kigi-CLI
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

refactor(auth): centralize inference-credential routing in CredentialAuthority

Browse Source 
One authority answers 'which credential may ride this request':
credential_class / manager_for / credential_for / bearer_resolver_for,
keyed by (platform, base_url). SessionCredential is an opaque type with
no production constructor, so a new call site cannot re-introduce the
session-bearer leak. Platform-scoped tests extended across all bearer
channels (session, aux, summary, subagent override).

Verified: cargo check --workspace --all-targets clean; kigi-shell and
kigi-tui suites green (6611+ tests).
This commit is contained in:
ZacharyZhang-NY
2026-07-22 15:12:00 -04:00
parent 2d00a4e6e6
commit 48d89c7830
35 changed files with 3158 additions and 1099 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-shell/src/agent/auth_method.rs
+59 -152
View File
@@ -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<kigi_models::PlatformId>,
base_url: &str,
) -> bool {
match platform {
Some(platform) => platform.uses_oauth(),
None => crate::util::is_effective_coding_endpoint_url(base_url),
}
}
pub const AUTH_ERROR_SESSION_EXPIRED: &str =
"Session expired. Run `kigi login` to re-authenticate.";
@@ -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,
crates/codegen/kigi-shell/src/agent/config.rs
+69 -61
View File
@@ -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<ResolvedCredentials> {
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<T>(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<String, ModelEntry>,
endpoints: &EndpointsConfig,
session_key: Option<&str>,
session_key: Option<&SessionCredential>,
alpha_test_key: Option<String>,
) -> Option<SamplerConfig> {
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<kigi_sampler::SharedBearerResolver>,
max_retries: Option<u32>,
) {
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<SamplerConfig>,
active_session_config: &SamplerConfig,
bearer_resolver: Option<kigi_sampler::SharedBearerResolver>,
max_retries: Option<u32>,
) -> (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<String, ModelEntry>,
session_key: Option<&str>,
alpha_test_key: Option<String>,
fallback_entry: Option<ModelEntry>,
) -> Option<SamplerConfig> {
let entry = find_model_by_id(models, model_id)
.cloned()
.or(fallback_entry)?;
let credentials = resolve_credentials(&entry, session_key);
Some(sampling_config_for_model(
&entry,
credentials,
alpha_test_key,
))
}
pub fn to_acp_model_info(
models: &IndexMap<String, ModelEntry>,
) -> IndexMap<acp::ModelId, acp::ModelInfo> {
@@ -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]
crates/codegen/kigi-shell/src/agent/handlers/model_switch.rs
+9
View File
@@ -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,
crates/codegen/kigi-shell/src/agent/models.rs
+190 -45
View File
@@ -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<kigi_models::PlatformId>,
}
// ── 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<String> {
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<String, ModelEntry>,
/// `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<String, ModelEntry>,
current_key: Option<&str>,
slug: &str,
) -> Option<String> {
) -> Option<&'a ModelEntry> {
entry_for_slug_resolution(models, current_key, slug).map(|(entry, _)| entry)
}
/// [`entry_for_slug`] plus whether the SESSION's own selected catalog key is
/// what resolved it. `false` means the entry came from the picker's slug scan —
/// a resolution that is only a GUESS when the slug collides across platforms
/// (H-b).
fn entry_for_slug_resolution<'a>(
models: &'a IndexMap<String, ModelEntry>,
current_key: Option<&str>,
slug: &str,
) -> Option<(&'a ModelEntry, bool)> {
if let Some(entry) = current_key.and_then(|key| models.get(key))
&& (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<String, ModelEntry>, slug: &str) -> bool {
if models.contains_key(slug) {
return false;
}
let mut seen: Option<Option<kigi_models::PlatformId>> = None;
for entry in models.values().filter(|entry| entry.info.model == slug) {
let platform = crate::auth::credential_authority::entry_platform(entry);
match seen {
None => seen = Some(platform),
Some(first) if first == platform => {}
Some(_) => return true,
}
}
false
}
/// The registry platform a routing slug resolves to, via [`entry_for_slug`].
/// `None` for a bare / `[model.*]` / unlisted model. Single definition shared by
/// 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<String, ModelEntry>,
current_key: Option<&str>,
slug: &str,
) -> Option<kigi_models::PlatformId> {
let key = managed_key_for_slug(models, current_key, slug);
kigi_models::parse_managed_model_key(key.as_deref().unwrap_or(slug))
.map(|(platform, _)| platform)
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<String, ModelEntry>,
session_model_id: &acp::ModelId,
) -> Option<String> {
resolve_catalog_key(models, session_model_id).map(|key| key.0.to_string())
}
/// Catalog key for a persisted session model id, restricted to **selectable**
crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs
+38 -9
View File
@@ -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;
crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs
+225 -92
View File
@@ -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<String, ModelEntry>,
slug: &str,
base_url: &str,
) -> Option<kigi_sampler::SharedBearerResolver> {
let auth_method = self.auth_method_id.load();
crate::session::acp_session::sampler_turn::aux_bearer_resolver_for(
&self.credential_authority(),
auth_method.as_deref(),
crate::agent::models::platform_for_slug(models, None, slug),
crate::agent::config::resolve_model_auth_facts(slug).byok,
base_url,
)
}
/// `true` for session-based ACP auth methods.
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 = <Kimi bearer>` with the THIRD-PARTY `base_url`, which
/// `SamplingClient` then builds into `Authorization: Bearer …`. It is
/// reachable with ZERO configuration: `default_models.json` bundles
/// `moonshot-cn/*` + `moonshot-ai/*` entries that a Kimi-subscription user
/// sees on first launch / offline;
/// - the first-party subscription channel (kimi-code, a `KIGI_CODE_BASE_URL`
/// deployment, a loopback proxy) uses the primary session manager, and
/// only under a session-based auth method — byte-identical to the pre-fix
/// path.
/// 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 = <Kimi bearer>` 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<crate::auth::KimiAuth> {
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<crate::auth::credential_authority::SessionCredential> {
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<kigi_models::PlatformId> {
self.sampling_config_platform.get()
}
/// H2/C1: stamp the shared `sampling_config`'s OWN governing session
/// credential — whatever the authority says that is for this config's model
/// and endpoint.
///
/// The shared config is not inert: `resolve_sampling_config_for_model`
/// returns it verbatim whenever a model id fails to resolve, and
/// `SubagentSpawnContext` clones it as every subagent's baseline — so an
/// `api_key` stamped here reaches the wire against whatever `base_url` the
/// config carries. The login/seed sites used to stamp it unconditionally,
/// exactly the mistake `authenticate_oauth_platform` already documents ("Do
/// NOT stamp this token onto the shared sampling_config").
///
/// C1: this function does NOT take a credential. The previous shape took
/// `key: String` — always the primary Kimi bearer — and guarded it with a
/// predicate asking whether **a** session credential may ride. For a
/// subscription-OAuth platform at its own host that is correctly `true`, but
/// the credential that may ride there is that platform's POOLED token: a
/// Claude Pro/Max user running `kigi login` stamped the Kimi subscription
/// bearer onto a config routed at `api.anthropic.com`. Asking
/// `credential_for` instead makes the question and the credential the same
/// object, so the pairing cannot be wrong — and there is no primary handle
/// here to hand-carry (M2).
///
/// `overwrite = false` keeps the historical "only if missing" seeding
/// behaviour; the login handlers pass `true` because a fresh login must
/// replace a stale bearer, and they call this AFTER the manager holds the
/// new token so it is read back through the authority.
///
/// SECURITY: the token is never logged.
pub(super) fn stamp_session_credential(&self, overwrite: bool) -> bool {
let (model, base_url) = {
let sampling_config = self.sampling_config.borrow();
if !overwrite && sampling_config.api_key.is_some() {
return false;
}
(
sampling_config.model.clone(),
sampling_config.base_url.clone(),
)
};
let platform = self.shared_config_platform();
let Some(credential) = self.credential_authority().credential_for(platform, &base_url)
else {
tracing::debug!(
model = model.as_str(),
"auth: no session credential governs the shared sampling config \
(its model + endpoint take none, or the governing provider has no session)"
);
return false;
};
self.sampling_config.borrow_mut().api_key = Some(credential.expose().to_owned());
true
}
/// Whether the shared `sampling_config` may receive a credential that
/// authorizes only the SESSION's own first-party endpoint — the house
/// `KIGI_API_KEY` the `xai.api_key` login handler reads from the
/// environment, which the authority does not own and so cannot produce.
///
/// The house key rides the [`CredentialClass::Primary`] channel and no
/// other. `Pooled` is deliberately excluded: a subscription-OAuth platform's
/// own host DOES take a session credential, but that credential is the
/// platform's pooled token, where the house key has no business (C1).
pub(super) fn shared_config_takes_house_key(&self) -> bool {
let base_url = self.sampling_config.borrow().base_url.clone();
matches!(
self.credential_authority()
.credential_class(self.shared_config_platform(), &base_url),
crate::auth::credential_authority::CredentialClass::Primary
)
}
/// Seed the global sampling config with login auth when available.
///
/// 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(),
crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs
+13
View File
@@ -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<SamplingConfig>,
/// The registry platform [`Self::sampling_config`] was BUILT from, captured
/// in the same `ModelsManager::sampling_config()` call that produced it.
///
/// H-a: the shared config is built ONCE (`Self::with_models`) and never
/// rebuilt, but `ModelsManager::current_model_id()` moves on every
/// non-Leader model switch and on catalog reselection. Its guards
/// (`stamp_session_credential`, `shared_config_takes_house_key`)
/// therefore read THIS cell, never the live one: after a switch the live
/// cell names a different entry, and re-resolving the config's bare slug
/// against it fell through to `resolve_catalog_key`'s `.rev()` scan —
/// answering the API-key twin, which takes no session credential, so a
/// successful `kigi login` silently failed to replace the expired bearer.
pub(crate) sampling_config_platform: std::cell::Cell<Option<kigi_models::PlatformId>>,
pub(crate) auth_manager: Arc<AuthManager>,
pub(crate) models_manager: crate::agent::models::ModelsManager,
/// Forwards pasted codes from `handle_auth_submit_code` to the auth flow.
crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs
+1
View File
@@ -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]
crates/codegen/kigi-shell/src/agent/mvp_agent/tests/api_key_channel_leak_tests.rs
+9 -8
View File
@@ -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 <Kimi OAuth token>` to `api.moonshot.cn`, which is NOT
/// first-party.
///
/// Revert-to-red: dropping the `platform_takes_session_credential` guard from
/// `session_token_for_model` makes every `api_key` below `Some(KIMI_TOKEN)`.
/// 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() {
crates/codegen/kigi-shell/src/agent/mvp_agent/tests/chokepoint_leak_tests.rs
+589
View File
@@ -0,0 +1,589 @@
//! LEAK GUARD (the chokepoint's own call sites) — H1, H2 and the H3 regression.
//!
//! The companion of `api_key_channel_leak_tests`, which drives
//! `MvpAgent::prepare_sampling_config_for_model`. These three pin the OTHER
//! producers of an outgoing credential: `ModelsManager::sampling_config()` (the
//! agent-wide baseline), the shared `MvpAgent::sampling_config` the login/seed
//! paths stamp, and the SESSION's effective coding endpoint as configured from
//! config.toml rather than the environment.
//!
//! Every assertion reads what the real resolution path produced; nothing is
//! hand-stamped.
use super::super::*;
use super::api_key_channel_leak_tests::{
KIMI_TOKEN, kimi_session_agent, platform_entry, without_ambient_byok_env,
};
use crate::agent::auth_method::CACHED_TOKEN_AUTH_METHOD_ID;
use crate::agent::config::{Config as AgentConfig, EndpointsConfig, ModelEntry};
use crate::auth::credential_authority::CredentialClass;
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
use kigi_sampler::BearerResolver;
use kigi_test_support::EnvGuard;
/// Re-seed the shared config the way `MvpAgent::with_models` does: the config
/// AND the platform it was BUILT from, from ONE `ModelsManager::sampling_config()`
/// against whatever `current_model_id` names right now. A test can therefore
/// never set one without the other — which is the whole point of H-a.
fn rebuild_shared_config(agent: &MvpAgent) {
let baseline = agent.models_manager.sampling_config();
agent.sampling_config_platform.set(baseline.platform);
*agent.sampling_config.borrow_mut() = baseline.config;
}
/// H1 — `ModelsManager::sampling_config()`, the OTHER `api_key` producer, and a
/// byte-for-byte repeat of the round-1 defect: it resolved the session bearer
/// itself (`platform.oauth()`, else `auth_manager.current_or_expired()`), so
/// every non-OAuth platform got the primary Kimi token.
///
/// This config is not incidental — it is the `MvpAgent` baseline
/// (`Self::with_models`), which `resolve_sampling_config_for_model` returns
/// verbatim for an unresolved model id and `SubagentSpawnContext` clones as
/// every subagent's baseline, so its `api_key` reaches the wire against its own
/// `base_url`. Zero-config repro: a Kimi subscription + a bundled
/// `moonshot-cn/*` default.
///
/// Revert-to-red (L: this edit COMPILES — the previous wording named a
/// `Option<String>` argument that the `Option<&SessionCredential>` signature
/// rejects, so it could never have been run): in
/// `ModelsManager::sampling_config`, ask the authority about the SESSION's
/// endpoint instead of the current model's own —
/// `.credential_for(None, &config.endpoints.proxy_url())` in place of
/// `.credential_for_model(current_model)`. That is the round-1 defect's shape
/// (the credential decided by something other than the model's own platform +
/// endpoint) and every `assert_ne!` below sees `Some(KIMI_TOKEN)`.
#[tokio::test]
#[serial_test::serial]
async fn models_manager_sampling_config_never_carries_the_kimi_bearer() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
for (catalog_key, slug, base_url) in [
(
"moonshot-cn/kimi-k2-turbo-preview",
"kimi-k2-turbo-preview",
"https://api.moonshot.cn/v1",
),
(
"deepseek/deepseek-chat",
"deepseek-chat",
"https://api.deepseek.com/v1",
),
("openai/gpt-5.2", "gpt-5.2", "https://api.openai.com/v1"),
] {
agent
.models_manager
.insert_test_entry(catalog_key, platform_entry(catalog_key, slug, base_url));
agent
.models_manager
.set_current_model_id(acp::ModelId::new(catalog_key));
let cfg = agent.models_manager.sampling_config().config;
assert_eq!(cfg.base_url, base_url, "{catalog_key}: routed to its own host");
assert_ne!(
cfg.api_key.as_deref(),
Some(KIMI_TOKEN),
"LEAK: ModelsManager::sampling_config sent the Kimi bearer to {base_url}"
);
}
// …and the first-party subscription channel is unchanged, which is what
// proves the Kimi bearer was reachable above.
let kimi_key = "kimi-code/kimi-for-coding";
agent.models_manager.insert_test_entry(
kimi_key,
platform_entry(
kimi_key,
"kimi-for-coding",
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
),
);
agent
.models_manager
.set_current_model_id(acp::ModelId::new(kimi_key));
assert_eq!(
agent
.models_manager
.sampling_config()
.config
.api_key
.as_deref(),
Some(KIMI_TOKEN),
"the kimi-code subscription channel must be byte-identical"
);
}
/// H2 — the SHARED `MvpAgent::sampling_config`. `seed_client_config_auth_if_available`
/// (from `new_session` / `load_session`) and the `cached_token` / `kimi.com/oidc`
/// login handlers all stamped `sampling_config.api_key = Some(<primary bearer>)`
/// with NO platform or endpoint guard — while that very config may point at a
/// third-party model, and is both the subagent baseline and the
/// unresolved-model fallback. `agent_ops`' generic-OAuth login handler already
/// documents the correct rule ("Do NOT stamp this token onto the shared
/// sampling_config"); the Kimi handlers violated it.
///
/// Revert-to-red: make `stamp_session_credential` skip the authority entirely —
/// `self.sampling_config.borrow_mut().api_key =
/// self.auth_manager.current_or_expired().map(|a| a.key); return true;` — and
/// the third-party rows below become `Some(KIMI_TOKEN)`.
#[tokio::test]
#[serial_test::serial]
async fn shared_sampling_config_is_never_stamped_off_the_session_endpoint() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
for (catalog_key, slug, base_url) in [
(
"deepseek/deepseek-chat",
"deepseek-chat",
"https://api.deepseek.com/v1",
),
(
"moonshot-cn/kimi-k2-turbo-preview",
"kimi-k2-turbo-preview",
"https://api.moonshot.cn/v1",
),
] {
agent
.models_manager
.insert_test_entry(catalog_key, platform_entry(catalog_key, slug, base_url));
agent
.models_manager
.set_current_model_id(acp::ModelId::new(catalog_key));
rebuild_shared_config(&agent);
{
let mut shared = agent.sampling_config.borrow_mut();
assert_eq!(shared.model, slug, "{catalog_key}: built from this entry");
assert_eq!(shared.base_url, base_url);
shared.api_key = None;
}
// The `new_session` / `load_session` seed…
agent.seed_client_config_auth_if_available();
assert_eq!(
agent.sampling_config.borrow().api_key, None,
"LEAK: seeding stamped the Kimi bearer onto a config routed at {base_url}"
);
// …and the login handlers, which overwrite rather than seed.
assert!(
!agent.stamp_session_credential(true),
"LEAK: a login handler stamped the Kimi bearer onto a config routed at {base_url}"
);
assert_eq!(agent.sampling_config.borrow().api_key, None);
}
// Byte-identical on the session's own endpoint: both paths still stamp.
let kimi_key = "kimi-code/kimi-for-coding";
agent.models_manager.insert_test_entry(
kimi_key,
platform_entry(
kimi_key,
"kimi-for-coding",
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
),
);
agent
.models_manager
.set_current_model_id(acp::ModelId::new(kimi_key));
rebuild_shared_config(&agent);
agent.sampling_config.borrow_mut().api_key = None;
agent.seed_client_config_auth_if_available();
assert_eq!(
agent.sampling_config.borrow().api_key.as_deref(),
Some(KIMI_TOKEN),
"the subscription endpoint must still be seeded (this is what makes the \
assertions above meaningful)"
);
}
/// C1 — THE CRITICAL. The login stamp used to pair the right QUESTION (*may
/// **a** session credential ride here?*) with the wrong CREDENTIAL (always
/// `auth_manager.current_or_expired().key`, the primary Kimi bearer). For a
/// subscription-OAuth platform at its OWN host the answer is correctly "yes" —
/// the credential that may ride there is that platform's POOLED token
/// ([`CredentialClass::Pooled`]) — so a Claude Pro/Max user whose current model is
/// `claude-pro-max/*` ran `kigi login` and the Kimi subscription bearer landed
/// on a config routed at `api.anthropic.com`. From there it reaches the wire via
/// `resolve_sampling_config_for_model`'s verbatim fallback (offline / stale
/// catalog) and via `SubagentSpawnContext`'s baseline clone.
///
/// All FOUR subscription platforms, each at its own registry host, derived from
/// the registry so the fixture cannot drift.
///
/// The pooled managers are empty here (`pool_home()` is a per-process temp path
/// under `cfg(test)`), so the correct answer is `None` — and `None` is also what
/// proves the primary is not being substituted, because the same agent DOES
/// stamp `KIMI_TOKEN` on its own coding endpoint at the end of the test.
///
/// Revert-to-red (production, compiles): restore the old shape in
/// `MvpAgent::stamp_session_credential` —
/// ```ignore
/// if self.credential_authority().credential_class(platform, &base_url)
/// == CredentialClass::None
/// {
/// return false;
/// }
/// let Some(auth) = self.auth_manager.current_or_expired() else { return false };
/// self.sampling_config.borrow_mut().api_key = Some(auth.key);
/// true
/// ```
/// and every OAuth row below becomes `Some(KIMI_TOKEN)`.
#[tokio::test]
#[serial_test::serial]
async fn oauth_platform_shared_config_never_receives_the_primary_on_login() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
for (platform_id, slug) in [
("claude-pro-max", "claude-opus-4-8"),
("openai-codex", "gpt-5.5-codex"),
("github-copilot", "gpt-4.1"),
("xai-grok", "grok-4.5"),
] {
let platform = kigi_models::PlatformId::parse(platform_id).expect("known platform");
let base_url = platform.base_url();
let catalog_key = format!("{platform_id}/{slug}");
agent
.models_manager
.insert_test_entry(&catalog_key, platform_entry(&catalog_key, slug, &base_url));
agent
.models_manager
.set_current_model_id(acp::ModelId::new(catalog_key.clone()));
rebuild_shared_config(&agent);
{
let mut shared = agent.sampling_config.borrow_mut();
assert_eq!(shared.model, slug, "{catalog_key}: built from this entry");
assert_eq!(shared.base_url, base_url);
shared.api_key = None;
}
// Precondition: this endpoint DOES take a session credential — that is
// exactly why guarding a hand-carried primary with "may a session
// credential ride?" was the bug. The class names WHICH one: the
// platform's POOLED token, never the primary / house key.
assert_eq!(
agent
.credential_authority()
.credential_class(Some(platform), &base_url),
CredentialClass::Pooled,
"{catalog_key}: precondition — its own host takes its POOLED token, \
and never the primary / house credential"
);
// `kigi login` (cached_token and kimi.com/oidc both land here) …
assert!(
!agent.stamp_session_credential(true),
"LEAK: a Kimi login stamped a credential onto a config routed at {base_url}"
);
assert_eq!(
agent.sampling_config.borrow().api_key,
None,
"LEAK: {catalog_key} received a bearer that is not its own pooled token"
);
// … and the `new_session` / `load_session` seed.
agent.seed_client_config_auth_if_available();
assert_ne!(
agent.sampling_config.borrow().api_key.as_deref(),
Some(KIMI_TOKEN),
"LEAK: seeding sent the Kimi subscription bearer to {base_url}"
);
}
// The first-party channel is untouched — this is what makes every
// assertion above meaningful (the Kimi bearer IS live and IS stampable).
let kimi_key = "kimi-code/kimi-for-coding";
agent.models_manager.insert_test_entry(
kimi_key,
platform_entry(
kimi_key,
"kimi-for-coding",
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
),
);
agent
.models_manager
.set_current_model_id(acp::ModelId::new(kimi_key));
rebuild_shared_config(&agent);
agent.sampling_config.borrow_mut().api_key = None;
assert!(agent.stamp_session_credential(true));
assert_eq!(
agent.sampling_config.borrow().api_key.as_deref(),
Some(KIMI_TOKEN),
"the Kimi subscription channel must be byte-identical"
);
}
/// H-a (AVAILABILITY, HIGH) — the stamp guard must resolve the model the shared
/// config was BUILT from, not the one `current_model_id()` names NOW.
///
/// The shared config is built ONCE (`MvpAgent::with_models`) and never rebuilt,
/// while `current_model_id()` moves on every non-Leader model switch
/// (`handlers/model_switch.rs`) and on catalog reselection. Once they drift, a
/// guard re-resolving the config's BARE slug against the live cell fails
/// `entry_for_slug_resolution`'s `entry.info.model == slug` test and falls
/// through to `resolve_catalog_key`'s `.rev()` scan.
///
/// `kimi-code` (subscription, `uses_oauth`) and `kimi-coding` (API-key twin,
/// SAME coding host, `uses_oauth: false`) list the same routing slug, and
/// `PlatformId::ALL` puts `kimi-code` 1st and `kimi-coding` 19th — so that scan
/// answers `kimi-coding`, which takes NO session credential and therefore has NO
/// governing manager. Because `stamp_session_credential` only overwrites ON
/// SUCCESS, a Kimi-subscription user who also has `KIMI_API_KEY` set kept the
/// EXPIRED bearer in the shared config after a successful re-login: every
/// unresolved-model fallback and every subagent baseline turn 401'd until
/// restart.
///
/// The switch happens AFTER the config is built — the previous version of this
/// test left both on the same model, so it could not catch this.
///
/// Revert-to-red (production, compiles): make `MvpAgent::shared_config_platform`
/// re-resolve from the live cell instead of returning the captured value —
/// ```ignore
/// fn shared_config_platform(&self) -> Option<kigi_models::PlatformId> {
/// let model = self.sampling_config.borrow().model.clone();
/// let current_key = self.models_manager.current_model_id();
/// crate::agent::models::platform_for_slug(
/// &self.models_manager.models(),
/// Some(current_key.0.as_ref()),
/// &model,
/// )
/// }
/// ```
/// and the re-login assertion below sees the stale bearer.
#[tokio::test]
#[serial_test::serial]
async fn relogin_restamps_the_shared_config_after_the_model_switched() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
let coding_host = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
let slug = "kimi-for-coding";
// Insertion order mirrors `PlatformId::ALL`: the subscription platform
// first, its API-key twin later — so the `.rev()` scan answers the twin.
for catalog_key in ["kimi-code/kimi-for-coding", "kimi-coding/kimi-for-coding"] {
agent
.models_manager
.insert_test_entry(catalog_key, platform_entry(catalog_key, slug, coding_host));
}
assert_eq!(
crate::agent::models::platform_for_slug(&agent.models_manager.models(), None, slug),
Some(kigi_models::PlatformId::KimiCoding),
"precondition: a `None`-keyed slug scan answers the API-key twin, which takes \
no session credential"
);
// Startup: the picker selected the SUBSCRIPTION entry, and the shared config
// was built from it (config + platform, one call).
agent
.models_manager
.set_current_model_id(acp::ModelId::new("kimi-code/kimi-for-coding"));
rebuild_shared_config(&agent);
assert_eq!(
agent.sampling_config.borrow().model,
slug,
"precondition: the shared config carries the BARE slug, which collides"
);
// The user switches model: `current_model_id` moves to the API-key twin
// while the shared config still represents the subscription entry.
agent
.models_manager
.set_current_model_id(acp::ModelId::new("kimi-coding/kimi-for-coding"));
// The session expires and `kigi login` re-mints it. The handlers overwrite
// (`stamp_session_credential(true)`) AFTER the manager holds the new token.
agent.sampling_config.borrow_mut().api_key = Some("expired-bearer".to_string());
assert!(
agent.stamp_session_credential(true),
"a successful re-login must restamp the shared config"
);
assert_eq!(
agent.sampling_config.borrow().api_key.as_deref(),
Some(KIMI_TOKEN),
"the re-login must REPLACE the expired bearer: resolving the API-key twin \
instead finds no governing manager, leaves the stale token in place, and \
401s every subagent baseline turn until restart"
);
// And the seed path (`new_session` / `load_session`) agrees.
agent.sampling_config.borrow_mut().api_key = None;
agent.seed_client_config_auth_if_available();
assert_eq!(
agent.sampling_config.borrow().api_key.as_deref(),
Some(KIMI_TOKEN),
);
}
/// H3 (REGRESSION) — a MANAGED deployment configures its coding endpoint with
/// `[endpoints] coding_api_base_url` in **config.toml** (what the managed-config
/// sync writes), NOT the `KIGI_CODE_BASE_URL` env var. The previous round's
/// predicate knew only the env var, so `EndpointsConfig::proxy_url()`'s
/// config-key branch was invisible: every model inheriting the managed endpoint
/// classified third-party, lost its api_key AND its resolver, and 401'd on every
/// turn.
///
/// NOTE: no env var is set anywhere in this test — that is the point.
///
/// Revert-to-red: drop the `proxy_url()` arm from
/// `CredentialAuthority::is_session_coding_endpoint` and both `assert_eq!`s
/// below become `None`.
#[tokio::test]
#[serial_test::serial]
async fn managed_config_toml_coding_endpoint_keeps_the_session_bearer() {
let _env = without_ambient_byok_env();
let _no_env_override = EnvGuard::unset("KIGI_CODE_BASE_URL");
let managed = "https://proxy.acme.example/v1";
let dir = tempfile::tempdir().expect("tempdir");
let auth_manager = std::sync::Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
auth_manager.hot_swap(KimiAuth {
key: KIMI_TOKEN.to_string(),
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..KimiAuth::test_default()
});
let cfg = AgentConfig {
endpoints: EndpointsConfig {
coding_api_base_url: Some(managed.to_string()),
..EndpointsConfig::default()
},
..AgentConfig::default()
};
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let agent = MvpAgent::new(GatewaySender::new(tx), &cfg, auth_manager, None)
.expect("valid test config");
agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID));
assert_eq!(
agent.models_manager.endpoints().proxy_url(),
managed,
"precondition: the session's effective coding endpoint is the config.toml key"
);
// A `[model.*]` entry inheriting the managed endpoint …
let mut bare = ModelEntry::fallback("kigi-4.5", &cfg.endpoints);
bare.info.id = None;
bare.info.base_url = managed.to_string();
assert_eq!(
agent
.prepare_sampling_config_for_model(&bare, None)
.api_key
.as_deref(),
Some(KIMI_TOKEN),
"a managed deployment must still receive the session bearer"
);
// … and the kimi-code catalog entry, whose base_url IS `proxy_url()`.
let kimi = platform_entry("kimi-code/kimi-for-coding", "kimi-for-coding", managed);
assert_eq!(
agent
.prepare_sampling_config_for_model(&kimi, None)
.api_key
.as_deref(),
Some(KIMI_TOKEN),
"the managed kimi-code entry must still receive the session bearer"
);
// The spec's requirement was "rides AND still refreshes": an api_key alone
// freezes at login and 401s unrecoverably ~1h in. The managed endpoint must
// also keep a LIVE manager — the primary's, so mid-session refresh and 401
// recovery run against the credential that actually owns that host.
for platform in [None, Some(kigi_models::PlatformId::KimiCode)] {
let manager = agent
.credential_authority()
.manager_for(platform, managed)
.unwrap_or_else(|| {
panic!("{platform:?}: a managed deployment must keep a live manager")
});
assert!(
std::sync::Arc::ptr_eq(&manager, &agent.auth_manager),
"{platform:?}: and it must be the session's OWN primary manager"
);
let resolver = agent
.credential_authority()
.bearer_resolver_for(platform, managed)
.unwrap_or_else(|| panic!("{platform:?}: … exposed as a live bearer_resolver"));
assert_eq!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"{platform:?}: the resolver reads the live primary session bearer"
);
}
// The guard still holds: a third-party host under the SAME managed config
// gets nothing — no key, and no resolver either.
let third_party = platform_entry(
"deepseek/deepseek-chat",
"deepseek-chat",
"https://api.deepseek.com/v1",
);
assert_eq!(
agent
.prepare_sampling_config_for_model(&third_party, None)
.api_key,
None,
"LEAK: a managed deployment must not widen the trust set to third parties"
);
assert!(
agent
.credential_authority()
.bearer_resolver_for(
Some(kigi_models::PlatformId::DeepSeek),
"https://api.deepseek.com/v1"
)
.is_none(),
"LEAK: nor hand it a live resolver over the primary"
);
}
/// M3 (COMPLETION) — the SUMMARY client's `bearer_resolver` must honour the
/// session-token gate, exactly as the session actor's aux path does.
///
/// `build_summary_client` set `cfg.bearer_resolver =
/// authority.bearer_resolver_for(platform_for_slug(…), &cfg.base_url)` with NO
/// gate while `SessionActor::aux_bearer_resolver` had one. `SamplingClient::post`
/// REPLACES the request's auth header from that resolver, so an api-key /
/// house-key session whose `[model.session-summary]` block carries its OWN
/// `env_key` on the session's own coding endpoint had that key overwritten by
/// the primary bearer on every summary request. Both now go through ONE rule,
/// `sampler_turn::aux_bearer_resolver_for`.
///
/// The summary slug is deliberately absent from the catalog and from any
/// config: it classifies `NotByok` definitively, so the only variable left is
/// the ACP auth method — which is the gate term under test.
///
/// Revert-to-red (production, compiles): in `MvpAgent::summary_bearer_resolver`,
/// return `self.credential_authority().bearer_resolver_for(platform, base_url)`
/// directly (the pre-fix shape) and the api-key row below resolves `KIMI_TOKEN`.
#[tokio::test]
#[serial_test::serial]
async fn summary_client_resolver_honours_the_session_gate() {
let _env = without_ambient_byok_env();
let (_dir, agent) = kimi_session_agent();
let coding_host = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
let slug = "kigi-summary-aux-not-in-any-catalog";
let models = agent.models_manager.models();
// A session-based method on the session's OWN endpoint: byte-identical, the
// summary model keeps a LIVE resolver over the primary.
agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID));
let resolver = agent
.summary_bearer_resolver(&models, slug, coding_host)
.expect("the first-party subscription summary channel keeps its resolver");
assert_eq!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"…and it reads the live primary session bearer"
);
// An API-KEY session: the gate is inactive, so the summary model's own key
// must survive to the wire instead of being replaced by the primary.
agent.set_auth_method(acp::AuthMethodId::new(
crate::agent::auth_method::XAI_API_KEY_METHOD_ID,
));
assert!(
agent
.summary_bearer_resolver(&models, slug, coding_host)
.is_none(),
"LEAK: an api-key session's summary model had its own key replaced on the \
wire by the primary bearer"
);
}
crates/codegen/kigi-shell/src/agent/subagent/mod.rs
+6 -11
View File
@@ -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());
crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs
+3 -2
View File
@@ -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");
crates/codegen/kigi-shell/src/auth/credential_authority.rs
+568
View File
@@ -0,0 +1,568 @@
//! THE credential chokepoint.
//!
//! One authority answers, for every outgoing inference request, the only
//! question that matters: **which credential — if any — may ride it?**
//! ([`CredentialAuthority::credential_class`]). Before this module the answer was
//! re-derived — differently — at every call site (`ModelsManager`, `MvpAgent`,
//! `SessionActor`, the aux/summary/subagent paths), and three separate rounds
//! of fixes each closed some sites and missed others.
//!
//! # How omission is structurally prevented
//!
//! 1. [`SessionCredential`] wraps the bearer and has **no production
//! constructor outside this module**. The only function in the crate that
//! can build one is [`CredentialAuthority::credential_for`], which *requires*
//! `(platform, base_url)` and holds the session's `EndpointsConfig` and
//! primary [`AuthManager`] privately.
//! 2. Every API that stamps a session credential onto a request —
//! `resolve_credentials`, `resolve_aux_model_sampling_config`,
//! `try_resolve_model_credentials`,
//! `resolve_chat_state_auth_type` — takes `Option<&SessionCredential>`,
//! never `Option<&str>`. A new call site therefore *cannot compile* a leak:
//! there is no way to produce the value without going through the rule.
//! 3. The authority owns the primary manager privately and exposes it only via
//! [`CredentialAuthority::manager_for`] /
//! [`CredentialAuthority::bearer_resolver_for`], which take the same
//! `(platform, base_url)` pair — so the `bearer_resolver` sink is funnelled
//! through the identical rule as the `api_key` sink.
//! 4. A guard asks [`CredentialAuthority::credential_class`] and MATCHES on the
//! answer. There is no second, similarly-named boolean to pick by mistake:
//! the round-3 defect (C1) was `takes_session_credential` — *may **a**
//! session credential ride?* — paired with a hand-carried PRIMARY bearer,
//! and the two predicates that made that pairing expressible are gone.
//!
//! SECURITY: no token is ever logged, `Debug`-printed or `Display`ed here.
use std::sync::Arc;
use crate::agent::config::{EndpointsConfig, ModelEntry};
use crate::auth::AuthManager;
/// A session bearer this authority has cleared for one specific request
/// endpoint.
///
/// Opaque by construction: the inner `String` is private, the type is not
/// `Debug`/`Clone`-into-`String`, and the only production constructor is
/// [`CredentialAuthority::credential_for`]. See the module docs for why that
/// matters.
pub(crate) struct SessionCredential(String);
impl SessionCredential {
/// The raw bearer. SECURITY: callers stamp this straight onto a request —
/// never log it.
pub(crate) fn expose(&self) -> &str {
&self.0
}
/// Test-only forgery, so unit tests can exercise the *downstream*
/// credential plumbing (`resolve_credentials`' BYOK-vs-session precedence,
/// aux config shapes) without standing up an `AuthManager`. Deliberately
/// `#[cfg(test)]`: production code has no way to build one.
#[cfg(test)]
pub(crate) fn for_test(key: &str) -> Self {
Self(key.to_owned())
}
}
/// WHICH credential — if any — may ride a request routed to a given
/// `(platform, base_url)` pair.
///
/// ONE question with three answers, replacing the two look-alike booleans
/// `takes_session_credential` / `takes_primary_credential` (identical
/// signatures, near-identical names, opposite answers on a subscription host).
/// C1 was caused by asking the first and stamping the credential the second
/// describes; with a single classifier a call site must MATCH on the answer, so
/// that mistake is no longer expressible.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum CredentialClass {
/// `platform`'s OWN pooled subscription-OAuth token, at its own registry
/// host. NEVER the primary session bearer and never the house key.
Pooled,
/// The credential that authorizes the SESSION's own coding endpoint:
/// the primary (`kimi-code` / platform-less) bearer.
///
/// Deliberately NOT split into a separate `HouseKey` variant: the house
/// `KIGI_API_KEY` is accepted by exactly this endpoint and no other, so it
/// rides precisely this class. A fourth variant would re-create the
/// two-similar-answers hazard this enum exists to remove.
Primary,
/// Nothing rides: every API-key registry platform, an OAuth platform
/// redirected off its own host, and any endpoint that is not the session's.
None,
}
impl CredentialClass {
/// Stable label for structured logs. SECURITY: names a channel, never a
/// token.
pub(crate) fn as_str(self) -> &'static str {
match self {
Self::Pooled => "pooled",
Self::Primary => "primary",
Self::None => "none",
}
}
}
/// The single authority over inference-time session credentials.
///
/// Construct one from the session's EFFECTIVE endpoints plus its primary
/// (first-party / Kimi) manager, then ask it about a request. Cheap to build
/// (a handful of `Option<String>` clones + an `Arc` clone).
#[derive(Clone)]
pub(crate) struct CredentialAuthority {
/// The session's effective `[endpoints]` — config.toml layered over env.
/// H3: `EndpointsConfig::proxy_url()` prefers `[endpoints]
/// coding_api_base_url` from **config.toml** and only then falls back to
/// `KIGI_CODE_BASE_URL`. A predicate that knows only the env var makes a
/// managed/enterprise deployment lose its session bearer entirely (401 on
/// every turn), so the endpoints are part of the authority's identity, not
/// an afterthought.
endpoints: EndpointsConfig,
/// The primary session manager. PRIVATE: nothing hands it back, so a path
/// holding a `CredentialAuthority` cannot reach `current_or_expired()`
/// without naming an endpoint.
primary: Option<Arc<AuthManager>>,
}
impl CredentialAuthority {
pub(crate) fn new(endpoints: EndpointsConfig, primary: Option<Arc<AuthManager>>) -> Self {
Self { endpoints, primary }
}
/// THE rule, stated once.
///
/// - a subscription-OAuth platform (claude-pro-max, openai-codex,
/// github-copilot, xai-grok) rides ITS OWN pooled manager — never the
/// primary — and only to its own registry host (L10: a
/// `[model."claude-pro-max/x"]` override keeps `info.id` but can point
/// `base_url` anywhere, and used to ship the Claude OAuth bearer there);
/// - `kimi-code` — the one `uses_oauth` platform with no `OAuthConfig` —
/// rides the PRIMARY session, and only at the session's own effective
/// coding endpoint;
/// - every API-key registry platform (deepseek, openai, anthropic,
/// moonshot-*, …) rides NOTHING: its credential is that platform's API
/// key, already resolved into the catalog entry;
/// - a platform-less model (a bare slug or a `[model.*]` block) is decided
/// purely by the ENDPOINT — BYOK detection probes `std::env::var` at call
/// time, so an unset/mistyped `env_key` must not turn into "send the
/// subscription bearer to `api.openai.com`".
pub(crate) fn credential_class(
&self,
platform: Option<kigi_models::PlatformId>,
base_url: &str,
) -> CredentialClass {
match platform {
Some(platform) => match platform.oauth() {
Some(_) if self.endpoint_is_platform_host(platform, base_url) => {
CredentialClass::Pooled
}
// An OAuth platform pointed at a host that is NOT its own.
Some(_) => CredentialClass::None,
None if platform.uses_oauth() && self.is_session_coding_endpoint(base_url) => {
CredentialClass::Primary
}
// `kimi-code` off the session's endpoint, and every API-key
// registry platform.
None => CredentialClass::None,
},
None if self.is_session_coding_endpoint(base_url) => CredentialClass::Primary,
None => CredentialClass::None,
}
}
/// The manager behind [`Self::credential_class`]. Derived from the class, so
/// the rule is stated exactly once and the two can never disagree.
fn governing_manager(
&self,
platform: Option<kigi_models::PlatformId>,
base_url: &str,
) -> Option<Arc<AuthManager>> {
match self.credential_class(platform, base_url) {
CredentialClass::Pooled => {
platform
.and_then(kigi_models::PlatformId::oauth)
.map(|oauth| {
crate::auth::oauth_registry::global_manager_for(
&crate::auth::oauth_registry::pool_home(),
oauth,
)
})
}
CredentialClass::Primary => self.primary.clone(),
CredentialClass::None => None,
}
}
/// Whether `base_url` is the SESSION's own coding endpoint: the effective
/// `[endpoints] coding_api_base_url` from **config.toml** (what a managed /
/// enterprise deployment actually sets — H3), the `models_base_url`
/// custom-endpoint mode, the `KIGI_CODE_BASE_URL` env override, a loopback
/// dev proxy, or the compiled production endpoint.
///
/// Deliberately NOT [`crate::util::is_first_party_url`], which is
/// production-only and would break every custom deployment.
fn is_session_coding_endpoint(&self, base_url: &str) -> bool {
if crate::util::is_effective_coding_endpoint_url(base_url) {
return true;
}
if crate::util::matches_trusted_base_url(base_url, &self.endpoints.proxy_url()) {
return true;
}
self.endpoints
.models_base_url
.as_deref()
.is_some_and(|models_base| crate::util::matches_trusted_base_url(base_url, models_base))
}
/// Whether `base_url` is `platform`'s own registry host — the guard that
/// keeps a subscription-OAuth bearer from riding a redirected `[model.*]`
/// override to a third party (L10).
fn endpoint_is_platform_host(&self, platform: kigi_models::PlatformId, base_url: &str) -> bool {
crate::util::matches_trusted_base_url(base_url, &platform.base_url())
}
/// The `AuthManager` that governs this request's bearer resolution,
/// mid-session refresh and 401 recovery — or `None` when no session
/// credential may ride (fail fast; never a silent fallback to the primary).
pub(crate) fn manager_for(
&self,
platform: Option<kigi_models::PlatformId>,
base_url: &str,
) -> Option<Arc<AuthManager>> {
self.governing_manager(platform, base_url)
}
/// The session bearer to stamp as this request's `api_key`, or `None`.
///
/// The ONLY production constructor of [`SessionCredential`].
pub(crate) fn credential_for(
&self,
platform: Option<kigi_models::PlatformId>,
base_url: &str,
) -> Option<SessionCredential> {
self.governing_manager(platform, base_url)
.and_then(|am| am.current_or_expired())
.map(|auth| SessionCredential(auth.key))
}
/// A live sampler `bearer_resolver` over the governing manager, so the
/// request keeps mid-session refresh / 401 recovery against the credential
/// that actually belongs to its host.
pub(crate) fn bearer_resolver_for(
&self,
platform: Option<kigi_models::PlatformId>,
base_url: &str,
) -> Option<kigi_sampler::SharedBearerResolver> {
self.manager_for(platform, base_url)
.map(crate::session::acp_session::sampler_turn::auth_manager_bearer_resolver)
}
/// [`Self::credential_for`] for a resolved catalog entry: derives the
/// platform and the base URL from the SAME entry, so the two can never be
/// mismatched by a call site.
pub(crate) fn credential_for_model(&self, entry: &ModelEntry) -> Option<SessionCredential> {
let info = entry.info();
self.credential_for(entry_platform(entry), &info.base_url)
}
/// [`Self::credential_for`] for the catalog model a routing slug resolves
/// to. `current_key` is the SESSION's own selected catalog key (see
/// [`crate::agent::models::entry_for_slug`]); pass `None` for aux /
/// override slugs, which are not the session's selection.
///
/// M5: a slug that is NOT in the catalog resolves through the SAME endpoint
/// rule against the aux fallback endpoint
/// (`EndpointsConfig::resolve_inference_base_url`, which is exactly where
/// `resolve_aux_model_sampling_config`'s Tier-2 entry routes) instead of
/// being handed the primary unconditionally — the old "first-party by
/// construction" justification was false once `models_base_url` could point
/// anywhere.
///
/// M6: the platform and the base URL come from ONE
/// [`crate::agent::models::entry_for_slug`] lookup, so they can no longer
/// disagree (the aux path used to resolve the platform with `current_key`
/// and the credential with a separate `find_model_by_id`).
pub(crate) fn credential_for_slug(
&self,
models: &indexmap::IndexMap<String, ModelEntry>,
current_key: Option<&str>,
slug: &str,
) -> Option<SessionCredential> {
match crate::agent::models::entry_for_slug(models, current_key, slug) {
Some(entry) => self.credential_for_model(entry),
None => self.credential_for(None, &self.endpoints.resolve_inference_base_url()),
}
}
}
/// The registry platform a catalog entry belongs to (`info.id` is the managed
/// key `{platform}/{model}`). `None` for a bare / `[model.*]` entry.
pub(crate) fn entry_platform(entry: &ModelEntry) -> Option<kigi_models::PlatformId> {
entry
.info()
.id
.as_deref()
.and_then(kigi_models::parse_managed_model_key)
.map(|(platform, _)| platform)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::{AuthMode, KimiAuth, KimiCodeConfig};
/// A primary holding a fixed in-memory bearer. The `TempDir` is returned so
/// the caller keeps it alive; the token is read from memory, so on-disk
/// contents are irrelevant.
fn primary(key: &str) -> (tempfile::TempDir, Arc<AuthManager>) {
let dir = tempfile::tempdir().unwrap();
let manager = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
manager.hot_swap(KimiAuth {
key: key.to_string(),
auth_mode: AuthMode::OAuth,
..KimiAuth::test_default()
});
(dir, manager)
}
fn authority(endpoints: EndpointsConfig, primary: Arc<AuthManager>) -> CredentialAuthority {
CredentialAuthority::new(endpoints, Some(primary))
}
fn platform(id: &str) -> kigi_models::PlatformId {
kigi_models::PlatformId::parse(id).expect("known platform")
}
/// H3 (REGRESSION): the effective coding endpoint is
/// `EndpointsConfig::proxy_url()`, which prefers `[endpoints]
/// coding_api_base_url` from **config.toml** — the key the managed-config
/// sync writes. A predicate that knows only `KIGI_CODE_BASE_URL` classifies
/// such a deployment as third-party, withholds the api_key AND the
/// resolver, and 401s on every turn.
///
/// Revert-to-red: drop the `proxy_url()` arm from
/// `is_session_coding_endpoint` (leaving only
/// `is_effective_coding_endpoint_url`) and every assertion here fails —
/// with NO env var set anywhere in the test.
#[test]
fn config_toml_coding_endpoint_still_rides_the_session_bearer() {
let (_d, kimi) = primary("kimi-tok");
let managed = "https://proxy.acme.com/v1";
let auth = authority(
EndpointsConfig {
coding_api_base_url: Some(managed.to_string()),
..EndpointsConfig::default()
},
kimi.clone(),
);
assert_eq!(
auth.credential_class(None, managed),
CredentialClass::Primary,
"a [model.*] entry inheriting the managed coding endpoint takes the session bearer"
);
assert_eq!(
auth.credential_for(None, managed)
.map(|c| c.expose().to_owned()),
Some("kimi-tok".to_string()),
"the managed deployment must still receive the session bearer"
);
assert!(
auth.manager_for(None, managed).is_some(),
"and must keep a live manager, or it loses refresh and 401 recovery"
);
// kimi-code entries route to `proxy_url()` too (models_fetch's
// `platform_fetch_base`), so the platform arm must honour it as well.
assert_eq!(
auth.credential_for(Some(platform("kimi-code")), managed)
.map(|c| c.expose().to_owned()),
Some("kimi-tok".to_string()),
);
// A DIFFERENT authority (no managed key configured) must NOT trust it.
let default_auth = authority(EndpointsConfig::default(), kimi);
assert_eq!(
default_auth.credential_class(None, managed),
CredentialClass::None,
"the managed host is only trusted for the session that configured it"
);
}
/// The `models_base_url` custom-endpoint mode is equally invisible to the
/// env-var-only predicate.
#[test]
fn config_toml_models_base_url_still_rides_the_session_bearer() {
let (_d, kimi) = primary("kimi-tok");
let custom = "https://models.acme.internal/v1";
let auth = authority(
EndpointsConfig {
models_base_url: Some(custom.to_string()),
..EndpointsConfig::default()
},
kimi,
);
assert_eq!(
auth.credential_for(None, custom)
.map(|c| c.expose().to_owned()),
Some("kimi-tok".to_string()),
);
}
/// The compiled production endpoint and loopback proxies are unchanged.
#[test]
fn production_and_loopback_endpoints_are_unchanged() {
let (_d, kimi) = primary("kimi-tok");
let auth = authority(EndpointsConfig::default(), kimi);
for url in [
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
"http://127.0.0.1:8080/v1",
"http://localhost:3000/v1",
"http://[::1]:9000/v1",
] {
assert_eq!(
auth.credential_class(None, url),
CredentialClass::Primary,
"{url}: the session's own endpoint is byte-identical"
);
}
}
/// LEAK guard: every API-key registry platform, and any platform-less model
/// on a third-party host, gets NO session credential and NO manager.
#[test]
fn third_party_endpoints_never_receive_the_primary() {
let (_d, kimi) = primary("kimi-tok");
let auth = authority(EndpointsConfig::default(), kimi);
for id in [
"deepseek",
"openai",
"anthropic",
"moonshot-cn",
"moonshot-ai",
] {
let p = platform(id);
assert!(
auth.credential_for(Some(p), &p.base_url()).is_none(),
"LEAK: {id} is an API-key platform — no session bearer may ride there"
);
assert!(auth.manager_for(Some(p), &p.base_url()).is_none());
}
for url in [
"https://api.openai.com/v1",
"https://api.deepseek.com/v1",
"https://api.moonshot.cn/v1",
"",
] {
assert!(
auth.credential_for(None, url).is_none(),
"LEAK: {url} is a third-party host"
);
}
}
/// L10: an OAuth platform whose `[model.*]` override redirects `base_url`
/// to a third-party host keeps `info.id` — and must NOT ship that
/// platform's pooled OAuth bearer there.
#[tokio::test]
async fn oauth_platform_redirected_to_a_third_party_host_gets_nothing() {
let (_d, kimi) = primary("kimi-tok");
let auth = authority(EndpointsConfig::default(), kimi);
for id in [
"claude-pro-max",
"openai-codex",
"github-copilot",
"xai-grok",
] {
let p = platform(id);
assert!(
auth.manager_for(Some(p), &p.base_url()).is_some(),
"{id} keeps its own pooled manager on its own host"
);
assert!(
auth.manager_for(Some(p), "https://third.party/v1")
.is_none(),
"LEAK: {id} redirected to a third-party host must ship no bearer"
);
assert_eq!(
auth.credential_class(Some(p), "https://third.party/v1"),
CredentialClass::None,
"LEAK: {id} redirected to a third-party host takes no session credential"
);
}
}
/// C1 — a subscription platform's own host DOES take a session credential,
/// but it is that platform's POOLED token, never the primary / house key.
/// Two look-alike booleans used to encode this, and picking the wrong one is
/// the whole defect; one classifier makes the distinction impossible to
/// mis-read.
#[test]
fn a_subscription_host_classifies_pooled_never_primary() {
let (_d, kimi) = primary("kimi-tok");
let auth = authority(EndpointsConfig::default(), kimi);
for id in [
"claude-pro-max",
"openai-codex",
"github-copilot",
"xai-grok",
] {
let p = platform(id);
assert_eq!(
auth.credential_class(Some(p), &p.base_url()),
CredentialClass::Pooled,
"LEAK: {id}'s own host takes its POOLED token — never the primary / house key"
);
}
// Every API-key registry platform: nothing at all.
for id in ["deepseek", "openai", "anthropic", "moonshot-cn"] {
let p = platform(id);
assert_eq!(
auth.credential_class(Some(p), &p.base_url()),
CredentialClass::None
);
}
// The primary channel is unchanged: kimi-code and a platform-less model
// on the session's own endpoint, and nothing on a third-party host.
for url in [
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
"http://127.0.0.1:8080/v1",
] {
assert_eq!(auth.credential_class(None, url), CredentialClass::Primary);
assert_eq!(
auth.credential_class(Some(platform("kimi-code")), url),
CredentialClass::Primary
);
}
assert_eq!(
auth.credential_class(None, "https://api.openai.com/v1"),
CredentialClass::None
);
}
/// The four subscription-OAuth platforms draw from their OWN pooled
/// managers — never the primary Kimi one, even under a Kimi session.
#[tokio::test]
async fn oauth_platforms_never_resolve_the_primary() {
let (_d, kimi) = primary("kimi-tok");
let auth = authority(EndpointsConfig::default(), kimi.clone());
for id in [
"claude-pro-max",
"openai-codex",
"github-copilot",
"xai-grok",
] {
let p = platform(id);
let resolved = auth
.manager_for(Some(p), &p.base_url())
.expect("pooled manager");
assert!(
!Arc::ptr_eq(&resolved, &kimi),
"{id} must NOT resolve the primary Kimi manager"
);
assert_ne!(
auth.credential_for(Some(p), &p.base_url())
.map(|c| c.expose().to_owned()),
Some("kimi-tok".to_string()),
"{id} must never receive the primary Kimi bearer"
);
}
}
}
crates/codegen/kigi-shell/src/auth/mod.rs
+1
View File
@@ -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;
crates/codegen/kigi-shell/src/auth/oauth_registry.rs
+99 -405
View File
@@ -15,8 +15,13 @@
//! on-disk token stays fresh and a 401 recovers via the provider's own manager.
//! Managers are built ON DEMAND from the on-disk token ([`global_manager_for`]),
//! so a login landing AFTER a session spawned self-heals — no frozen per-session
//! snapshot. [`manager_for_model`] routes a managed catalog key to the pool
//! (oauth platform) or to the session's primary (everything else).
//! snapshot.
//!
//! ROUTING LIVES ELSEWHERE. This module is only the pool; the decision of which
//! credential governs a request belongs to the single chokepoint,
//! [`crate::auth::credential_authority::CredentialAuthority`]. Keeping the two
//! apart is deliberate: three rounds of leaks came from routing rules being
//! re-derived per call site.
//!
//! SECURITY: access/refresh tokens and resolved bearers are NEVER logged here.
@@ -38,26 +43,53 @@ fn oauth_manager_pool() -> &'static Mutex<HashMap<&'static str, Arc<AuthManager>
POOL.get_or_init(|| Mutex::new(HashMap::new()))
}
/// The kigi home every OAuth-pool call site resolves from. Single definition so
/// the pool, the aux/summary token routing and the session's inference manager
/// can never read different homes.
/// The kigi home EVERY OAuth-provider construction in this crate resolves from
/// — the pool, the catalog fetch's per-platform token resolution, the
/// aux/summary token routing and the session's inference manager — so they can
/// never read different homes.
///
/// Production: [`crate::util::kigi_home::kigi_home`]. LIB TESTS: a
/// process-lifetime `TempDir`, unconditionally — the pool is process-global and
/// every manager it builds starts a never-cancelled proactive-refresh loop, so
/// a unit test resolving the real `~/.kigi` would read the developer's stored
/// OAuth tokens and, 60 s later, fire REAL refresh requests against them.
/// Deliberately not a per-test opt-in that can be forgotten: `kigi_home()` is
/// itself a `OnceLock` an earlier test has usually already resolved to the real
/// home, so setting `KIGI_SHARE_DIR` in a test cannot pin it after the fact.
/// per-process path under the system temp dir that is deliberately **never
/// created**. The pool is process-global and every manager it builds starts a
/// never-cancelled proactive-refresh loop, so a unit test resolving the real
/// `~/.kigi` would read the developer's stored OAuth tokens and, 60 s later,
/// fire REAL refresh requests against them. Deliberately not a per-test opt-in
/// that can be forgotten: `kigi_home()` is itself a `OnceLock` an earlier test
/// has usually already resolved to the real home, so setting `KIGI_SHARE_DIR`
/// in a test cannot pin it after the fact.
///
/// M4 — THE LIMIT, STATED: `cfg(test)` is set only for THIS crate's `--lib`
/// tests. `crates/codegen/kigi-shell/tests/*.rs` link the library built WITHOUT
/// it, so for an integration test this resolves the real home unless that test
/// binary itself isolates one, which it must do through the two overrides the
/// auth stack already honours and BEFORE anything resolves `kigi_home()`:
/// `KIGI_SHARE_DIR` (read by `kigi_home()`, a `OnceLock`) or `KIGI_AUTH_PATH`
/// (read by [`AuthManager::new_oauth_provider`], which pins the token file
/// outright and so overrides this home entirely). 12 of the 28 integration
/// binaries under `crates/codegen/kigi-shell/tests/` set `KIGI_SHARE_DIR`; the
/// other 16 never reach an OAuth-platform inference path today, which is a
/// property of those tests, not a guarantee of this function. No
/// production-readable env override is added here on purpose: a knob that
/// redirects where OAuth tokens are read from is not worth a test convenience.
///
/// M8: this used to be a `static OnceLock<TempDir>`. Statics are never dropped,
/// so that leaked one temp directory per test binary — against the project's
/// "tests are TempDir self-cleaning" discipline. Nothing is created here
/// instead, and nothing in the lib-test suite creates it: a manager reads a
/// missing `auth.json` as "no session", and the only two paths that WRITE one
/// are a successful token refresh (which needs a stored refresh token that by
/// construction does not exist here) and a completed device login
/// ([`crate::agent::mvp_agent::MvpAgent::authenticate_oauth_platform`], which
/// M4 repointed at this same home). Both require the network, so no lib test
/// performs either. That is an observation about the suite, not an invariant of
/// this function — [`tests::test_pool_home_is_disposable_and_never_the_real_home`]
/// asserts the directory does not exist and is the tripwire if one ever does
/// (the path is per-PROCESS under the system temp dir, so the blast radius of a
/// future login-driving test is one disposable directory, never `~/.kigi`).
pub(crate) fn pool_home() -> std::path::PathBuf {
#[cfg(test)]
{
static TEST_HOME: OnceLock<tempfile::TempDir> = OnceLock::new();
TEST_HOME
.get_or_init(|| tempfile::tempdir().expect("tempdir for the test OAuth pool"))
.path()
.to_path_buf()
std::env::temp_dir().join(format!("kigi-oauth-pool-test-{}", std::process::id()))
}
#[cfg(not(test))]
crate::util::kigi_home::kigi_home()
@@ -89,411 +121,73 @@ pub(crate) fn global_manager_for(
manager
}
/// The `AuthManager` that governs INFERENCE auth for `managed_key`
/// (`{platform}/{model}`, e.g. `xai-grok/grok-4-latest`).
///
/// A generic device-code OAuth platform routes to ITS OWN scope-keyed manager
/// from the process-global pool ([`global_manager_for`], built on demand from
/// the on-disk token); every other key (Kimi, API-key platforms, `[model.*]`
/// entries, or an unprefixed bare id) routes to `primary`.
///
/// The pool is the single source of truth — there is no per-session snapshot to
/// freeze at spawn, so a grok login that happens AFTER a session spawned is
/// resolved correctly on the next grok turn. A grok key NEVER resolves to
/// `primary`: even before the user logs into grok the pooled manager simply
/// holds no token (its bearer / api_key is then `None`), so the Kimi
/// subscription bearer can never reach a third-party host — fail-fast, never a
/// silent fallback to the Kimi manager.
pub(crate) fn manager_for_model(
kigi_home: &Path,
managed_key: &str,
primary: Option<&Arc<AuthManager>>,
) -> Option<Arc<AuthManager>> {
if let Some((platform, _)) = kigi_models::parse_managed_model_key(managed_key)
&& let Some(oauth) = platform.oauth()
{
return Some(global_manager_for(kigi_home, oauth));
}
primary.cloned()
}
/// The SESSION token (the raw bearer/key string) that may ride an INFERENCE
/// request routed to `platform` at `base_url`. Used by the aux-model, summary
/// and subagent-override wire paths, where the result is stamped straight into
/// [`crate::agent::config::resolve_credentials`] as the request's `api_key`.
///
/// - a generic device-code OAuth platform (xai-grok, claude-pro-max,
/// github-copilot, openai-codex) draws from ITS OWN pooled manager; when that
/// provider has no stored session the result is `None` — never `primary`;
/// - `kimi-code`, and a platform-less model whose endpoint IS the session's own
/// coding endpoint (incl. a `KIGI_CODE_BASE_URL` deployment or a loopback
/// proxy), yield the primary's current-or-expired token — byte-identical to
/// reading it directly;
/// - every API-key registry platform, and every `[model.*]` block pointed at a
/// third-party host, yields `None`. Handing them `primary` put the user's
/// Kimi subscription bearer on `api.deepseek.com` / `api.moonshot.cn` / …
/// as the request's `api_key`.
///
/// SECURITY: the resolved token is never logged.
pub(crate) fn session_key_for_endpoint(
platform: Option<kigi_models::PlatformId>,
base_url: &str,
primary: Option<&Arc<AuthManager>>,
) -> Option<String> {
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
return global_manager_for(&pool_home(), oauth)
.current_or_expired()
.map(|a| a.key);
}
if !crate::agent::auth_method::platform_takes_session_credential(platform, base_url) {
return None;
}
primary
.and_then(|am| am.current_or_expired())
.map(|a| a.key)
}
/// [`session_key_for_endpoint`] for the catalog model whose routing slug (or
/// catalog key) is `slug`.
///
/// A slug absent from the catalog keeps the pre-registry behaviour: the aux
/// resolver's Tier-2 fallback builds its entry against
/// `EndpointsConfig::resolve_inference_base_url` (first-party), so the primary
/// still governs.
pub(crate) fn session_key_for_catalog_model(
models: &indexmap::IndexMap<String, crate::agent::config::ModelEntry>,
slug: &str,
primary: Option<&Arc<AuthManager>>,
) -> Option<String> {
let Some(entry) = crate::agent::config::find_model_by_id(models, slug) else {
return primary
.and_then(|am| am.current_or_expired())
.map(|a| a.key);
};
let info = entry.info();
session_key_for_endpoint(
info.id
.as_deref()
.and_then(kigi_models::parse_managed_model_key)
.map(|(platform, _)| platform),
&info.base_url,
primary,
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::KimiCodeConfig;
use crate::auth::{AuthMode, KimiAuth};
/// A Kimi manager holding a fixed in-memory bearer, standing in for a
/// session's primary. The `TempDir` is returned so the caller keeps it
/// alive; the token is read from memory (`current_or_expired`), so disk
/// contents are irrelevant to the assertion.
fn primary_with_token(key: &str) -> (tempfile::TempDir, Arc<AuthManager>) {
let dir = tempfile::tempdir().unwrap();
let manager = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
manager.hot_swap(KimiAuth {
key: key.to_string(),
auth_mode: AuthMode::OAuth,
..KimiAuth::test_default()
});
(dir, manager)
}
fn xai_oauth() -> &'static kigi_models::OAuthConfig {
kigi_models::PlatformId::XaiGrok
fn oauth_for(id: &str) -> &'static kigi_models::OAuthConfig {
kigi_models::PlatformId::parse(id)
.expect("known platform")
.oauth()
.expect("xai-grok carries an OAuthConfig")
.expect("subscription-OAuth platform carries an OAuthConfig")
}
fn claude_oauth() -> &'static kigi_models::OAuthConfig {
kigi_models::PlatformId::ClaudeProMax
.oauth()
.expect("claude-pro-max carries an OAuthConfig")
}
fn copilot_oauth() -> &'static kigi_models::OAuthConfig {
kigi_models::PlatformId::GithubCopilot
.oauth()
.expect("github-copilot carries an OAuthConfig")
}
fn codex_oauth() -> &'static kigi_models::OAuthConfig {
kigi_models::PlatformId::OpenaiCodex
.oauth()
.expect("openai-codex carries an OAuthConfig")
}
/// `session_key_for_endpoint` for a managed catalog key, resolving the
/// platform and its base URL from the registry exactly as the catalog entry
/// would.
fn session_key_for_key(
managed_key: &str,
primary: Option<&Arc<AuthManager>>,
) -> Option<String> {
let platform = kigi_models::parse_managed_model_key(managed_key).map(|(p, _)| p);
let base_url = platform
.map(kigi_models::PlatformId::base_url)
.unwrap_or_default();
session_key_for_endpoint(platform, &base_url, primary)
}
/// An `openai-codex/<model>` turn resolves to the process-global pooled
/// openai-codex manager (its OWN `oauth/openai-codex` scope), NEVER the
/// primary Kimi manager — the same leak-safe routing as the other OAuth
/// platforms, and a DISTINCT pool entry from each. Fail-fast: even with a
/// Kimi primary, a codex turn never yields the Kimi bearer.
/// Each subscription-OAuth platform gets its OWN process-global pooled
/// manager, and no two share one. (Which credential governs a REQUEST is
/// not decided here — see
/// [`crate::auth::credential_authority::CredentialAuthority`].)
#[tokio::test]
async fn openai_codex_model_resolves_to_its_own_manager_not_kimi() {
let (_kd, kimi) = primary_with_token("kimi-tok");
async fn every_oauth_scope_gets_its_own_pooled_manager() {
let home = tempfile::tempdir().unwrap();
let resolved = manager_for_model(home.path(), "openai-codex/gpt-5.5", Some(&kimi))
.expect("openai-codex model resolves to its pooled manager");
assert!(
!Arc::ptr_eq(&resolved, &kimi),
"openai-codex must NOT resolve to the Kimi manager"
);
assert!(
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), codex_oauth())),
"openai-codex must resolve to its OWN process-global pooled manager"
);
assert!(
!Arc::ptr_eq(&resolved, &global_manager_for(home.path(), copilot_oauth())),
"openai-codex and github-copilot must not share a pooled manager"
);
assert!(
!Arc::ptr_eq(&resolved, &global_manager_for(home.path(), claude_oauth())),
"openai-codex and claude-pro-max must not share a pooled manager"
);
assert_ne!(
session_key_for_key("openai-codex/gpt-5.5", Some(&kimi)),
Some("kimi-tok".to_string()),
"an openai-codex model must never receive the primary Kimi token"
);
}
/// A `github-copilot/<model>` turn resolves to the process-global pooled
/// github-copilot manager (its OWN `oauth/github-copilot` scope), NEVER the
/// primary Kimi manager — the same leak-safe routing as xai-grok /
/// claude-pro-max, and a DISTINCT pool entry from either.
#[tokio::test]
async fn github_copilot_model_resolves_to_its_own_manager_not_kimi() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
let resolved = manager_for_model(home.path(), "github-copilot/gpt-4.1", Some(&kimi))
.expect("github-copilot model resolves to its pooled manager");
assert!(
!Arc::ptr_eq(&resolved, &kimi),
"github-copilot must NOT resolve to the Kimi manager"
);
assert!(
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), copilot_oauth())),
"github-copilot must resolve to its OWN process-global pooled manager"
);
assert!(
!Arc::ptr_eq(&resolved, &global_manager_for(home.path(), claude_oauth())),
"github-copilot and claude-pro-max must not share a pooled manager"
);
// Fail-fast: even with a Kimi primary, a copilot turn never yields the
// Kimi bearer — it draws from the copilot pool (its own token, or None).
assert_ne!(
session_key_for_key("github-copilot/gpt-4.1", Some(&kimi)),
Some("kimi-tok".to_string()),
"a github-copilot model must never receive the primary Kimi token"
);
}
/// A `claude-pro-max/<model>` turn resolves to the process-global pooled
/// claude-pro-max manager (its OWN `oauth/claude-pro-max` scope), NEVER the
/// primary Kimi manager — the same leak-safe routing as xai-grok, and a
/// DISTINCT pool entry from the xai manager.
#[tokio::test]
async fn claude_pro_max_model_resolves_to_its_own_manager_not_kimi() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
let resolved =
manager_for_model(home.path(), "claude-pro-max/claude-opus-4-8", Some(&kimi))
.expect("claude-pro-max model resolves to its pooled manager");
assert!(
!Arc::ptr_eq(&resolved, &kimi),
"claude-pro-max must NOT resolve to the Kimi manager"
);
assert!(
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), claude_oauth())),
"claude-pro-max must resolve to its OWN process-global pooled manager"
);
// And it is a DIFFERENT manager than xai-grok's pooled one.
assert!(
!Arc::ptr_eq(&resolved, &global_manager_for(home.path(), xai_oauth())),
"claude-pro-max and xai-grok must not share a pooled manager"
);
}
/// Fail-fast (no Kimi fallback): a claude-pro-max key with a Kimi primary
/// never yields the Kimi session token — it draws from the claude pool (its
/// own token, or `None`), so the Kimi bearer can never reach api.anthropic.
#[tokio::test]
async fn session_key_for_claude_pro_max_is_never_the_kimi_primary() {
let (_kd, kimi) = primary_with_token("kimi-tok");
assert_ne!(
session_key_for_key("claude-pro-max/claude-opus-4-8", Some(&kimi)),
Some("kimi-tok".to_string()),
"a claude-pro-max model must never receive the primary Kimi session token"
);
}
/// A non-OAuth managed key (moonshot-cn/…) and an unprefixed bare id both
/// route to the primary Kimi manager — the Kimi / first-party path is
/// untouched and never consults the pool (no runtime needed).
#[test]
fn non_oauth_and_bare_models_route_to_primary() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
for key in ["moonshot-cn/kimi-k2", "kimi-k2-0905-preview"] {
let resolved = manager_for_model(home.path(), key, Some(&kimi))
.expect("non-oauth key routes to the primary");
let ids = [
"xai-grok",
"claude-pro-max",
"github-copilot",
"openai-codex",
];
let managers: Vec<_> = ids
.iter()
.map(|id| global_manager_for(home.path(), oauth_for(id)))
.collect();
for (i, a) in managers.iter().enumerate() {
assert!(
Arc::ptr_eq(&resolved, &kimi),
"{key} must resolve to the primary manager"
Arc::ptr_eq(a, &global_manager_for(home.path(), oauth_for(ids[i]))),
"{}: the pool must return the SAME manager for a scope",
ids[i]
);
assert_eq!(resolved.current_or_expired().unwrap().key, "kimi-tok");
for (j, b) in managers.iter().enumerate() {
if i != j {
assert!(
!Arc::ptr_eq(a, b),
"{} and {} must not share a pooled manager",
ids[i],
ids[j]
);
}
}
}
}
/// The primary being `None` (test / BYOK sessions) still yields `None` for a
/// non-oauth key, never a panic — and without touching the pool.
#[test]
fn none_primary_is_passed_through_for_non_oauth() {
let home = tempfile::tempdir().unwrap();
assert!(manager_for_model(home.path(), "kimi-k2", None).is_none());
}
/// An `xai-grok/<model>` turn resolves to the process-global pooled xai
/// manager, NEVER the primary Kimi manager — the pool is the single source.
#[tokio::test]
async fn grok_model_resolves_to_pooled_xai_manager_not_kimi() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi))
.expect("grok model resolves to the pooled xai manager");
assert!(
!Arc::ptr_eq(&resolved, &kimi),
"grok model must NOT resolve to the Kimi manager"
);
assert!(
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), xai_oauth())),
"grok model must resolve to the process-global pooled xai manager"
);
}
/// Facet B guard: the resolver routes purely by the model's platform, with
/// no auth-method input — so even when the session's primary is a Kimi
/// (session) manager holding "kimi-tok", a grok model never resolves that
/// Kimi token.
#[tokio::test]
async fn grok_model_under_kimi_primary_never_yields_kimi_token() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-fast", Some(&kimi))
.expect("grok model resolves to its own pooled manager regardless of primary");
assert!(!Arc::ptr_eq(&resolved, &kimi));
assert_ne!(
resolved.current_or_expired().map(|a| a.key),
Some("kimi-tok".to_string()),
"the Kimi bearer must never be what a grok turn resolves"
);
}
/// Fail-fast: a grok key resolves to the pooled xai manager (never the Kimi
/// primary) even with no stored grok session in the pool — the pooled
/// manager then simply holds no token, so nothing (least of all the Kimi
/// bearer) is sent to api.x.ai.
#[tokio::test]
async fn grok_never_falls_back_to_kimi_primary() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi))
.expect("grok routes to the pooled xai manager, not None");
assert!(
!Arc::ptr_eq(&resolved, &kimi),
"an OAuth platform must never fall back to the primary Kimi manager"
);
}
/// `session_key_for_endpoint`: the endpoints that genuinely ride the
/// PRIMARY session — `kimi-code` (the subscription channel) and a
/// platform-less model routed at the session's own coding endpoint (a
/// `KIGI_CODE_BASE_URL` deployment or a loopback dev proxy) — yield the
/// primary token exactly as reading it directly would. No runtime / pool
/// touched.
#[test]
fn session_key_for_the_sessions_own_endpoint_is_the_primary_token() {
let (_kd, kimi) = primary_with_token("kimi-tok");
assert_eq!(
session_key_for_key("kimi-code/kimi-for-coding", Some(&kimi)),
Some("kimi-tok".to_string()),
"kimi-code rides the primary session, unchanged"
);
for url in [
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
"http://127.0.0.1:4000/v1",
] {
assert_eq!(
session_key_for_endpoint(None, url, Some(&kimi)),
Some("kimi-tok".to_string()),
"{url}: a platform-less model on the session's own endpoint is unchanged"
);
}
}
/// LEAK guard (aux / summary / subagent-override `api_key` channel): an
/// API-key registry platform, and a `[model.*]` block pointed at a
/// third-party host, must yield NO session token. Handing them the primary
/// stamped the user's Kimi subscription bearer onto `api.moonshot.cn` /
/// `api.deepseek.com` as the request's `api_key` — the channel the
/// `bearer_resolver` guard alone does not close.
/// M8: the test pool home is a per-process path that is never created, so a
/// test binary leaves nothing behind (and never resolves the developer's
/// real `~/.kigi`, whose stored OAuth tokens the pool would otherwise read
/// and proactively refresh over the network).
///
/// Revert-to-red: dropping the `platform_takes_session_credential` term
/// from `session_key_for_endpoint` returns `Some("kimi-tok")` here.
/// This is also the tripwire for the cleanup claim in [`pool_home`]: a
/// completed device login through `authenticate_oauth_platform` WOULD create
/// this directory, so if a lib test ever drives one, this assertion fires
/// and the cleanup has to be added rather than silently regressing the
/// "tests are TempDir self-cleaning" discipline.
#[test]
fn session_key_for_a_third_party_endpoint_is_never_the_primary_token() {
let (_kd, kimi) = primary_with_token("kimi-tok");
for key in [
"moonshot-cn/kimi-k2",
"deepseek/deepseek-chat",
"openai/gpt-5",
] {
assert_eq!(
session_key_for_key(key, Some(&kimi)),
None,
"LEAK: {key} is an API-key platform — no session token may ride there"
);
}
assert_eq!(
session_key_for_endpoint(None, "https://api.openai.com/v1", Some(&kimi)),
None,
"LEAK: a [model.*] block on a third-party host gets no session token"
fn test_pool_home_is_disposable_and_never_the_real_home() {
let home = pool_home();
assert!(
home.starts_with(std::env::temp_dir()),
"the test pool home must live under the system temp dir, got {home:?}"
);
}
/// LEAK guard (aux-model + subagent-override token routing): a grok key with
/// a Kimi primary NEVER yields the primary Kimi token — it draws from the
/// pooled xai manager (its own token, or `None`). This is the exact source
/// the aux `session_key` and the override `session_key` now use.
#[tokio::test]
async fn session_key_for_grok_is_never_the_kimi_primary() {
let (_kd, kimi) = primary_with_token("kimi-tok");
assert_ne!(
session_key_for_key("xai-grok/grok-4-latest", Some(&kimi)),
Some("kimi-tok".to_string()),
"a grok aux/override model must never receive the primary Kimi session token"
);
// Even with `None` primary the routing is unchanged: grok → pool, never a panic.
assert_ne!(
session_key_for_key("xai-grok/grok-4-fast", None),
Some("kimi-tok".to_string()),
assert!(
!home.exists(),
"the test pool home must not be created — nothing to clean up"
);
}
}
crates/codegen/kigi-shell/src/session/acp_session.rs
+15
View File
@@ -417,6 +417,21 @@ pub(crate) struct SessionActor {
/// [`SessionActor::model_auth_facts`].
pub(crate) model_auth_facts:
std::cell::RefCell<Option<(String, crate::agent::config::ModelAuthFacts)>>,
/// The catalog KEY this session's model was selected by (`{platform}/{model}`
/// for a registry model), owned PER SESSION.
///
/// H4: `SamplingConfig::model` is the bare routing slug, and duplicate slugs
/// across an API-key platform and its subscription-OAuth twin
/// (`xai`/`xai-grok`, `anthropic`/`claude-pro-max`, `openai`/`openai-codex`)
/// are BY DESIGN, so the slug alone cannot name the platform. This used to
/// be read from `ModelsManager::current_model_id()` — a single
/// PROCESS-GLOBAL cell that Leader mode never writes
/// (`agent/handlers/model_switch.rs`) and that is last-writer-wins across
/// concurrent sessions, so both collision directions resolved the wrong
/// platform: the subscription session lost its resolver (unrecoverable 401
/// ~1h in) and the API-key session got the pooled OAuth bearer stamped over
/// its own `sk-…` key. Written at spawn and on every `SetSessionModel`.
pub(crate) selected_catalog_key: std::cell::RefCell<Option<String>>,
/// 401-attribution callback. Joined with the bearer the
/// sampler sends on the wire to emit an `auth 401 attribution`
/// event at each of the six `OaiCompatClient` 401 arms in
crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs
+8 -4
View File
@@ -5,12 +5,18 @@ impl SessionActor {
pub(super) async fn handle_set_session_model(
&self,
sampling_config: kigi_sampler::SamplerConfig,
catalog_key: Option<String>,
use_concise: bool,
apply_prompt_override: bool,
skip_prompt_rewrite: bool,
auto_compact_threshold_percent: u8,
) -> Result<acp::ModelId, acp::Error> {
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,
crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs
+11 -12
View File
@@ -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}"
crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs
+21 -8
View File
@@ -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
crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs
+275 -148
View File
@@ -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<kigi_models::PlatformId>,
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<kigi_models::PlatformId>,
model_byok: crate::agent::auth_method::ModelByok,
base_url: &str,
) -> Option<kigi_sampler::SharedBearerResolver> {
let is_primary_channel = platform.is_none_or(|p| p.oauth().is_none());
if is_primary_channel
&& !SessionTokenAuthGate::new(auth_method_id, model_byok, base_url, platform, authority)
.active()
{
return None;
}
authority.bearer_resolver_for(platform, base_url)
}
/// Run a tool call; on an auth-shaped failure, attempt recovery via
/// `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<crate::auth::AuthManager>);
@@ -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<crate::auth::AuthManager>,
) -> kigi_sampler::SharedBearerResolver {
std::sync::Arc::new(AuthManagerBearerResolver(am))
}
/// The `bearer_resolver` an AUX / summary `SamplerConfig` may carry, given the
/// SESSION model's resolver (`stamped`) and the AUX model's own platform +
/// endpoint. ONE decision shared by image-describe, the auto-mode classifier
/// (via [`SessionActor::repoint_aux_bearer_resolver`]) and
/// `MvpAgent::build_summary_client`.
///
/// [`crate::agent::config::stamp_session_local_sampler_fields`] copies the
/// session resolver onto every aux config, and `SamplingClient::post` REPLACES
/// the request's auth header from it — so an aux model on a DIFFERENT provider
/// would have its own correctly-resolved key overwritten by the session bearer
/// on the AUX host (H3/H4). Resolve by the aux model instead:
/// - an OAuth platform → a live resolver over ITS OWN pooled manager (so a grok
/// / claude-pro-max / copilot / codex aux model keeps mid-session refresh);
/// - `kimi-code`, or a platform-less model on the session's own coding endpoint
/// → the stamped session resolver, byte-identical;
/// - every API-key registry platform, and any `[model.*]` block pointed at a
/// third-party host → `None`, so the aux model's own key survives to the wire.
///
/// SECURITY: no token is logged.
pub(crate) fn aux_bearer_resolver(
stamped: Option<kigi_sampler::SharedBearerResolver>,
platform: Option<kigi_models::PlatformId>,
base_url: &str,
) -> Option<kigi_sampler::SharedBearerResolver> {
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
return Some(auth_manager_bearer_resolver(
crate::auth::oauth_registry::global_manager_for(
&crate::auth::oauth_registry::pool_home(),
oauth,
),
));
}
if crate::agent::auth_method::platform_takes_session_credential(platform, base_url) {
return stamped;
}
None
}
impl SessionActor {
pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
let mcp_wait_start = std::time::Instant::now();
@@ -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<std::sync::Arc<crate::auth::AuthManager>> {
// `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<String> {
/// 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<crate::auth::credential_authority::SessionCredential> {
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<String> {
self.selected_catalog_key.borrow().clone()
}
/// Keep the session's own selected catalog key consistent with an
/// `OverrideModelName` rename: KEEP it when it still names `model_name`
/// (same entry, new routing name), otherwise CLEAR it.
///
/// H-c: `OverrideModelName` is the one command that rewrites
/// `SamplingConfig::model` without going through `SetSessionModel`, so it
/// used to leave the field naming a model the session is no longer on.
/// Clearing rather than re-resolving is deliberate: re-resolving would put
/// `resolve_catalog_key`'s `.rev()` guess INTO the field the whole rule
/// treats as the session's deliberate selection, and a cleared field
/// refuses a collided slug instead of guessing its OAuth twin (H-b).
pub(super) fn retain_selected_catalog_key_for(&self, model_name: &str) {
let models = self.models_manager.models();
let still_names_it = self.selected_catalog_key().is_some_and(|key| {
key == model_name
|| models
.get(key.as_str())
.is_some_and(|entry| entry.info.model == model_name)
});
if !still_names_it {
*self.selected_catalog_key.borrow_mut() = None;
}
}
/// The registry platform the routing slug `model` belongs to, from the SAME
/// 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<kigi_models::PlatformId> {
pub(super) fn model_platform(&self, model: &str) -> Option<kigi_models::PlatformId> {
let models = self.models_manager.models();
let current = self.models_manager.current_model_id();
crate::agent::models::platform_for_slug(&models, Some(current.0.as_ref()), model)
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<kigi_sampler::SharedBearerResolver> {
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(&current_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(&current_model_id).is_some()
&& !self.model_has_own_credential(&current_model_id)
{
return;
}
let Some(new_key) = self.reload_api_key_from_config(&current_model_id) else {
@@ -1125,6 +1243,15 @@ impl SessionActor {
creds.api_key = Some(new_key);
self.chat_state_handle.update_credentials(creds);
}
/// Whether the live catalog entry for `slug` carries its own credential —
/// an `api_key`/`env_key` from a `[model.*]` block, which a config edit can
/// rotate mid-session. A platform entry whose key came from the platform
/// credential has none.
fn model_has_own_credential(&self, slug: &str) -> bool {
let models = self.models_manager.models();
crate::agent::config::find_model_by_id(&models, slug)
.is_some_and(crate::agent::config::ModelEntry::has_own_credentials)
}
fn reload_api_key_from_config(&self, current_model_id: &str) -> Option<String> {
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();
crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs
+11
View File
@@ -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>| 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,
crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs
+32 -17
View File
@@ -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!(
crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs
+4
View File
@@ -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,
crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs
+1
View File
@@ -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(
crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs
+3
View File
@@ -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(
crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs
+1
View File
@@ -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,
crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs
+1
View File
@@ -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,
crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_platform_tests.rs
+770 -69
View File
@@ -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<String> {
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<acp::ModelId>` written by whichever session switched last. With one
/// session on `xai-grok/grok-4.5` and another on `xai/grok-4.5` — same routing
/// slug, by design — the loser resolved the OTHER session's platform:
/// the subscription session lost its live resolver (unrecoverable 401 ~1h in)
/// and the API-key session got the pooled OAuth bearer stamped over its own
/// `sk-…` key, which the provider rejects.
///
/// Here the global cell is deliberately set to the API-key twin for BOTH
/// sessions (last writer wins, and it was the API-key one). Each session must
/// still resolve ITS OWN selection.
///
/// Revert-to-red: replace `self.selected_catalog_key()` in
/// `SessionActor::model_platform` with
/// `Some(self.models_manager.current_model_id().0.as_ref())`. Under THIS
/// fixture both of the subscription session's assertions fail. (L: the fixture
/// is what makes that true — `managed_entry` carries `api_key: None` and
/// `actor_with_catalog` pins `NotByok`, which is exactly what a FETCHED registry
/// entry resolves to. A user who additionally sets `ANTHROPIC_API_KEY` /
/// `[model.*] env_key` classifies `Byok`, the gate is inactive for that reason
/// alone, and only the `anthropic_oauth` assertion would still catch the
/// mis-resolution — hence the env guard below.)
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn concurrent_sessions_on_a_colliding_slug_each_resolve_their_own_platform() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let api_key_twin = "anthropic/claude-opus-4-8";
let oauth_twin = "claude-pro-max/claude-opus-4-8";
let slug = "claude-opus-4-8";
let host = &anthropic_collision_host();
let catalog = || {
vec![
// API-key platform FIRST, exactly as `PlatformId::ALL` orders them.
managed_entry(api_key_twin, slug, host),
managed_entry(oauth_twin, slug, host),
]
};
let (_d1, subscription, _r1) =
actor_with_catalog(catalog(), oauth_twin, "unused").await;
let (_d2, api_key, _r2) =
actor_with_catalog(catalog(), api_key_twin, "sk-ant-user").await;
// The other session switched last: the process-global cell now names
// the API-key twin for BOTH.
for actor in [&subscription, &api_key] {
actor
.models_manager
.set_current_model_id(acp::ModelId::new(api_key_twin.to_string()));
}
let sub_cfg = subscription.reconstruct_full_config().await;
let resolver = sub_cfg.bearer_resolver.as_ref().expect(
"the subscription session must keep its own live bearer_resolver even when \
another session switched the process-global model last",
);
assert_ne!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"it must read the claude-pro-max pool, never the Kimi primary"
);
assert!(
sub_cfg.anthropic_oauth,
"the Claude OAuth Messages adaptation must follow the SUBSCRIPTION session"
);
let api_cfg = api_key.reconstruct_full_config().await;
assert!(
api_cfg.bearer_resolver.is_none(),
"LEAK: the API-key session must get no session bearer_resolver"
);
assert!(
!api_cfg.anthropic_oauth,
"the API-key session must not get the OAuth Messages adaptation"
);
assert_eq!(
api_cfg.api_key.as_deref(),
Some("sk-ant-user"),
"the API-key session keeps its own provider key"
);
})
.await;
}
/// H4 in LEADER mode, where `agent/handlers/model_switch.rs` skips
/// `set_current_model_id` ENTIRELY, so the process-global cell is frozen at the
/// startup default for the whole process lifetime. `platform_for_slug` then
/// fell through to the `.rev()` scan, which returns the LAST match — the OAuth
/// twin — so a Leader-mode session on the API-KEY twin was handed the pooled
/// OAuth bearer plus the Messages adaptation, and Anthropic rejects both.
///
/// Revert-to-red: same edit as above; with the global cell naming the startup
/// default (not in this catalog) the `.rev()` fallback resolves
/// `claude-pro-max/*` and both assertions fail.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn leader_mode_session_resolves_its_own_platform_without_the_global_cell() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let slug = "claude-opus-4-8";
let host = &anthropic_collision_host();
let (_dir, actor, _rx) = actor_with_catalog(
vec![
managed_entry("anthropic/claude-opus-4-8", slug, host),
managed_entry("claude-pro-max/claude-opus-4-8", slug, host),
],
"anthropic/claude-opus-4-8",
"sk-ant-user",
)
.await;
// Leader mode never writes the global cell: it still names the
// startup default, which is not in this catalog at all.
assert!(
!actor
.models_manager
.models()
.contains_key(actor.models_manager.current_model_id().0.as_ref()),
"precondition: the process-global model id is stale (Leader mode)"
);
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_none(),
"LEAK: a Leader-mode API-key session must get no session bearer_resolver"
);
assert!(
!cfg.anthropic_oauth,
"a Leader-mode API-key session must not get the OAuth Messages adaptation"
);
})
.await;
}
/// L13 — a user whose ACP auth method is an API-KEY registry platform (e.g.
/// `deepseek`) can still SELECT a subscription-OAuth model, and the chokepoint
/// hands it that platform's pooled bearer as the request's `api_key`. The gate
/// keys on the primary method, which is not session-based, so the config used
/// to carry NO `bearer_resolver`: the pooled token froze at selection time and
/// the session died with an unrecoverable 401 once it expired (~1h).
///
/// The model's own credential now makes the gate session-based, confined to
/// that platform's own host by the gate's `credential_class` conjunct.
#[tokio::test(flavor = "current_thread")]
async fn oauth_model_under_an_api_key_auth_method_keeps_its_pooled_resolver() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let base_url = kigi_models::PlatformId::ClaudeProMax.base_url();
let (_dir, actor, _rx) = actor_with_catalog(
vec![managed_entry(
"claude-pro-max/claude-opus-4-8",
"claude-opus-4-8",
&base_url,
)],
"claude-pro-max/claude-opus-4-8",
"unused",
)
.await;
// The PRIMARY ACP method is an API-key registry platform login.
actor
.auth_method_id
.store(Some(std::sync::Arc::new(acp::AuthMethodId::new(
"deepseek",
))));
let cfg = actor.reconstruct_full_config().await;
let resolver = cfg.bearer_resolver.as_ref().expect(
"a subscription-OAuth model keeps a live resolver whatever the primary \
ACP auth method is, or it cannot refresh mid-session",
);
assert_ne!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"and it reads the claude-pro-max pool, never the primary"
);
// An API-key-platform model under the same method stays resolver-free.
let (_d2, deepseek, _r2) = actor_with_catalog(
vec![managed_entry(
"deepseek/deepseek-chat",
"deepseek-chat",
"https://api.deepseek.com/v1",
)],
"deepseek/deepseek-chat",
"sk-deepseek",
)
.await;
deepseek
.auth_method_id
.store(Some(std::sync::Arc::new(acp::AuthMethodId::new(
"deepseek",
))));
assert!(
deepseek
.reconstruct_full_config()
.await
.bearer_resolver
.is_none(),
"LEAK: an API-key-platform model must never get a session resolver"
);
})
.await;
}
/// H-b — a `None` or STALE per-session catalog key must REFUSE, not degrade to
/// the subscription-OAuth twin.
///
/// `model_platform` falls through to `resolve_catalog_key`'s `.rev()` scan when
/// the session's own key does not name the slug, and that scan returns the LAST
/// match — the OAuth twin, because `PlatformId::ALL` orders every API-key
/// platform first. Combined with the L13 disjunct (a model whose own credential
/// is a pooled OAuth session is session-based BY ITSELF), an API-KEY session on
/// `anthropic/claude-opus-4-8` with no per-session key got
/// `is_session_based = true`, `credential_class = Pooled` (same
/// host) and, at `NotByok`, an ACTIVE gate — so `manager_for` handed it the
/// Claude POOLED manager, whose `bearer_resolver` REPLACES the user's own
/// `sk-ant-…` on the wire, plus the OAuth Messages adaptation. Anthropic rejects
/// both. This is exactly what H4 prevents, reached through the `None` path.
///
/// A key that is absent or names a different model is not evidence for either
/// twin: resolve to NO platform, which the chokepoint then decides purely by the
/// ENDPOINT (the OAuth host is not this session's coding endpoint ⇒ nothing
/// rides).
///
/// Revert-to-red (production, compiles): delete the
/// `platform.oauth().is_some() && !disambiguated && slug_collides_across_platforms(..)`
/// refusal from `crate::agent::models::platform_for_slug` and every
/// `bearer_resolver` / `anthropic_oauth` assertion below fails.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn a_missing_or_stale_session_key_refuses_instead_of_guessing_the_oauth_twin() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let api_key_twin = "anthropic/claude-opus-4-8";
let slug = "claude-opus-4-8";
let host = anthropic_collision_host();
let catalog = || {
vec![
// API-key platform FIRST, exactly as `PlatformId::ALL` orders them.
managed_entry(api_key_twin, slug, &host),
managed_entry("claude-pro-max/claude-opus-4-8", slug, &host),
]
};
for (case, stale_key) in [
// No key at all: a session spawned on a model that left the
// catalog, or one an older build never seeded.
("absent", None),
// Stale: an `OverrideModelName` rename, or a key naming a model
// this session is no longer on.
("stale", Some("claude-pro-max/some-other-model".to_string())),
] {
let (_dir, actor, _rx) =
actor_with_catalog(catalog(), api_key_twin, "sk-ant-user").await;
*actor.selected_catalog_key.borrow_mut() = stale_key;
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_none(),
"{case}: LEAK — an unresolvable selection must get NO session bearer \
resolver; the pooled OAuth bearer would REPLACE the user's own key"
);
assert!(
!cfg.anthropic_oauth,
"{case}: nor the Claude OAuth Messages adaptation"
);
assert_eq!(
cfg.api_key.as_deref(),
Some("sk-ant-user"),
"{case}: the user's own provider key must survive untouched"
);
assert!(
actor
.credential_authority()
.manager_for(
crate::agent::models::platform_for_slug(
&actor.models_manager.models(),
actor.selected_catalog_key().as_deref(),
slug,
),
&host,
)
.is_none(),
"{case}: and no manager either — refuse, never guess"
);
}
// …while a session that DID select the OAuth twin still gets its
// pooled resolver: the refusal is about the guess, not the platform.
let (_dir, selected, _rx) =
actor_with_catalog(catalog(), "claude-pro-max/claude-opus-4-8", "unused").await;
let cfg = selected.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_some() && cfg.anthropic_oauth,
"a DELIBERATE subscription selection keeps its pooled resolver and \
adaptation (this is what makes the refusals above meaningful)"
);
})
.await;
}
/// H-c — coverage for the FIRST of the two production writers of
/// `selected_catalog_key`: the spawn seed
/// (`crate::agent::models::selected_catalog_key_for_spawn`, called from
/// `spawn.rs`). Every other test in this module sets the field by hand, so a
/// wrong seed was silent.
///
/// Both spawn shapes are covered: a FRESH session, spawned on the catalog key
/// the picker resolved, and a RESUME/LOAD, which spawns with the RAW persisted
/// `summary.current_model_id` — a BARE routing slug after any `SetSessionModel`,
/// since `handle_set_session_model` persists `sampling_config.model`. This seed
/// is where that slug becomes a key. The assertion is end-to-end: the seeded key
/// is fed to the very function the auth layer keys on.
#[test]
fn spawn_seeds_the_session_key_the_auth_layer_keys_on() {
let slug = "claude-opus-4-8";
let host = anthropic_collision_host();
let api_key_twin = "anthropic/claude-opus-4-8";
let oauth_twin = "claude-pro-max/claude-opus-4-8";
let models: indexmap::IndexMap<String, crate::agent::config::ModelEntry> = [
managed_entry(api_key_twin, slug, &host),
managed_entry(oauth_twin, slug, &host),
]
.into_iter()
.collect();
// FRESH: spawned on the catalog key. Idempotent, and it disambiguates.
for selected in [api_key_twin, oauth_twin] {
let seeded = crate::agent::models::selected_catalog_key_for_spawn(
&models,
&acp::ModelId::new(selected.to_string()),
);
assert_eq!(
seeded.as_deref(),
Some(selected),
"a fresh session must record the catalog key it was spawned with"
);
assert_eq!(
crate::agent::models::platform_for_slug(&models, seeded.as_deref(), slug),
kigi_models::parse_managed_model_key(selected).map(|(p, _)| p),
"…and that key must resolve THIS session's own platform for the bare slug"
);
}
// RESUME/LOAD: `acp_agent::load_session` spawns with the RAW persisted id,
// which after any model switch is the bare routing slug. THIS seed is what
// turns it into a key — the picker's `.rev()` answer, which is the resume
// default for a collided slug.
let resumed = crate::agent::models::selected_catalog_key_for_spawn(
&models,
&acp::ModelId::new(slug.to_string()),
);
assert_eq!(
resumed.as_deref(),
Some(oauth_twin),
"a bare persisted slug resolves through the picker's own lookup"
);
// A model that is no longer in the catalog seeds NOTHING, which (H-b) then
// refuses rather than guessing a twin.
assert_eq!(
crate::agent::models::selected_catalog_key_for_spawn(
&models,
&acp::ModelId::new("gone/model".to_string()),
),
None,
"a model that left the catalog must not seed a key"
);
}
/// H-c — coverage for the SECOND production writer: `SetSessionModel`
/// (`handle_set_session_model`), the picker's own path. The existing test
/// through this handler passes `None`, so a handler that dropped the key on the
/// floor stayed green.
///
/// End-to-end: after the switch the session's per-turn config must carry the
/// SELECTED twin's pooled resolver and adaptation, even though the bare slug in
/// the config is ambiguous.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn set_session_model_records_the_key_the_next_turn_resolves_on() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let slug = "claude-opus-4-8";
let host = anthropic_collision_host();
let api_key_twin = "anthropic/claude-opus-4-8";
let oauth_twin = "claude-pro-max/claude-opus-4-8";
let (_dir, actor, _rx) = actor_with_catalog(
vec![
managed_entry(api_key_twin, slug, &host),
managed_entry(oauth_twin, slug, &host),
],
api_key_twin,
"sk-ant-user",
)
.await;
// Switch to the SUBSCRIPTION twin, exactly as
// `agent/handlers/model_switch.rs` does: the ambiguous slug in the
// sampler config plus the catalog KEY the picker resolved.
let models = actor.models_manager.models();
let entry = models.get(oauth_twin).expect("catalog entry");
let sampler = crate::agent::config::sampling_config_for_model(
entry,
crate::agent::config::resolve_credentials(entry, None),
None,
);
actor
.handle_set_session_model(
sampler,
Some(oauth_twin.to_string()),
false,
false,
true,
85,
)
.await
.expect("model switch");
assert_eq!(
actor.selected_catalog_key().as_deref(),
Some(oauth_twin),
"SetSessionModel must record the picker's catalog key"
);
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_some(),
"the switched-to subscription model must keep a live pooled resolver"
);
assert!(
cfg.anthropic_oauth,
"…and the Claude OAuth Messages adaptation"
);
// And back: switching to the API-key twin must UNDO both.
let entry = models.get(api_key_twin).expect("catalog entry");
let sampler = crate::agent::config::sampling_config_for_model(
entry,
crate::agent::config::resolve_credentials(entry, None),
None,
);
actor
.handle_set_session_model(
sampler,
Some(api_key_twin.to_string()),
false,
false,
true,
85,
)
.await
.expect("model switch");
let cfg = actor.reconstruct_full_config().await;
assert!(
cfg.bearer_resolver.is_none() && !cfg.anthropic_oauth,
"LEAK: switching back to the API-key twin must drop the pooled resolver \
and the OAuth adaptation"
);
})
.await;
}
/// H-c — `OverrideModelName` is the one command that rewrites
/// `SamplingConfig::model` WITHOUT going through `SetSessionModel`, so it used
/// to leave `selected_catalog_key` naming a model the session is no longer on.
/// It must keep the field consistent: KEEP when the key still names the new
/// routing name, CLEAR otherwise — never re-resolve, which would put the
/// `.rev()` guess into the field the rule treats as a deliberate selection.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn override_model_name_keeps_the_session_key_consistent() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let slug = "claude-opus-4-8";
let host = anthropic_collision_host();
let oauth_twin = "claude-pro-max/claude-opus-4-8";
let (_dir, actor, _rx) = actor_with_catalog(
vec![
managed_entry("anthropic/claude-opus-4-8", slug, &host),
managed_entry(oauth_twin, slug, &host),
],
oauth_twin,
"unused",
)
.await;
// A rename to the SAME model's routing slug (or to its catalog key)
// keeps the selection.
for same in [slug, oauth_twin] {
actor.retain_selected_catalog_key_for(same);
assert_eq!(
actor.selected_catalog_key().as_deref(),
Some(oauth_twin),
"{same}: still names the selected entry — keep it"
);
}
// A rename to a DIFFERENT name makes the key stale: clear it, so the
// collided slug refuses (H-b) instead of resolving the old model.
actor.retain_selected_catalog_key_for("some-harness-model-name");
assert_eq!(
actor.selected_catalog_key(),
None,
"a stale key must be cleared, not carried into the next turn's \
platform lookup"
);
})
.await;
}
/// M3 — the FIRST-PARTY aux case must honour the session gate, which is what
/// the old shape did implicitly.
///
/// `stamp_session_local_sampler_fields` used to copy
/// `active_session_config.bearer_resolver`, and that field is `None` whenever
/// the gate is inactive. Re-pointing the aux resolver at the chokepoint (the
/// LEAK 1b fix) made it `Some(primary)` for the session's own coding endpoint
/// REGARDLESS of the gate — so a BYOK / api-key session with a `[model.*]` aux
/// entry carrying its own key on that endpoint had that key REPLACED by the
/// primary bearer on every image-describe / auto-mode-classifier / summary
/// request (`SamplingClient::post` overrides the auth header from the resolver).
///
/// Revert-to-red (production, compiles): delete the
/// `if is_primary_channel && !SessionTokenAuthGate::new(…).active()` early
/// return from `sampler_turn::aux_bearer_resolver_for` and the first two rows
/// below resolve `KIMI_TOKEN`.
#[tokio::test(flavor = "current_thread")]
async fn first_party_aux_resolver_honours_the_session_gate() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let coding_host = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
let aux_slug = "kigi-aux";
let mut info = crate::agent::config::ModelInfo::fallback(aux_slug);
info.id = None; // a `[model.kigi-aux]` block, not a registry entry
info.base_url = coding_host.to_string();
let aux_entry = crate::agent::config::ModelEntry {
info,
api_key: None,
env_key: None,
api_base_url: None,
};
// (case, ACP auth method, the aux model's own BYOK status, expected)
for (case, auth_method, byok, expect_resolver) in [
(
"an API-key session: the aux model's own key must survive",
"deepseek",
crate::agent::auth_method::ModelByok::NotByok,
false,
),
(
"a BYOK aux entry under a session method: its env_key wins",
"cached_token",
crate::agent::auth_method::ModelByok::Byok,
false,
),
(
"the first-party subscription aux channel: byte-identical",
"cached_token",
crate::agent::auth_method::ModelByok::NotByok,
true,
),
] {
let (_dir, actor, _rx) = actor_with_catalog(
vec![(aux_slug.to_string(), aux_entry.clone())],
aux_slug,
"",
)
.await;
actor
.auth_method_id
.store(Some(Arc::new(acp::AuthMethodId::new(auth_method))));
actor.model_auth_facts.replace(Some((
aux_slug.to_string(),
crate::agent::config::ModelAuthFacts {
byok,
auth_scheme: Default::default(),
},
)));
let resolved = actor.aux_bearer_resolver(aux_slug, coding_host);
assert_eq!(
resolved.is_some(),
expect_resolver,
"{case}: aux resolver presence on the session's own endpoint"
);
if let Some(resolver) = resolved {
assert_eq!(
resolver.current_bearer(),
Some(KIMI_TOKEN.to_string()),
"{case}: and when it IS kept it is the primary's, live"
);
}
}
})
.await;
}
/// M-aux (REGRESSION this remediation introduced) — an aux call must NOT evict
/// the SESSION model's memoized auth facts.
///
/// `SessionActor::model_auth_facts` is a SINGLE slot. When `aux_bearer_resolver`
/// began asking it about the AUX slug, a definite result overwrote the session
/// model's entry, and:
/// (a) the next `reconstruct_full_config` re-paid `load_effective_config()` +
/// `resolve_model_list()` — the per-turn disk read M7/M9 removed — on top
/// of the one the aux call itself paid; and
/// (b) the memo's documented purpose (a transient `Unknown` falling back to
/// the last DEFINITE value FOR THE SAME model_id) was defeated: with the
/// aux slug in the slot, the session model's `Unknown` degrades to
/// `endpoint_is_first_party`, which is `false` for every
/// subscription-OAuth host — the session loses its `bearer_resolver` and
/// 401s unrecoverably ~1h in, the failure L13 exists to prevent.
///
/// Round 3's deleted `repoint_aux_bearer_resolver` never touched the memo.
///
/// Revert-to-red (production, compiles): make `SessionActor::aux_bearer_resolver`
/// call `self.model_auth_facts(slug)` instead of `self.aux_model_auth_facts(slug)`
/// — the slot then names the aux slug and both assertions below fail.
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn an_aux_call_does_not_evict_the_session_models_auth_facts() {
let _env = anthropic_collision_env_guard();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let session_slug = "claude-opus-4-8";
let host = anthropic_collision_host();
let oauth_twin = "claude-pro-max/claude-opus-4-8";
let (_dir, actor, _rx) = actor_with_catalog(
vec![managed_entry(oauth_twin, session_slug, &host)],
oauth_twin,
"unused",
)
.await;
// The session model's DEFINITE facts, as a turn would have memoized
// them.
actor.model_auth_facts.replace(Some((
session_slug.to_string(),
crate::agent::config::ModelAuthFacts {
byok: crate::agent::auth_method::ModelByok::NotByok,
auth_scheme: Default::default(),
},
)));
// An aux turn: the auto-mode classifier / image-describe slug, which
// is NOT the session's model.
let _ = actor.aux_bearer_resolver("kigi-aux-classifier", &host);
let memo = actor.model_auth_facts.borrow();
let (cached_id, facts) = memo
.as_ref()
.expect("the session model's memo must survive an aux call");
assert_eq!(
cached_id, session_slug,
"an aux call evicted the SESSION model's memo: the next turn re-reads \
config from disk, and a transient Unknown loses its definite fallback"
);
assert_eq!(facts.byok, crate::agent::auth_method::ModelByok::NotByok);
})
.await;
}
crates/codegen/kigi-shell/src/session/acp_session_tests/session_bearer_leak_tests.rs
+29 -28
View File
@@ -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 <KIMI_TOKEN>` on this request.
#[tokio::test(flavor = "multi_thread")]
async fn deepseek_turn_under_a_kimi_session_sends_no_kimi_bearer_on_the_wire() {
@@ -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();
crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs
+1
View File
@@ -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,
crates/codegen/kigi-shell/src/session/commands.rs
+7
View File
@@ -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<String>,
use_concise: bool,
/// When `false`, skip the system prompt rewrite (concise/default swap).
/// Set to `false` for forked sessions so mid-session model switches
crates/codegen/kigi-shell/src/session/compaction.rs
+1
View File
@@ -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,