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");