update
This commit is contained in:
@@ -302,14 +302,26 @@ impl ModelByok {
|
||||
/// buffered token on every turn and 401s with `bad-credentials` until restart.
|
||||
/// It refreshes when `endpoint_is_first_party` — the request targets the
|
||||
/// first-party API, where sending the session token cannot leak to a
|
||||
/// third-party BYOK endpoint. A definite `NotByok` always refreshes (it only
|
||||
/// ever routes to the session endpoint); a definite `Byok` never does.
|
||||
/// third-party BYOK endpoint. A definite `NotByok` refreshes (within the
|
||||
/// session's own providers); a definite `Byok` never does.
|
||||
///
|
||||
/// `endpoint_takes_session_credential` ([`platform_takes_session_credential`])
|
||||
/// is the outer, non-negotiable guard and the reason the `NotByok` arm is safe:
|
||||
/// it is `false` for every API-key registry platform (deepseek, openai,
|
||||
/// anthropic, moonshot-*, …), whose models classify `NotByok` (they carry no
|
||||
/// `[model.*]` key) yet route to a THIRD-PARTY inference host while
|
||||
/// `oauth_registry::manager_for_model` falls through to the primary (Kimi)
|
||||
/// manager. Without this term the mainstream "session auth + API-key-platform
|
||||
/// model" configuration stamps the user's Kimi subscription bearer on every
|
||||
/// request to that host.
|
||||
pub fn session_token_auth_gate(
|
||||
is_session_based_method: bool,
|
||||
model_byok: ModelByok,
|
||||
endpoint_is_first_party: bool,
|
||||
endpoint_takes_session_credential: bool,
|
||||
) -> bool {
|
||||
is_session_based_method
|
||||
&& endpoint_takes_session_credential
|
||||
&& match model_byok {
|
||||
ModelByok::NotByok => true,
|
||||
ModelByok::Byok => false,
|
||||
@@ -317,6 +329,42 @@ pub fn session_token_auth_gate(
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a session bearer may EVER be stamped on a request routed to
|
||||
/// `platform` at `base_url` — i.e. whether the credential
|
||||
/// `oauth_registry::manager_for_model` resolves for such a model belongs to the
|
||||
/// host that receives it.
|
||||
///
|
||||
/// - a `uses_oauth` platform: `kimi-code` rides the PRIMARY session; each of the
|
||||
/// four subscription-OAuth platforms (claude-pro-max, openai-codex,
|
||||
/// github-copilot, xai-grok) rides its OWN pooled `AuthManager`. In both cases
|
||||
/// the resolved bearer belongs to the host being called, and mid-session
|
||||
/// refresh/401-recovery must stay live — so this is `true` even though those
|
||||
/// four have non-first-party base URLs.
|
||||
/// - every API-key registry platform: its credential is that platform's API key
|
||||
/// (resolved into the catalog entry), never a session bearer, and
|
||||
/// `manager_for_model` has no platform manager to route to. `false`.
|
||||
/// - `None` — a bare slug or a `[model.*]` config entry, which carries no
|
||||
/// platform at all — is decided by the ENDPOINT, never blanket-allowed: BYOK
|
||||
/// is `has_own_credentials()`, which probes `std::env::var` at call time, so a
|
||||
/// `[model.gpt-4o]` block with `base_url = "https://api.openai.com/v1"` and an
|
||||
/// unset / mistyped `env_key` classifies `NotByok` and would otherwise hand
|
||||
/// the Kimi subscription bearer to `api.openai.com`. Only the SESSION's own
|
||||
/// coding endpoint qualifies: [`crate::util::is_effective_coding_endpoint_url`]
|
||||
/// = the *effective* `KIGI_CODE_BASE_URL` deployment (so custom Kimi
|
||||
/// deployments keep the session bearer) plus loopback (local dev proxies and
|
||||
/// test mocks) plus the compiled production endpoint. Deliberately NOT
|
||||
/// `is_first_party_url`, which is production-only and would break every
|
||||
/// `KIGI_CODE_BASE_URL` deployment.
|
||||
pub fn platform_takes_session_credential(
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> bool {
|
||||
match platform {
|
||||
Some(platform) => platform.uses_oauth(),
|
||||
None => crate::util::is_effective_coding_endpoint_url(base_url),
|
||||
}
|
||||
}
|
||||
|
||||
pub const AUTH_ERROR_SESSION_EXPIRED: &str =
|
||||
"Session expired. Run `kigi login` to re-authenticate.";
|
||||
|
||||
@@ -683,15 +731,117 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn session_token_auth_gate_matrix() {
|
||||
// Session method + NotByok → refresh.
|
||||
assert!(session_token_auth_gate(true, ModelByok::NotByok, false));
|
||||
// Session method + NotByok → refresh (endpoint takes the session cred).
|
||||
assert!(session_token_auth_gate(
|
||||
true,
|
||||
ModelByok::NotByok,
|
||||
false,
|
||||
true
|
||||
));
|
||||
// Session method + Byok → never.
|
||||
assert!(!session_token_auth_gate(true, ModelByok::Byok, true));
|
||||
assert!(!session_token_auth_gate(true, ModelByok::Byok, true, true));
|
||||
// Session method + Unknown → only on first-party endpoints.
|
||||
assert!(session_token_auth_gate(true, ModelByok::Unknown, true));
|
||||
assert!(!session_token_auth_gate(true, ModelByok::Unknown, false));
|
||||
assert!(session_token_auth_gate(
|
||||
true,
|
||||
ModelByok::Unknown,
|
||||
true,
|
||||
true
|
||||
));
|
||||
assert!(!session_token_auth_gate(
|
||||
true,
|
||||
ModelByok::Unknown,
|
||||
false,
|
||||
true
|
||||
));
|
||||
// Non-session method → never.
|
||||
assert!(!session_token_auth_gate(false, ModelByok::NotByok, true));
|
||||
assert!(!session_token_auth_gate(
|
||||
false,
|
||||
ModelByok::NotByok,
|
||||
true,
|
||||
true
|
||||
));
|
||||
// An endpoint that does not take the session credential (every API-key
|
||||
// registry platform) is refused on EVERY arm — this is the outer guard
|
||||
// that keeps the Kimi subscription bearer off third-party hosts.
|
||||
for byok in [ModelByok::NotByok, ModelByok::Byok, ModelByok::Unknown] {
|
||||
for first_party in [false, true] {
|
||||
assert!(
|
||||
!session_token_auth_gate(true, byok, first_party, false),
|
||||
"byok={byok:?} first_party={first_party}: an API-key-platform \
|
||||
endpoint must never receive a session bearer"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The platform → "may a session bearer ride here?" classification.
|
||||
/// `kimi-code` rides the PRIMARY session; the four subscription-OAuth
|
||||
/// platforms ride their OWN pooled managers (so they keep a live
|
||||
/// bearer_resolver despite non-first-party base URLs); every API-key
|
||||
/// registry platform is refused, whatever the endpoint.
|
||||
#[test]
|
||||
fn platform_takes_session_credential_matrix() {
|
||||
let first_party = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
|
||||
for id in [
|
||||
"kimi-code",
|
||||
"claude-pro-max",
|
||||
"openai-codex",
|
||||
"github-copilot",
|
||||
"xai-grok",
|
||||
] {
|
||||
let platform = kigi_models::PlatformId::parse(id).expect("known platform");
|
||||
for url in [first_party, "https://api.anthropic.com/v1"] {
|
||||
assert!(
|
||||
platform_takes_session_credential(Some(platform), url),
|
||||
"{id} rides a session credential (primary or its own pool)"
|
||||
);
|
||||
}
|
||||
}
|
||||
for platform in kigi_models::PlatformId::ALL {
|
||||
if platform.uses_oauth() {
|
||||
continue;
|
||||
}
|
||||
for url in [first_party, "https://api.deepseek.com/v1"] {
|
||||
assert!(
|
||||
!platform_takes_session_credential(Some(platform), url),
|
||||
"{} is an API-key platform — no session bearer may ride to it",
|
||||
platform.as_str()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// C2 regression: a platform-less model (a bare slug or a `[model.*]` entry)
|
||||
/// is decided by the ENDPOINT, never blanket-allowed. A `[model.gpt-4o]`
|
||||
/// block whose `env_key` is unset classifies `NotByok`, so before this the
|
||||
/// `None` arm handed the Kimi subscription bearer to `api.openai.com`.
|
||||
/// Custom `KIGI_CODE_BASE_URL` deployments and local dev proxies must still
|
||||
/// keep it (that is why the predicate is not `is_first_party_url`).
|
||||
#[test]
|
||||
fn platform_less_model_takes_the_session_credential_only_on_its_own_endpoint() {
|
||||
for url in [
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
"http://127.0.0.1:8080/v1",
|
||||
"http://localhost:3000/v1",
|
||||
"http://[::1]:9000/v1",
|
||||
] {
|
||||
assert!(
|
||||
platform_takes_session_credential(None, url),
|
||||
"{url} is the session's own endpoint (or a local proxy) — unchanged"
|
||||
);
|
||||
}
|
||||
for url in [
|
||||
"https://api.openai.com/v1",
|
||||
"https://api.deepseek.com/v1",
|
||||
"https://api.anthropic.com/v1",
|
||||
"https://api.moonshot.cn/v1",
|
||||
"",
|
||||
] {
|
||||
assert!(
|
||||
!platform_takes_session_credential(None, url),
|
||||
"LEAK: {url} is a third-party host — no session bearer may ride there"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// RAII guard restoring an env var on drop (panic-safe).
|
||||
|
||||
@@ -2779,13 +2779,20 @@ pub fn default_model_entries(endpoints: &EndpointsConfig) -> IndexMap<String, Mo
|
||||
}
|
||||
/// Resolve a model against the available model map.
|
||||
/// Checks the map key (id) first, then falls back to a slug scan.
|
||||
///
|
||||
/// The slug scan takes the LAST match, exactly like the picker's
|
||||
/// [`crate::agent::models::resolve_catalog_key`]. Duplicate slugs across
|
||||
/// platforms are by design (`moonshot-cn` and `moonshot-ai` list the same ids;
|
||||
/// so do every API-key platform and its subscription-OAuth twin), and a
|
||||
/// first-match scan here made the auth layer resolve a DIFFERENT entry than the
|
||||
/// one the picker selected — the H5 slug collision. One direction, one answer.
|
||||
pub fn find_model_by_id<'a>(
|
||||
models: &'a IndexMap<String, ModelEntry>,
|
||||
model_id: &str,
|
||||
) -> Option<&'a ModelEntry> {
|
||||
models
|
||||
.get(model_id)
|
||||
.or_else(|| models.values().find(|m| m.model == model_id))
|
||||
.or_else(|| models.values().rev().find(|m| m.model == model_id))
|
||||
}
|
||||
/// Whether the EFFECTIVE Auto-mode classifier model supports reasoning effort:
|
||||
/// the model actually routed to (`aux_model` when the aux sampler resolved) else
|
||||
@@ -4643,14 +4650,22 @@ reasoning_effort = "low"
|
||||
});
|
||||
(dir, manager)
|
||||
}
|
||||
/// LEAK 1a (aux/summary model): the aux `session_key` is resolved by the aux
|
||||
/// model's OWN platform (as `build_summary_client` / `resolve_aux_sampler_config`
|
||||
/// now do). A grok (oauth-platform) aux model's resolved sampler `api_key` is
|
||||
/// therefore grok's own token or `None` — NEVER the primary Kimi key — while a
|
||||
/// first-party / non-oauth aux model still gets the primary (byte-identical).
|
||||
/// LEAK 1a (aux/summary model `api_key` channel): the aux `session_key` is
|
||||
/// resolved by the aux model's OWN platform AND endpoint (as
|
||||
/// `build_summary_client` / `resolve_aux_sampler_config` now do), so it is
|
||||
/// never the primary Kimi key on a host that does not own it.
|
||||
///
|
||||
/// - a grok (oauth-platform) aux model → grok's own pooled token or `None`;
|
||||
/// - a `moonshot-cn` (API-key platform) aux model → NO session key at all
|
||||
/// (pre-fix it received `kimi-tok`, which `resolve_credentials` then
|
||||
/// stamped as the `api_key` on an `api.moonshot.cn` request);
|
||||
/// - a `kimi-code` aux model → the primary, byte-identical.
|
||||
///
|
||||
/// Revert-to-red: dropping the `platform_takes_session_credential` term
|
||||
/// from `session_key_for_endpoint` makes the moonshot assertion see
|
||||
/// `Some("kimi-tok")`.
|
||||
#[tokio::test]
|
||||
async fn aux_model_session_key_is_platform_scoped_never_leaking_kimi() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let (_kd, kimi) = kimi_primary("kimi-tok");
|
||||
let endpoints = EndpointsConfig::default();
|
||||
// A grok aux catalog entry (managed id → oauth platform), no own key.
|
||||
@@ -4658,9 +4673,9 @@ reasoning_effort = "low"
|
||||
grok.info.id = Some("xai-grok/grok-4-latest".to_string());
|
||||
let mut grok_catalog = IndexMap::new();
|
||||
grok_catalog.insert("grok".to_string(), grok);
|
||||
let grok_key = crate::auth::oauth_registry::session_key_for_model(
|
||||
home.path(),
|
||||
"xai-grok/grok-4-latest",
|
||||
let grok_key = crate::auth::oauth_registry::session_key_for_catalog_model(
|
||||
&grok_catalog,
|
||||
"grok",
|
||||
Some(&kimi),
|
||||
);
|
||||
let grok_cfg = resolve_aux_model_sampling_config(
|
||||
@@ -4675,10 +4690,11 @@ reasoning_effort = "low"
|
||||
Some("kimi-tok"),
|
||||
"a grok aux model must never receive the primary Kimi session token",
|
||||
);
|
||||
// A non-oauth aux catalog entry still resolves to the primary token.
|
||||
// An API-key registry platform gets NO session key — the pre-fix
|
||||
// behaviour handed it the primary Kimi token on api.moonshot.cn.
|
||||
let mut k2 = test_model_entry(
|
||||
"kimi-k2-0905-preview",
|
||||
"https://vendor/v1",
|
||||
"https://api.moonshot.cn/v1",
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
@@ -4686,23 +4702,43 @@ reasoning_effort = "low"
|
||||
k2.info.id = Some("moonshot-cn/kimi-k2".to_string());
|
||||
let mut k2_catalog = IndexMap::new();
|
||||
k2_catalog.insert("k2".to_string(), k2);
|
||||
let k2_key = crate::auth::oauth_registry::session_key_for_model(
|
||||
home.path(),
|
||||
"moonshot-cn/kimi-k2",
|
||||
assert_eq!(
|
||||
crate::auth::oauth_registry::session_key_for_catalog_model(
|
||||
&k2_catalog,
|
||||
"k2",
|
||||
Some(&kimi),
|
||||
),
|
||||
None,
|
||||
"LEAK: an API-key-platform aux model must receive no session token",
|
||||
);
|
||||
// The first-party subscription channel is byte-identical.
|
||||
let mut kimi_code = test_model_entry(
|
||||
"kimi-for-coding",
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
kimi_code.info.id = Some("kimi-code/kimi-for-coding".to_string());
|
||||
let mut kimi_catalog = IndexMap::new();
|
||||
kimi_catalog.insert("kimi-for-coding".to_string(), kimi_code);
|
||||
let kimi_key = crate::auth::oauth_registry::session_key_for_catalog_model(
|
||||
&kimi_catalog,
|
||||
"kimi-for-coding",
|
||||
Some(&kimi),
|
||||
);
|
||||
let k2_cfg = resolve_aux_model_sampling_config(
|
||||
"k2",
|
||||
&k2_catalog,
|
||||
let kimi_cfg = resolve_aux_model_sampling_config(
|
||||
"kimi-for-coding",
|
||||
&kimi_catalog,
|
||||
&endpoints,
|
||||
k2_key.as_deref(),
|
||||
kimi_key.as_deref(),
|
||||
None,
|
||||
)
|
||||
.expect("non-oauth aux resolves via the primary session token");
|
||||
.expect("a kimi-code aux model resolves via the primary session token");
|
||||
assert_eq!(
|
||||
k2_cfg.api_key.as_deref(),
|
||||
kimi_cfg.api_key.as_deref(),
|
||||
Some("kimi-tok"),
|
||||
"a non-oauth aux model must still receive the primary session token",
|
||||
"the first-party subscription aux path must be byte-identical",
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
|
||||
@@ -1917,6 +1917,51 @@ pub(crate) fn resolve_catalog_key(
|
||||
.map(|(key, _)| acp::ModelId::new(key.clone()))
|
||||
}
|
||||
|
||||
/// The managed catalog key (`{platform}/{model}`) a routing slug belongs to.
|
||||
///
|
||||
/// H5: `SamplingConfig::model` is the BARE routing slug, never the catalog key,
|
||||
/// and duplicate slugs across platforms are BY DESIGN — the registry guarantees
|
||||
/// an API-key platform and its subscription-OAuth twin list the SAME ids
|
||||
/// (`xai`/`xai-grok`, `anthropic`/`claude-pro-max`, `openai`/`openai-codex`),
|
||||
/// and `PlatformId::ALL` orders every API-key platform FIRST. A slug scan
|
||||
/// therefore resolves the WRONG platform for a user holding both credentials:
|
||||
/// the OAuth twin loses its live `bearer_resolver` (no mid-session refresh → an
|
||||
/// unrecoverable 401 ~1h in), its Messages adaptation and its Copilot/Codex
|
||||
/// identity headers.
|
||||
///
|
||||
/// `current_key` — [`ModelsManager::current_model_id`], the catalog key the
|
||||
/// picker actually selected — is therefore authoritative whenever it names this
|
||||
/// slug. Anything else (aux models, subagent overrides, unlisted slugs) falls
|
||||
/// back to the picker's OWN lookup, [`resolve_catalog_key`], so the auth layer
|
||||
/// and the picker can never resolve different entries.
|
||||
pub(crate) fn managed_key_for_slug(
|
||||
models: &IndexMap<String, ModelEntry>,
|
||||
current_key: Option<&str>,
|
||||
slug: &str,
|
||||
) -> Option<String> {
|
||||
if let Some(entry) = current_key.and_then(|key| models.get(key))
|
||||
&& (entry.info.model == slug || current_key == Some(slug))
|
||||
{
|
||||
return entry.info.id.clone();
|
||||
}
|
||||
let key = resolve_catalog_key(models, &acp::ModelId::new(slug.to_string()))?;
|
||||
models.get(key.0.as_ref())?.info.id.clone()
|
||||
}
|
||||
|
||||
/// The registry platform a routing slug resolves to, via [`managed_key_for_slug`].
|
||||
/// `None` for a bare / `[model.*]` / unlisted model. Single definition shared by
|
||||
/// the session actor's inference-auth chokepoints and the aux/summary paths, so
|
||||
/// the gate, the manager and the wire adaptations can never disagree.
|
||||
pub(crate) fn platform_for_slug(
|
||||
models: &IndexMap<String, ModelEntry>,
|
||||
current_key: Option<&str>,
|
||||
slug: &str,
|
||||
) -> Option<kigi_models::PlatformId> {
|
||||
let key = managed_key_for_slug(models, current_key, slug);
|
||||
kigi_models::parse_managed_model_key(key.as_deref().unwrap_or(slug))
|
||||
.map(|(platform, _)| platform)
|
||||
}
|
||||
|
||||
/// Catalog key for a persisted session model id, restricted to **selectable**
|
||||
/// entries. A selectable exact-key match wins (as in [`resolve_catalog_key`]);
|
||||
/// otherwise the last selectable entry whose routing slug matches `id`, so a
|
||||
@@ -3668,6 +3713,52 @@ mod tests {
|
||||
assert!(!info.visible_for_auth(false));
|
||||
}
|
||||
|
||||
/// SHIP-BLOCKER regression: a user who signed in with ONLY a Claude Pro/Max
|
||||
/// subscription has no PRIMARY (Kimi) session, so `is_session_auth()` is
|
||||
/// false. Stamping `supported_in_api = !uses_oauth()` therefore hid every
|
||||
/// one of their models — `available()` returned an empty picker and the
|
||||
/// whole subscription-OAuth feature was dead for its target user. The
|
||||
/// claude-pro-max entry must be visible with no primary session at all.
|
||||
#[test]
|
||||
fn claude_pro_max_only_user_sees_their_models_in_the_picker() {
|
||||
let wire: kigi_models::WireModel =
|
||||
serde_json::from_value(serde_json::json!({ "id": "claude-opus-4-8" }))
|
||||
.expect("wire model fixture");
|
||||
let entry_config = crate::agent::models_fetch::platform_wire_model_to_entry(
|
||||
kigi_models::PlatformId::ClaudeProMax,
|
||||
wire,
|
||||
"https://api.anthropic.com/v1",
|
||||
);
|
||||
let entry = ModelEntry::from_config_entry(&entry_config);
|
||||
let key = entry
|
||||
.info
|
||||
.id
|
||||
.clone()
|
||||
.expect("platform entries carry a managed catalog key");
|
||||
assert_eq!(key, "claude-pro-max/claude-opus-4-8");
|
||||
let mut catalog = IndexMap::new();
|
||||
catalog.insert(key.clone(), entry);
|
||||
|
||||
// Empty home ⇒ the primary AuthManager holds no credential at all.
|
||||
let home = tempfile::tempdir().expect("tempdir");
|
||||
let mgr = ModelsManager::new(
|
||||
None,
|
||||
catalog,
|
||||
acp::ModelId::new(Arc::from(key.clone())),
|
||||
Arc::new(AuthManager::new(home.path(), KimiCodeConfig::default())),
|
||||
config::Config::default(),
|
||||
);
|
||||
assert!(
|
||||
!mgr.is_session_auth(),
|
||||
"a claude-pro-max-only user has no PRIMARY (Kimi) OAuth session"
|
||||
);
|
||||
assert!(
|
||||
mgr.available()
|
||||
.contains_key(&acp::ModelId::new(Arc::from(key.clone()))),
|
||||
"the claude-pro-max model must reach the picker without a primary session"
|
||||
);
|
||||
}
|
||||
|
||||
// ── duplicate model slug re-keying (A/B experiment "auto" alias) ──
|
||||
|
||||
fn make_entry_config(model: &str, name: Option<&str>) -> config::ModelEntryConfig {
|
||||
|
||||
@@ -645,9 +645,15 @@ pub(crate) fn platform_wire_model_to_entry(
|
||||
inference_idle_timeout_secs: None,
|
||||
max_retries: None,
|
||||
hidden: false,
|
||||
// Subscription models require the OAuth session; open-platform
|
||||
// models are usable by API-key users.
|
||||
supported_in_api: !platform.uses_oauth(),
|
||||
// `supported_in_api: false` hides a model unless the PRIMARY session is
|
||||
// an OAuth session (`ModelInfo::visible_for_auth`). Only `kimi-code`
|
||||
// rides that primary session, so only it may be gated on it. Every
|
||||
// other OAuth platform (claude-pro-max, openai-codex, github-copilot,
|
||||
// xai-grok) carries its OWN pooled credential, and its models only
|
||||
// enter the catalog once THAT provider is signed in — gating them on
|
||||
// the Kimi session would hide every model from a user who signed in
|
||||
// with only a Claude/ChatGPT/Copilot/Grok subscription.
|
||||
supported_in_api: platform != kigi_models::PlatformId::KimiCode,
|
||||
supports_backend_search: false,
|
||||
compactions_remaining: None,
|
||||
compaction_at_tokens: None,
|
||||
@@ -898,6 +904,16 @@ fn get_string_map(
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Whether a freshly-fetched platform entry shows up in the picker for a
|
||||
/// user whose PRIMARY session is NOT an OAuth session (`is_session_auth ==
|
||||
/// false`): an API-key user, or — the case that made this a ship-blocker —
|
||||
/// someone who signed in with ONLY a Claude Pro/Max, ChatGPT, Copilot, or
|
||||
/// Grok subscription. `ModelInfo::visible_for_auth` is the picker's real
|
||||
/// predicate (`agent/models.rs` → `available()`).
|
||||
fn visible_to_non_primary_session_user(entry: &crate::agent::config::ModelEntryConfig) -> bool {
|
||||
crate::agent::config::ModelEntry::from_config_entry(entry).visible_for_auth(false)
|
||||
}
|
||||
|
||||
/// OpenAI-cycle e2e (mock wire): a polluted bare-id `/models` listing +
|
||||
/// a models.dev refresh produce a catalog with ONLY chat models, enriched
|
||||
/// context windows / efforts, and the Responses backend — the full
|
||||
@@ -1277,8 +1293,14 @@ mod tests {
|
||||
"enrichment fills the context window from models.dev anthropic"
|
||||
);
|
||||
assert!(
|
||||
!opus.supported_in_api,
|
||||
"subscription (uses_oauth) models require the OAuth session"
|
||||
opus.supported_in_api,
|
||||
"claude-pro-max carries its OWN pooled credential — it must NOT be \
|
||||
gated on the primary (Kimi) session"
|
||||
);
|
||||
assert!(
|
||||
visible_to_non_primary_session_user(opus),
|
||||
"a Claude-Pro/Max-only user has no primary OAuth session; their \
|
||||
models must still appear in the picker"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1409,8 +1431,14 @@ mod tests {
|
||||
"context window comes from models.dev github-copilot enrichment"
|
||||
);
|
||||
assert!(
|
||||
!entry.supported_in_api,
|
||||
"subscription (uses_oauth) models require the OAuth session"
|
||||
entry.supported_in_api,
|
||||
"github-copilot carries its OWN pooled credential — it must NOT be \
|
||||
gated on the primary (Kimi) session"
|
||||
);
|
||||
assert!(
|
||||
visible_to_non_primary_session_user(entry),
|
||||
"a Copilot-only user has no primary OAuth session; their models \
|
||||
must still appear in the picker"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1484,8 +1512,14 @@ mod tests {
|
||||
assert_eq!(sol.context_window.get(), 272_000);
|
||||
assert_eq!(sol.name.as_deref(), Some("GPT-5.6-Sol"));
|
||||
assert!(
|
||||
!sol.supported_in_api,
|
||||
"subscription (uses_oauth) models require the OAuth session"
|
||||
sol.supported_in_api,
|
||||
"openai-codex carries its OWN pooled credential — it must NOT be \
|
||||
gated on the primary (Kimi) session"
|
||||
);
|
||||
assert!(
|
||||
visible_to_non_primary_session_user(sol),
|
||||
"a ChatGPT/Codex-only user has no primary OAuth session; their \
|
||||
models must still appear in the picker"
|
||||
);
|
||||
assert!(sol.supports_reasoning_effort);
|
||||
assert_eq!(
|
||||
@@ -3563,8 +3597,14 @@ mod tests {
|
||||
"an OAuth channel carries no api-key env"
|
||||
);
|
||||
assert!(
|
||||
!entry.supported_in_api,
|
||||
"subscription models require the OAuth session (not the public API)"
|
||||
entry.supported_in_api,
|
||||
"xai-grok carries its OWN pooled credential — it must NOT be gated \
|
||||
on the primary (Kimi) session"
|
||||
);
|
||||
assert!(
|
||||
visible_to_non_primary_session_user(entry),
|
||||
"a Grok-only user has no primary OAuth session; their models must \
|
||||
still appear in the picker"
|
||||
);
|
||||
// Passthrough dialect (identical to the API-key xai wire).
|
||||
let model_entry = crate::agent::config::ModelEntry::from_config_entry(entry);
|
||||
|
||||
@@ -25,17 +25,19 @@ impl MvpAgent {
|
||||
primary: &SamplingConfig,
|
||||
) -> Result<(OaiCompatClient, String), acp::Error> {
|
||||
let slug = self.resolve_session_summary_model();
|
||||
// Resolve the aux token by the summary model's OWN platform: a grok
|
||||
// (oauth-platform) summary model draws its pooled grok token or `None`
|
||||
// — NEVER the primary Kimi session token (which `resolve_credentials`
|
||||
// would otherwise stamp onto an api.x.ai request). A first-party /
|
||||
// non-oauth summary model still gets the primary (byte-identical).
|
||||
let session_key = crate::auth::oauth_registry::session_key_for_model(
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
let models = self.models_manager.models();
|
||||
// Resolve the aux token by the summary model's OWN platform AND
|
||||
// endpoint: a grok (oauth-platform) summary model draws its pooled grok
|
||||
// token or `None`, and an API-key registry platform draws NOTHING —
|
||||
// NEVER the primary Kimi session token (which `resolve_credentials`
|
||||
// would otherwise stamp onto an api.x.ai / api.deepseek.com request).
|
||||
// The first-party subscription channel still gets the primary
|
||||
// (byte-identical).
|
||||
let session_key = crate::auth::oauth_registry::session_key_for_catalog_model(
|
||||
&models,
|
||||
&slug,
|
||||
Some(&self.auth_manager),
|
||||
);
|
||||
let models = self.models_manager.models();
|
||||
let endpoints = self.models_manager.endpoints();
|
||||
let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone();
|
||||
let config = match crate::agent::config::resolve_aux_model_sampling_config(
|
||||
@@ -47,7 +49,17 @@ impl MvpAgent {
|
||||
) {
|
||||
Some(mut cfg) => {
|
||||
cfg.attribution_callback = primary.attribution_callback.clone();
|
||||
cfg.bearer_resolver = primary.bearer_resolver.clone();
|
||||
// H4: the SESSION model's bearer_resolver must not ride to a
|
||||
// summary model on a different provider —
|
||||
// `SamplingClient::post` REPLACES the request's auth header
|
||||
// from it, overwriting the summary model's own resolved key on
|
||||
// ITS host. Route through the one aux-resolver decision.
|
||||
cfg.bearer_resolver =
|
||||
crate::session::acp_session::sampler_turn::aux_bearer_resolver(
|
||||
primary.bearer_resolver.clone(),
|
||||
crate::agent::models::platform_for_slug(&models, None, &slug),
|
||||
&cfg.base_url,
|
||||
);
|
||||
cfg.max_retries = primary.max_retries;
|
||||
cfg
|
||||
}
|
||||
@@ -669,31 +681,51 @@ impl MvpAgent {
|
||||
);
|
||||
Ok(entry.clone())
|
||||
}
|
||||
/// Resolve the SESSION token for `model` by the model's OWN platform — the
|
||||
/// single guard against the api_key-channel token leak.
|
||||
/// Resolve the SESSION token for `model` by the model's OWN platform AND
|
||||
/// endpoint — the single guard against the api_key-channel token leak.
|
||||
///
|
||||
/// An oauth-platform model (xai-grok) draws its session token from ITS OWN
|
||||
/// process-global pool manager (build-on-demand from the on-disk grok token,
|
||||
/// proactively refreshed), INDEPENDENT of the primary `auth_method`; when
|
||||
/// that provider has no stored session the token is `None` — NEVER the
|
||||
/// primary Kimi key. Every other model (first-party / Kimi) uses the primary
|
||||
/// session manager, and only under a session-based auth method —
|
||||
/// byte-identical to the pre-fix path. SECURITY: the resolved token is never
|
||||
/// logged.
|
||||
/// - an oauth-platform model (xai-grok, claude-pro-max, github-copilot,
|
||||
/// openai-codex) draws its session token from ITS OWN process-global pool
|
||||
/// manager (built on demand from the on-disk token, proactively
|
||||
/// refreshed), INDEPENDENT of the primary `auth_method`; when that
|
||||
/// provider has no stored session the token is `None` — NEVER the primary
|
||||
/// Kimi key;
|
||||
/// - an endpoint that does not take a session credential at all
|
||||
/// ([`crate::agent::auth_method::platform_takes_session_credential`] —
|
||||
/// every API-key registry platform, and any `[model.*]` block pointed at a
|
||||
/// third-party host) gets `None`. This is C1: `resolve_credentials` takes
|
||||
/// the `else if let Some(key) = session_key` arm and sets
|
||||
/// `api_key = <Kimi bearer>` with the THIRD-PARTY `base_url`, which
|
||||
/// `SamplingClient` then builds into `Authorization: Bearer …`. It is
|
||||
/// reachable with ZERO configuration: `default_models.json` bundles
|
||||
/// `moonshot-cn/*` + `moonshot-ai/*` entries that a Kimi-subscription user
|
||||
/// sees on first launch / offline;
|
||||
/// - the first-party subscription channel (kimi-code, a `KIGI_CODE_BASE_URL`
|
||||
/// deployment, a loopback proxy) uses the primary session manager, and
|
||||
/// only under a session-based auth method — byte-identical to the pre-fix
|
||||
/// path.
|
||||
///
|
||||
/// SECURITY: the resolved token is never logged.
|
||||
fn session_token_for_model(&self, model: &ModelEntry) -> Option<crate::auth::KimiAuth> {
|
||||
if let Some(oauth) = model
|
||||
.info()
|
||||
let info = model.info();
|
||||
let platform = info
|
||||
.id
|
||||
.as_deref()
|
||||
.and_then(kigi_models::parse_managed_model_key)
|
||||
.and_then(|(platform, _)| platform.oauth())
|
||||
{
|
||||
.map(|(platform, _)| platform);
|
||||
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
|
||||
return crate::auth::oauth_registry::global_manager_for(
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
oauth,
|
||||
)
|
||||
.current_or_expired();
|
||||
}
|
||||
if !crate::agent::auth_method::platform_takes_session_credential(
|
||||
platform,
|
||||
&info.base_url,
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
if self.is_session_based_auth() {
|
||||
self.auth_manager.current_or_expired()
|
||||
} else {
|
||||
|
||||
@@ -1197,18 +1197,26 @@ async fn prepare_sampling_config_never_stamps_kimi_key_on_grok_model() {
|
||||
|
||||
let endpoints = EndpointsConfig::default();
|
||||
|
||||
// First-party Kimi model (non-oauth platform): the primary session key IS
|
||||
// its api_key — the byte-identical primary path, and proof the Kimi token is
|
||||
// First-party SUBSCRIPTION model (kimi-code): the primary session key IS its
|
||||
// api_key — the byte-identical primary path, and proof the Kimi token is
|
||||
// live (so it WOULD leak if mis-routed onto a grok request). This assertion
|
||||
// also confirms the session-based primary path is active.
|
||||
let mut kimi_model = ModelEntry::fallback("kimi-k2-0905-preview", &endpoints);
|
||||
kimi_model.info.id = Some("moonshot-cn/kimi-k2-0905-preview".to_string());
|
||||
//
|
||||
// This used to use `moonshot-cn/kimi-k2-0905-preview` and assert the SAME
|
||||
// thing, which encoded the C1 defect: moonshot-cn is an API-key registry
|
||||
// platform on `api.moonshot.cn`, NOT first-party, so "must carry the primary
|
||||
// session key" was asserting the leak. `api_key_channel_leak_tests` now pins
|
||||
// the opposite for every moonshot entry.
|
||||
let mut kimi_model = ModelEntry::fallback("kimi-for-coding", &endpoints);
|
||||
kimi_model.info.id = Some("kimi-code/kimi-for-coding".to_string());
|
||||
kimi_model.info.base_url = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url.to_string();
|
||||
assert!(!kimi_model.has_own_credentials());
|
||||
let kimi_cfg = agent.prepare_sampling_config_for_model(&kimi_model, None);
|
||||
assert_eq!(
|
||||
kimi_cfg.api_key.as_deref(),
|
||||
Some(KIMI_KEY),
|
||||
"a first-party Kimi model must carry the primary session key (primary path unchanged)"
|
||||
"the first-party subscription model must carry the primary session key \
|
||||
(primary path unchanged)"
|
||||
);
|
||||
|
||||
// xai-grok model (oauth platform): the session token resolves from its OWN
|
||||
@@ -1280,6 +1288,9 @@ async fn ensure_plugin_registry_lazily_populates_snapshot() {
|
||||
);
|
||||
}
|
||||
mod subagent_spawn_context_tests;
|
||||
/// LEAK guard for the `api_key` channel (C1/C2), through the real
|
||||
/// `prepare_sampling_config_for_model` resolution path.
|
||||
mod api_key_channel_leak_tests;
|
||||
/// No load in flight and no session → the wait returns immediately
|
||||
/// (the caller then surfaces "unknown session id" exactly as before).
|
||||
#[tokio::test]
|
||||
|
||||
@@ -0,0 +1,313 @@
|
||||
//! LEAK GUARD (`api_key` channel) — C1/C2, driven through the REAL resolution
|
||||
//! path, `MvpAgent::prepare_sampling_config_for_model`.
|
||||
//!
|
||||
//! This is the channel the `bearer_resolver` guard does NOT close, and the one
|
||||
//! the first round of leak tests assumed away by hand-stamping a provider key
|
||||
//! into chat state. The chain: `session_token_for_model` used to fall through to
|
||||
//! `self.auth_manager.current_or_expired()` (the primary Kimi bearer) for every
|
||||
//! non-OAuth model, `resolve_credentials` then took its
|
||||
//! `else if let Some(key) = session_key` arm and set `api_key = <Kimi token>`
|
||||
//! with the THIRD-PARTY `base_url`, and `SamplingClient` builds
|
||||
//! `Authorization: Bearer <api_key>` straight into `default_headers` — which
|
||||
//! `post()` only overrides when a resolver exists, so `bearer_resolver: None`
|
||||
//! does not save it.
|
||||
//!
|
||||
//! Nothing here stamps a credential by hand: every assertion reads what the
|
||||
//! resolution path actually produced.
|
||||
|
||||
use super::super::*;
|
||||
use crate::agent::auth_method::{
|
||||
CACHED_TOKEN_AUTH_METHOD_ID, HOUSE_API_KEY_ENV_VAR, LEGACY_XAI_API_KEY_ENV_VAR,
|
||||
XAI_API_KEY_ENV_VAR,
|
||||
};
|
||||
use crate::agent::config::{Config as AgentConfig, EndpointsConfig, EnvKeys, ModelEntry};
|
||||
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
|
||||
use kigi_test_support::EnvGuard;
|
||||
|
||||
const KIMI_TOKEN: &str = "kimi-subscription-token-DO-NOT-LEAK";
|
||||
|
||||
/// Ambient BYOK env vars unset, so a model with no resolvable credential ends up
|
||||
/// with `api_key == None` rather than a global-key fallback that could mask the
|
||||
/// leak under test. Every test holding these must be `#[serial]`.
|
||||
fn without_ambient_byok_env() -> [EnvGuard; 3] {
|
||||
[
|
||||
EnvGuard::unset(HOUSE_API_KEY_ENV_VAR),
|
||||
EnvGuard::unset(XAI_API_KEY_ENV_VAR),
|
||||
EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR),
|
||||
]
|
||||
}
|
||||
|
||||
/// An `MvpAgent` on a session-based (`cached_token`) ACP method holding a live
|
||||
/// Kimi subscription bearer — the mainstream configuration in which the leak
|
||||
/// fires. `(tempdir, agent)`; the tempdir is the auth store and is returned so
|
||||
/// the caller keeps it alive.
|
||||
fn kimi_session_agent() -> (tempfile::TempDir, MvpAgent) {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let auth_manager = std::sync::Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
|
||||
auth_manager.hot_swap(KimiAuth {
|
||||
key: KIMI_TOKEN.to_string(),
|
||||
auth_mode: AuthMode::OAuth,
|
||||
refresh_token: Some("rt".into()),
|
||||
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
|
||||
..KimiAuth::test_default()
|
||||
});
|
||||
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let agent = MvpAgent::new(
|
||||
GatewaySender::new(tx),
|
||||
&AgentConfig::default(),
|
||||
auth_manager,
|
||||
None,
|
||||
)
|
||||
.expect("valid test config");
|
||||
agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID));
|
||||
(dir, agent)
|
||||
}
|
||||
|
||||
/// A catalog entry as `resolve_model_list` builds one for a fetched registry
|
||||
/// model: managed catalog key, platform base URL, no credential of its own.
|
||||
fn platform_entry(catalog_key: &str, slug: &str, base_url: &str) -> ModelEntry {
|
||||
let mut entry = ModelEntry::fallback(slug, &EndpointsConfig::default());
|
||||
entry.info.id = Some(catalog_key.to_string());
|
||||
entry.info.base_url = base_url.to_string();
|
||||
entry
|
||||
}
|
||||
|
||||
/// C1, the ZERO-CONFIGURATION repro. `default_models.json` bundles
|
||||
/// `moonshot-cn/*` and `moonshot-ai/*` entries with `api_key: None`, and
|
||||
/// `resolve_model_list` keeps the bundled defaults whenever no catalog fetch has
|
||||
/// succeeded — so on first launch / offline a Kimi-subscription user sees them
|
||||
/// in the picker with no configuration whatsoever. Selecting one used to send
|
||||
/// `Authorization: Bearer <Kimi OAuth token>` to `api.moonshot.cn`, which is NOT
|
||||
/// first-party.
|
||||
///
|
||||
/// Revert-to-red: dropping the `platform_takes_session_credential` guard from
|
||||
/// `session_token_for_model` makes every `api_key` below `Some(KIMI_TOKEN)`.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn bundled_default_moonshot_models_never_carry_the_kimi_bearer() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let (_dir, agent) = kimi_session_agent();
|
||||
|
||||
let bundled = crate::agent::config::default_model_entries(&EndpointsConfig::default());
|
||||
let moonshot: Vec<_> = bundled
|
||||
.iter()
|
||||
.filter(|(key, _)| key.starts_with("moonshot-cn/") || key.starts_with("moonshot-ai/"))
|
||||
.collect();
|
||||
assert_eq!(
|
||||
moonshot.len(),
|
||||
4,
|
||||
"default_models.json still bundles the four moonshot open-platform entries"
|
||||
);
|
||||
|
||||
for (key, entry) in moonshot {
|
||||
assert!(
|
||||
!entry.has_own_credentials(),
|
||||
"{key}: the bundled entry carries no credential of its own"
|
||||
);
|
||||
assert!(
|
||||
!crate::util::is_effective_coding_endpoint_url(&entry.info().base_url),
|
||||
"{key}: routes to a third-party host ({})",
|
||||
entry.info().base_url
|
||||
);
|
||||
let cfg = agent.prepare_sampling_config_for_model(entry, None);
|
||||
assert_ne!(
|
||||
cfg.api_key.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"LEAK: selecting the bundled {key} sent the Kimi subscription bearer to {}",
|
||||
entry.info().base_url
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.base_url,
|
||||
entry.info().base_url,
|
||||
"{key}: still routes to its own host (the fix must not reroute traffic)"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// C1 across the API-key registry platform shapes a fetched catalog produces.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn api_key_platform_models_never_carry_the_kimi_bearer_as_api_key() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let (_dir, agent) = kimi_session_agent();
|
||||
|
||||
for (catalog_key, slug, base_url) in [
|
||||
("deepseek/deepseek-chat", "deepseek-chat", "https://api.deepseek.com/v1"),
|
||||
("openai/gpt-5.2", "gpt-5.2", "https://api.openai.com/v1"),
|
||||
("anthropic/claude-opus-4-8", "claude-opus-4-8", "https://api.anthropic.com/v1"),
|
||||
("groq/llama-4", "llama-4", "https://api.groq.com/openai/v1"),
|
||||
("xai/grok-4.5", "grok-4.5", "https://api.x.ai/v1"),
|
||||
] {
|
||||
let entry = platform_entry(catalog_key, slug, base_url);
|
||||
let cfg = agent.prepare_sampling_config_for_model(&entry, None);
|
||||
assert_ne!(
|
||||
cfg.api_key.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"LEAK: {catalog_key} carried the primary Kimi session bearer to {base_url}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// C2, the `[model.*]` repro. A `[model.gpt-4o]` block has `info.id == None`, so
|
||||
/// it has no platform at all — which used to be a blanket allow. BYOK is
|
||||
/// `has_own_credentials()`, which probes `std::env::var` AT CALL TIME, so an
|
||||
/// unset (or mistyped) `env_key` classifies the model NotByok and the Kimi
|
||||
/// bearer went to `api.openai.com` on BOTH channels.
|
||||
///
|
||||
/// Revert-to-red: making the `None` arm of `platform_takes_session_credential`
|
||||
/// return `true` again makes `api_key` here `Some(KIMI_TOKEN)`.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn config_model_with_an_unset_env_key_never_carries_the_kimi_bearer() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let _typo = EnvGuard::unset("OPENAI_API_KEY_TYPO");
|
||||
let (_dir, agent) = kimi_session_agent();
|
||||
|
||||
let mut entry = ModelEntry::fallback("gpt-4o", &EndpointsConfig::default());
|
||||
entry.info.id = None; // a `[model.gpt-4o]` config block
|
||||
entry.info.base_url = "https://api.openai.com/v1".to_string();
|
||||
entry.env_key = Some(EnvKeys::single("OPENAI_API_KEY_TYPO"));
|
||||
assert!(
|
||||
!entry.has_own_credentials(),
|
||||
"the env var is unset, so this classifies NotByok — the precondition of the defect"
|
||||
);
|
||||
|
||||
let cfg = agent.prepare_sampling_config_for_model(&entry, None);
|
||||
assert_ne!(
|
||||
cfg.api_key.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"LEAK: a [model.*] block with an unset env_key sent the Kimi bearer to api.openai.com"
|
||||
);
|
||||
assert_eq!(cfg.api_key, None, "no credential resolves — fail fast");
|
||||
}
|
||||
|
||||
/// The first-party subscription channel must stay BYTE-IDENTICAL: `kimi-code/*`
|
||||
/// (and a `[model.*]` block on the session's own coding endpoint, including a
|
||||
/// `KIGI_CODE_BASE_URL` deployment / a local dev proxy) still carries the
|
||||
/// primary session key. This assertion is also what proves the Kimi token is
|
||||
/// live in the tests above — it WOULD leak if the guard were missing.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn the_first_party_subscription_channel_still_carries_the_session_key() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let (_dir, agent) = kimi_session_agent();
|
||||
|
||||
let kimi = platform_entry(
|
||||
"kimi-code/kimi-for-coding",
|
||||
"kimi-for-coding",
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
);
|
||||
assert_eq!(
|
||||
agent
|
||||
.prepare_sampling_config_for_model(&kimi, None)
|
||||
.api_key
|
||||
.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"the kimi-code subscription channel must be unchanged"
|
||||
);
|
||||
|
||||
for base_url in [
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
"http://127.0.0.1:4141/v1",
|
||||
"http://localhost:8080/v1",
|
||||
] {
|
||||
let mut bare = ModelEntry::fallback("kigi-4.5", &EndpointsConfig::default());
|
||||
bare.info.id = None;
|
||||
bare.info.base_url = base_url.to_string();
|
||||
assert_eq!(
|
||||
agent
|
||||
.prepare_sampling_config_for_model(&bare, None)
|
||||
.api_key
|
||||
.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"{base_url}: a custom deployment / local proxy keeps the session key"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// A subscription-OAuth model draws its `api_key` from ITS OWN pooled manager,
|
||||
/// never the Kimi primary — and never falls back to it when that provider has no
|
||||
/// stored session (the pool home is an empty TempDir under `cfg(test)`).
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn oauth_platform_models_never_carry_the_kimi_bearer_as_api_key() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let (_dir, agent) = kimi_session_agent();
|
||||
|
||||
for (catalog_key, slug, base_url) in [
|
||||
("xai-grok/grok-4-latest", "grok-4-latest", "https://api.x.ai/v1"),
|
||||
(
|
||||
"claude-pro-max/claude-opus-4-8",
|
||||
"claude-opus-4-8",
|
||||
"https://api.anthropic.com/v1",
|
||||
),
|
||||
("github-copilot/gpt-4.1", "gpt-4.1", "https://api.githubcopilot.com"),
|
||||
(
|
||||
"openai-codex/gpt-5.5",
|
||||
"gpt-5.5",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
] {
|
||||
let entry = platform_entry(catalog_key, slug, base_url);
|
||||
let cfg = agent.prepare_sampling_config_for_model(&entry, None);
|
||||
assert_ne!(
|
||||
cfg.api_key.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"LEAK: {catalog_key} carried the primary Kimi session bearer to {base_url}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// H5 at the api_key channel: `resolve_model_id` (the picker's own lookup) must
|
||||
/// hand `prepare_sampling_config_for_model` the entry the user SELECTED, even
|
||||
/// when an API-key platform and its subscription-OAuth twin list the same
|
||||
/// routing slug in `PlatformId::ALL` order. Selecting the OAuth twin by catalog
|
||||
/// key must not resolve the API-key twin — and vice versa.
|
||||
#[tokio::test]
|
||||
#[serial_test::serial]
|
||||
async fn dual_credential_slug_collision_resolves_the_selected_catalog_key() {
|
||||
let _env = without_ambient_byok_env();
|
||||
let (_dir, agent) = kimi_session_agent();
|
||||
|
||||
// API-key platform FIRST, exactly as `PlatformId::ALL` orders them.
|
||||
for key in ["xai/grok-4.5", "xai-grok/grok-4.5"] {
|
||||
agent.models_manager.insert_test_entry(
|
||||
key,
|
||||
platform_entry(key, "grok-4.5", "https://api.x.ai/v1"),
|
||||
);
|
||||
}
|
||||
|
||||
for key in ["xai/grok-4.5", "xai-grok/grok-4.5"] {
|
||||
let resolved = agent
|
||||
.resolve_model_id(&acp::ModelId::new(key))
|
||||
.expect("both twins resolve");
|
||||
assert_eq!(
|
||||
resolved.info().id.as_deref(),
|
||||
Some(key),
|
||||
"selecting {key} must resolve THAT catalog entry, not its slug twin"
|
||||
);
|
||||
}
|
||||
|
||||
// And the bare slug resolves the same entry the picker's `resolve_catalog_key`
|
||||
// does — one direction, one answer (the auth layer used to first-match).
|
||||
let by_slug = agent
|
||||
.resolve_model_id(&acp::ModelId::new("grok-4.5"))
|
||||
.expect("the bare slug resolves");
|
||||
let models = agent.models_manager.models();
|
||||
let picker_key = crate::agent::models::resolve_catalog_key(
|
||||
&models,
|
||||
&acp::ModelId::new("grok-4.5"),
|
||||
)
|
||||
.expect("the picker resolves the bare slug");
|
||||
assert_eq!(
|
||||
by_slug.info().id.as_deref(),
|
||||
Some(picker_key.0.as_ref()),
|
||||
"the auth layer and the picker must resolve the SAME entry for one slug"
|
||||
);
|
||||
assert_eq!(
|
||||
crate::agent::config::find_model_by_id(&models, "grok-4.5")
|
||||
.and_then(|e| e.info().id.as_deref()),
|
||||
Some(picker_key.0.as_ref()),
|
||||
"find_model_by_id must agree with resolve_catalog_key by construction"
|
||||
);
|
||||
}
|
||||
@@ -1003,16 +1003,22 @@ fn resolve_model_override_to_config(
|
||||
} else {
|
||||
acp::ModelId::new(entry.info().model.clone())
|
||||
};
|
||||
// Resolve the child's session token by the OVERRIDE model's OWN platform,
|
||||
// not the parent's primary auth: a grok (oauth-platform) override draws its
|
||||
// pooled grok token or `None` — NEVER the primary Kimi session token (which
|
||||
// `resolve_credentials` would otherwise stamp onto the child's api.x.ai
|
||||
// credentials, leaking it in the logout-mid-session edge). A first-party /
|
||||
// non-oauth override still resolves to the primary (byte-identical).
|
||||
let managed_key = entry.info().id.as_deref().unwrap_or(model_id);
|
||||
let session_key = crate::auth::oauth_registry::session_key_for_model(
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
managed_key,
|
||||
// Resolve the child's session token by the OVERRIDE model's OWN platform
|
||||
// AND endpoint, not the parent's primary auth: a grok (oauth-platform)
|
||||
// override draws its pooled grok token or `None`, and an API-key registry
|
||||
// platform / a third-party `[model.*]` host draws NOTHING — NEVER the
|
||||
// primary Kimi session token, which `resolve_credentials` would otherwise
|
||||
// stamp onto the child's api.x.ai / api.moonshot.cn credentials. The
|
||||
// first-party subscription channel still resolves to the primary
|
||||
// (byte-identical).
|
||||
let session_key = crate::auth::oauth_registry::session_key_for_endpoint(
|
||||
entry
|
||||
.info()
|
||||
.id
|
||||
.as_deref()
|
||||
.and_then(kigi_models::parse_managed_model_key)
|
||||
.map(|(platform, _)| platform),
|
||||
&entry.info().base_url,
|
||||
Some(&ctx.auth_manager),
|
||||
);
|
||||
let has_session_key = session_key.is_some();
|
||||
|
||||
@@ -3131,19 +3131,25 @@ fn fresh_tool_model_rejects_unavailable_exact_key_over_visible_slug_collision()
|
||||
"validation must inspect the unavailable exact-key entry selected by execution"
|
||||
);
|
||||
}
|
||||
/// Validation must inspect the SAME slug-collision entry execution selects.
|
||||
/// Both go through `find_model_by_id`, whose slug scan takes the LAST match —
|
||||
/// aligned with the picker's `resolve_catalog_key` so the auth layer and the
|
||||
/// picker can never resolve different platforms for one slug (the H5 collision).
|
||||
/// So a blocked LAST entry must be rejected even though an available earlier one
|
||||
/// shares the slug.
|
||||
#[test]
|
||||
fn fresh_tool_model_rejects_unavailable_first_slug_collision() {
|
||||
fn fresh_tool_model_rejects_unavailable_last_slug_collision() {
|
||||
let mut models = indexmap::IndexMap::new();
|
||||
let mut unavailable_first = test_model_entry("shared-routing-slug");
|
||||
unavailable_first.info.user_selectable = false;
|
||||
models.insert("blocked-first".to_string(), unavailable_first);
|
||||
models.insert("visible-second".to_string(), test_model_entry("shared-routing-slug"));
|
||||
models.insert("visible-first".to_string(), test_model_entry("shared-routing-slug"));
|
||||
let mut unavailable_last = test_model_entry("shared-routing-slug");
|
||||
unavailable_last.info.user_selectable = false;
|
||||
models.insert("blocked-last".to_string(), unavailable_last);
|
||||
assert_eq!(
|
||||
super::handle_request::task_model_override_error(Some("shared-routing-slug"),
|
||||
ModelOverrideProvenance::Tool, false, & models, false,).as_deref(),
|
||||
Some("Unknown Task.model slug 'shared-routing-slug'. Valid model slugs: \
|
||||
visible-second. Omit `model` to inherit the parent model."),
|
||||
"validation must inspect the first routing-slug entry selected by execution"
|
||||
visible-first. Omit `model` to inherit the parent model."),
|
||||
"validation must inspect the last routing-slug entry selected by execution"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
|
||||
@@ -2511,14 +2511,48 @@ async fn subagent_override_grok_model_never_leaks_kimi_session_token() {
|
||||
"a grok override must never receive the primary Kimi session token",
|
||||
);
|
||||
}
|
||||
/// Byte-identical guard: a non-oauth override with a Kimi primary still resolves
|
||||
/// to the primary session token — passes both before and after the fix (the
|
||||
/// non-oauth path is unchanged).
|
||||
/// Byte-identical guard: an override on the SESSION's own first-party endpoint
|
||||
/// (the kimi-code subscription channel) still resolves to the primary session
|
||||
/// token — the primary path is unchanged.
|
||||
#[tokio::test]
|
||||
async fn subagent_override_non_oauth_model_still_gets_primary_token() {
|
||||
async fn subagent_override_first_party_model_still_gets_primary_token() {
|
||||
let (_kd, manager) = kimi_primary_with_token("kimi-secret");
|
||||
let mut entry = test_model_entry("kimi-for-coding");
|
||||
entry.info.id = Some("kimi-code/kimi-for-coding".to_string());
|
||||
entry.info.base_url = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url.to_string();
|
||||
let mut models = indexmap::IndexMap::new();
|
||||
models.insert("kfc".to_string(), entry);
|
||||
let mut ctx = ctx_with_toggle(HashMap::new());
|
||||
ctx.available_models = models;
|
||||
ctx.auth = Some(crate::auth::KimiAuth {
|
||||
key: "kimi-secret".to_string(),
|
||||
auth_mode: crate::auth::AuthMode::OAuth,
|
||||
..crate::auth::KimiAuth::test_default()
|
||||
});
|
||||
ctx.auth_manager = manager;
|
||||
let (config, _model_id) = resolve_model_override_to_config("kfc", &ctx)
|
||||
.expect("first-party override resolves to a config");
|
||||
assert_eq!(
|
||||
config.api_key.as_deref(),
|
||||
Some("kimi-secret"),
|
||||
"a first-party override must still receive the primary session token",
|
||||
);
|
||||
}
|
||||
/// LEAK guard (C1, subagent-override `api_key` channel): an API-key registry
|
||||
/// platform override must NOT receive the parent's primary Kimi session token —
|
||||
/// `resolve_credentials` would stamp it as the child's `api_key` on
|
||||
/// `api.moonshot.cn`. This test previously asserted the opposite
|
||||
/// (`subagent_override_non_oauth_model_still_gets_primary_token`), which encoded
|
||||
/// the defect.
|
||||
///
|
||||
/// Revert-to-red: dropping the `platform_takes_session_credential` term from
|
||||
/// `oauth_registry::session_key_for_endpoint` makes `api_key` `Some("kimi-secret")`.
|
||||
#[tokio::test]
|
||||
async fn subagent_override_api_key_platform_never_gets_the_primary_token() {
|
||||
let (_kd, manager) = kimi_primary_with_token("kimi-secret");
|
||||
let mut entry = test_model_entry("kimi-k2-0905-preview");
|
||||
entry.info.id = Some("moonshot-cn/kimi-k2".to_string());
|
||||
entry.info.base_url = "https://api.moonshot.cn/v1".to_string();
|
||||
let mut models = indexmap::IndexMap::new();
|
||||
models.insert("k2".to_string(), entry);
|
||||
let mut ctx = ctx_with_toggle(HashMap::new());
|
||||
@@ -2529,12 +2563,12 @@ async fn subagent_override_non_oauth_model_still_gets_primary_token() {
|
||||
..crate::auth::KimiAuth::test_default()
|
||||
});
|
||||
ctx.auth_manager = manager;
|
||||
let (config, _model_id) =
|
||||
resolve_model_override_to_config("k2", &ctx).expect("non-oauth override resolves to a config");
|
||||
assert_eq!(
|
||||
let (config, _model_id) = resolve_model_override_to_config("k2", &ctx)
|
||||
.expect("an API-key-platform override still resolves to a config");
|
||||
assert_ne!(
|
||||
config.api_key.as_deref(),
|
||||
Some("kimi-secret"),
|
||||
"a non-oauth override must still receive the primary session token",
|
||||
"LEAK: an API-key-platform override must never receive the primary Kimi session token",
|
||||
);
|
||||
}
|
||||
/// An unresolvable `AgentDefinition.model` pin (model absent from
|
||||
|
||||
@@ -13,11 +13,10 @@
|
||||
//! generic-oauth scope, each wired with the SAME lifecycle as the primary Kimi
|
||||
//! manager (`configure_refresher()` + `start_proactive_refresh()`) so the
|
||||
//! on-disk token stays fresh and a 401 recovers via the provider's own manager.
|
||||
//! Managers are built ON DEMAND: the first grok turn (or model switch) reads the
|
||||
//! on-disk token via [`global_manager_for`], so a login that lands AFTER a
|
||||
//! session spawned self-heals — there is no frozen per-session snapshot to go
|
||||
//! stale. [`manager_for_model`] routes a managed catalog key to the pool (oauth
|
||||
//! platform) or to the session's primary (everything else).
|
||||
//! Managers are built ON DEMAND from the on-disk token ([`global_manager_for`]),
|
||||
//! so a login landing AFTER a session spawned self-heals — no frozen per-session
|
||||
//! snapshot. [`manager_for_model`] routes a managed catalog key to the pool
|
||||
//! (oauth platform) or to the session's primary (everything else).
|
||||
//!
|
||||
//! SECURITY: access/refresh tokens and resolved bearers are NEVER logged here.
|
||||
|
||||
@@ -39,6 +38,31 @@ fn oauth_manager_pool() -> &'static Mutex<HashMap<&'static str, Arc<AuthManager>
|
||||
POOL.get_or_init(|| Mutex::new(HashMap::new()))
|
||||
}
|
||||
|
||||
/// The kigi home every OAuth-pool call site resolves from. Single definition so
|
||||
/// the pool, the aux/summary token routing and the session's inference manager
|
||||
/// can never read different homes.
|
||||
///
|
||||
/// Production: [`crate::util::kigi_home::kigi_home`]. LIB TESTS: a
|
||||
/// process-lifetime `TempDir`, unconditionally — the pool is process-global and
|
||||
/// every manager it builds starts a never-cancelled proactive-refresh loop, so
|
||||
/// a unit test resolving the real `~/.kigi` would read the developer's stored
|
||||
/// OAuth tokens and, 60 s later, fire REAL refresh requests against them.
|
||||
/// Deliberately not a per-test opt-in that can be forgotten: `kigi_home()` is
|
||||
/// itself a `OnceLock` an earlier test has usually already resolved to the real
|
||||
/// home, so setting `KIGI_SHARE_DIR` in a test cannot pin it after the fact.
|
||||
pub(crate) fn pool_home() -> std::path::PathBuf {
|
||||
#[cfg(test)]
|
||||
{
|
||||
static TEST_HOME: OnceLock<tempfile::TempDir> = OnceLock::new();
|
||||
TEST_HOME
|
||||
.get_or_init(|| tempfile::tempdir().expect("tempdir for the test OAuth pool"))
|
||||
.path()
|
||||
.to_path_buf()
|
||||
}
|
||||
#[cfg(not(test))]
|
||||
crate::util::kigi_home::kigi_home()
|
||||
}
|
||||
|
||||
/// Get-or-create the process-global manager for `oauth`, wiring the same
|
||||
/// refresher + proactive-refresh lifecycle as the primary Kimi manager the
|
||||
/// FIRST time a scope is seen. The manager reads the on-disk token at
|
||||
@@ -93,27 +117,70 @@ pub(crate) fn manager_for_model(
|
||||
primary.cloned()
|
||||
}
|
||||
|
||||
/// The SESSION token (the raw bearer/key string) that governs INFERENCE auth
|
||||
/// for `managed_key`, resolved by the model's OWN platform. Thin wrapper over
|
||||
/// [`manager_for_model`] used by the aux-model and subagent-override wire paths
|
||||
/// so a `{platform}/{model}` key never receives the primary token of a
|
||||
/// DIFFERENT provider.
|
||||
/// The SESSION token (the raw bearer/key string) that may ride an INFERENCE
|
||||
/// request routed to `platform` at `base_url`. Used by the aux-model, summary
|
||||
/// and subagent-override wire paths, where the result is stamped straight into
|
||||
/// [`crate::agent::config::resolve_credentials`] as the request's `api_key`.
|
||||
///
|
||||
/// A generic device-code OAuth platform (xai-grok) draws its token from ITS OWN
|
||||
/// pooled manager; when that provider has no stored session the result is
|
||||
/// `None` — NEVER the primary Kimi key. Every other key routes to `primary` and
|
||||
/// yields the primary's current-or-expired token, byte-identical to reading it
|
||||
/// directly. SECURITY: the resolved token is never logged.
|
||||
pub(crate) fn session_key_for_model(
|
||||
kigi_home: &Path,
|
||||
managed_key: &str,
|
||||
/// - a generic device-code OAuth platform (xai-grok, claude-pro-max,
|
||||
/// github-copilot, openai-codex) draws from ITS OWN pooled manager; when that
|
||||
/// provider has no stored session the result is `None` — never `primary`;
|
||||
/// - `kimi-code`, and a platform-less model whose endpoint IS the session's own
|
||||
/// coding endpoint (incl. a `KIGI_CODE_BASE_URL` deployment or a loopback
|
||||
/// proxy), yield the primary's current-or-expired token — byte-identical to
|
||||
/// reading it directly;
|
||||
/// - every API-key registry platform, and every `[model.*]` block pointed at a
|
||||
/// third-party host, yields `None`. Handing them `primary` put the user's
|
||||
/// Kimi subscription bearer on `api.deepseek.com` / `api.moonshot.cn` / …
|
||||
/// as the request's `api_key`.
|
||||
///
|
||||
/// SECURITY: the resolved token is never logged.
|
||||
pub(crate) fn session_key_for_endpoint(
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
primary: Option<&Arc<AuthManager>>,
|
||||
) -> Option<String> {
|
||||
manager_for_model(kigi_home, managed_key, primary)
|
||||
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
|
||||
return global_manager_for(&pool_home(), oauth)
|
||||
.current_or_expired()
|
||||
.map(|a| a.key);
|
||||
}
|
||||
if !crate::agent::auth_method::platform_takes_session_credential(platform, base_url) {
|
||||
return None;
|
||||
}
|
||||
primary
|
||||
.and_then(|am| am.current_or_expired())
|
||||
.map(|a| a.key)
|
||||
}
|
||||
|
||||
/// [`session_key_for_endpoint`] for the catalog model whose routing slug (or
|
||||
/// catalog key) is `slug`.
|
||||
///
|
||||
/// A slug absent from the catalog keeps the pre-registry behaviour: the aux
|
||||
/// resolver's Tier-2 fallback builds its entry against
|
||||
/// `EndpointsConfig::resolve_inference_base_url` (first-party), so the primary
|
||||
/// still governs.
|
||||
pub(crate) fn session_key_for_catalog_model(
|
||||
models: &indexmap::IndexMap<String, crate::agent::config::ModelEntry>,
|
||||
slug: &str,
|
||||
primary: Option<&Arc<AuthManager>>,
|
||||
) -> Option<String> {
|
||||
let Some(entry) = crate::agent::config::find_model_by_id(models, slug) else {
|
||||
return primary
|
||||
.and_then(|am| am.current_or_expired())
|
||||
.map(|a| a.key);
|
||||
};
|
||||
let info = entry.info();
|
||||
session_key_for_endpoint(
|
||||
info.id
|
||||
.as_deref()
|
||||
.and_then(kigi_models::parse_managed_model_key)
|
||||
.map(|(platform, _)| platform),
|
||||
&info.base_url,
|
||||
primary,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -159,6 +226,20 @@ mod tests {
|
||||
.expect("openai-codex carries an OAuthConfig")
|
||||
}
|
||||
|
||||
/// `session_key_for_endpoint` for a managed catalog key, resolving the
|
||||
/// platform and its base URL from the registry exactly as the catalog entry
|
||||
/// would.
|
||||
fn session_key_for_key(
|
||||
managed_key: &str,
|
||||
primary: Option<&Arc<AuthManager>>,
|
||||
) -> Option<String> {
|
||||
let platform = kigi_models::parse_managed_model_key(managed_key).map(|(p, _)| p);
|
||||
let base_url = platform
|
||||
.map(kigi_models::PlatformId::base_url)
|
||||
.unwrap_or_default();
|
||||
session_key_for_endpoint(platform, &base_url, primary)
|
||||
}
|
||||
|
||||
/// An `openai-codex/<model>` turn resolves to the process-global pooled
|
||||
/// openai-codex manager (its OWN `oauth/openai-codex` scope), NEVER the
|
||||
/// primary Kimi manager — the same leak-safe routing as the other OAuth
|
||||
@@ -187,7 +268,7 @@ mod tests {
|
||||
"openai-codex and claude-pro-max must not share a pooled manager"
|
||||
);
|
||||
assert_ne!(
|
||||
session_key_for_model(home.path(), "openai-codex/gpt-5.5", Some(&kimi)),
|
||||
session_key_for_key("openai-codex/gpt-5.5", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"an openai-codex model must never receive the primary Kimi token"
|
||||
);
|
||||
@@ -218,7 +299,7 @@ mod tests {
|
||||
// Fail-fast: even with a Kimi primary, a copilot turn never yields the
|
||||
// Kimi bearer — it draws from the copilot pool (its own token, or None).
|
||||
assert_ne!(
|
||||
session_key_for_model(home.path(), "github-copilot/gpt-4.1", Some(&kimi)),
|
||||
session_key_for_key("github-copilot/gpt-4.1", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"a github-copilot model must never receive the primary Kimi token"
|
||||
);
|
||||
@@ -256,9 +337,8 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn session_key_for_claude_pro_max_is_never_the_kimi_primary() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
assert_ne!(
|
||||
session_key_for_model(home.path(), "claude-pro-max/claude-opus-4-8", Some(&kimi)),
|
||||
session_key_for_key("claude-pro-max/claude-opus-4-8", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"a claude-pro-max model must never receive the primary Kimi session token"
|
||||
);
|
||||
@@ -342,22 +422,62 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// `session_key_for_model`: a non-oauth / bare key yields the primary Kimi
|
||||
/// token exactly as reading it directly would — byte-identical to the
|
||||
/// pre-fix aux/override wire path (no runtime / pool touched).
|
||||
/// `session_key_for_endpoint`: the endpoints that genuinely ride the
|
||||
/// PRIMARY session — `kimi-code` (the subscription channel) and a
|
||||
/// platform-less model routed at the session's own coding endpoint (a
|
||||
/// `KIGI_CODE_BASE_URL` deployment or a loopback dev proxy) — yield the
|
||||
/// primary token exactly as reading it directly would. No runtime / pool
|
||||
/// touched.
|
||||
#[test]
|
||||
fn session_key_for_non_oauth_is_the_primary_token() {
|
||||
fn session_key_for_the_sessions_own_endpoint_is_the_primary_token() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
for key in ["moonshot-cn/kimi-k2", "kimi-k2-0905-preview"] {
|
||||
assert_eq!(
|
||||
session_key_for_key("kimi-code/kimi-for-coding", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"kimi-code rides the primary session, unchanged"
|
||||
);
|
||||
for url in [
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
"http://127.0.0.1:4000/v1",
|
||||
] {
|
||||
assert_eq!(
|
||||
session_key_for_model(home.path(), key, Some(&kimi)),
|
||||
session_key_for_endpoint(None, url, Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"{key} (non-oauth) must yield the primary token unchanged"
|
||||
"{url}: a platform-less model on the session's own endpoint is unchanged"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// LEAK guard (aux / summary / subagent-override `api_key` channel): an
|
||||
/// API-key registry platform, and a `[model.*]` block pointed at a
|
||||
/// third-party host, must yield NO session token. Handing them the primary
|
||||
/// stamped the user's Kimi subscription bearer onto `api.moonshot.cn` /
|
||||
/// `api.deepseek.com` as the request's `api_key` — the channel the
|
||||
/// `bearer_resolver` guard alone does not close.
|
||||
///
|
||||
/// Revert-to-red: dropping the `platform_takes_session_credential` term
|
||||
/// from `session_key_for_endpoint` returns `Some("kimi-tok")` here.
|
||||
#[test]
|
||||
fn session_key_for_a_third_party_endpoint_is_never_the_primary_token() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
for key in [
|
||||
"moonshot-cn/kimi-k2",
|
||||
"deepseek/deepseek-chat",
|
||||
"openai/gpt-5",
|
||||
] {
|
||||
assert_eq!(
|
||||
session_key_for_key(key, Some(&kimi)),
|
||||
None,
|
||||
"LEAK: {key} is an API-key platform — no session token may ride there"
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
session_key_for_endpoint(None, "https://api.openai.com/v1", Some(&kimi)),
|
||||
None,
|
||||
"LEAK: a [model.*] block on a third-party host gets no session token"
|
||||
);
|
||||
}
|
||||
|
||||
/// LEAK guard (aux-model + subagent-override token routing): a grok key with
|
||||
/// a Kimi primary NEVER yields the primary Kimi token — it draws from the
|
||||
/// pooled xai manager (its own token, or `None`). This is the exact source
|
||||
@@ -365,15 +485,14 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn session_key_for_grok_is_never_the_kimi_primary() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
assert_ne!(
|
||||
session_key_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi)),
|
||||
session_key_for_key("xai-grok/grok-4-latest", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"a grok aux/override model must never receive the primary Kimi session token"
|
||||
);
|
||||
// Even with `None` primary the routing is unchanged: grok → pool, never a panic.
|
||||
assert_ne!(
|
||||
session_key_for_model(home.path(), "xai-grok/grok-4-fast", None),
|
||||
session_key_for_key("xai-grok/grok-4-fast", None),
|
||||
Some("kimi-tok".to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -135,7 +135,7 @@ use prompt_build::*;
|
||||
mod session_mode;
|
||||
use session_mode::*;
|
||||
#[path = "acp_session_impl/sampler_turn.rs"]
|
||||
mod sampler_turn;
|
||||
pub(crate) mod sampler_turn;
|
||||
use sampler_turn::*;
|
||||
#[path = "acp_session_impl/tool_dispatch.rs"]
|
||||
mod tool_dispatch;
|
||||
@@ -1256,6 +1256,17 @@ mod rewind_synthetic_turn_tests;
|
||||
#[cfg(test)]
|
||||
#[path = "acp_session_tests/rewrite_zero_turn_prefix_tests.rs"]
|
||||
mod rewrite_zero_turn_prefix_tests;
|
||||
/// The same guard for the model→platform lookup (the dual-credential slug
|
||||
/// collision) and for the stamped aux/summary configs.
|
||||
#[cfg(test)]
|
||||
#[path = "acp_session_tests/session_bearer_leak_platform_tests.rs"]
|
||||
mod session_bearer_leak_platform_tests;
|
||||
/// LEAK guard: the primary Kimi subscription bearer must never ride a request
|
||||
/// to an API-key registry platform's host, while the subscription-OAuth
|
||||
/// platforms keep a live resolver from their OWN pooled manager.
|
||||
#[cfg(test)]
|
||||
#[path = "acp_session_tests/session_bearer_leak_tests.rs"]
|
||||
mod session_bearer_leak_tests;
|
||||
/// Pins the `SubagentFinished` usage-fold attribution gate.
|
||||
#[cfg(test)]
|
||||
#[path = "acp_session_tests/subagent_usage_fold_tests.rs"]
|
||||
|
||||
@@ -634,14 +634,12 @@ impl SessionActor {
|
||||
&active_session_config,
|
||||
Some(self.max_retries),
|
||||
);
|
||||
// A grok (oauth-platform) image-describe model must not inherit the
|
||||
// session (Kimi) bearer_resolver stamped by `finalize_*`; re-point it at
|
||||
// grok's own manager. No-op for a first-party / non-oauth model.
|
||||
// An image-describe model on another provider must not inherit the
|
||||
// session (Kimi) bearer_resolver stamped by `finalize_*`: re-point it at
|
||||
// an OAuth model's own manager, or clear it for an API-key-platform /
|
||||
// third-party endpoint. No-op for the first-party subscription channel.
|
||||
if aux_resolved {
|
||||
self.repoint_aux_bearer_resolver_for_oauth(
|
||||
&mut sampler_config,
|
||||
&self.image_description_model,
|
||||
);
|
||||
self.repoint_aux_bearer_resolver(&mut sampler_config, &self.image_description_model);
|
||||
}
|
||||
let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| {
|
||||
acp::Error::internal_error().data(format!(
|
||||
|
||||
@@ -36,20 +36,37 @@ struct SessionTokenAuthGate {
|
||||
/// BYOK status still refresh against the first-party cli-chat-proxy hosts without
|
||||
/// risking a session-token leak to a third-party BYOK endpoint.
|
||||
endpoint_is_first_party: bool,
|
||||
/// Whether this model's platform is one whose endpoint accepts a session
|
||||
/// bearer at all (see
|
||||
/// [`crate::agent::auth_method::platform_takes_session_credential`]). False
|
||||
/// for every API-key registry platform, which keeps the primary Kimi bearer
|
||||
/// off `api.deepseek.com` / `api.openai.com` / … .
|
||||
endpoint_takes_session_credential: bool,
|
||||
}
|
||||
impl SessionTokenAuthGate {
|
||||
/// Single place `is_session_based` / `endpoint_is_first_party` are derived,
|
||||
/// so all call sites assemble the gate identically.
|
||||
/// so all call sites assemble the gate identically. `model_platform` is the
|
||||
/// registry platform the model routes to (`None` for a bare / `[model.*]`
|
||||
/// entry) — it MUST be derived from the same lookup
|
||||
/// ([`SessionActor::managed_key_for_model`]) that
|
||||
/// [`SessionActor::auth_manager_for_model`] uses, so the gate's verdict and
|
||||
/// the manager actually wrapped as the bearer resolver can never disagree.
|
||||
fn new(
|
||||
auth_method_id: Option<&acp::AuthMethodId>,
|
||||
model_byok: crate::agent::auth_method::ModelByok,
|
||||
base_url: &str,
|
||||
model_platform: Option<kigi_models::PlatformId>,
|
||||
) -> Self {
|
||||
Self {
|
||||
is_session_based: auth_method_id
|
||||
.is_some_and(crate::agent::auth_method::is_session_based_method),
|
||||
model_byok,
|
||||
endpoint_is_first_party: crate::util::is_first_party_url(base_url),
|
||||
endpoint_takes_session_credential:
|
||||
crate::agent::auth_method::platform_takes_session_credential(
|
||||
model_platform,
|
||||
base_url,
|
||||
),
|
||||
}
|
||||
}
|
||||
fn active(self) -> bool {
|
||||
@@ -57,6 +74,7 @@ impl SessionTokenAuthGate {
|
||||
self.is_session_based,
|
||||
self.model_byok,
|
||||
self.endpoint_is_first_party,
|
||||
self.endpoint_takes_session_credential,
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -127,6 +145,43 @@ fn auth_manager_bearer_resolver(
|
||||
) -> kigi_sampler::SharedBearerResolver {
|
||||
std::sync::Arc::new(AuthManagerBearerResolver(am))
|
||||
}
|
||||
/// The `bearer_resolver` an AUX / summary `SamplerConfig` may carry, given the
|
||||
/// SESSION model's resolver (`stamped`) and the AUX model's own platform +
|
||||
/// endpoint. ONE decision shared by image-describe, the auto-mode classifier
|
||||
/// (via [`SessionActor::repoint_aux_bearer_resolver`]) and
|
||||
/// `MvpAgent::build_summary_client`.
|
||||
///
|
||||
/// [`crate::agent::config::stamp_session_local_sampler_fields`] copies the
|
||||
/// session resolver onto every aux config, and `SamplingClient::post` REPLACES
|
||||
/// the request's auth header from it — so an aux model on a DIFFERENT provider
|
||||
/// would have its own correctly-resolved key overwritten by the session bearer
|
||||
/// on the AUX host (H3/H4). Resolve by the aux model instead:
|
||||
/// - an OAuth platform → a live resolver over ITS OWN pooled manager (so a grok
|
||||
/// / claude-pro-max / copilot / codex aux model keeps mid-session refresh);
|
||||
/// - `kimi-code`, or a platform-less model on the session's own coding endpoint
|
||||
/// → the stamped session resolver, byte-identical;
|
||||
/// - every API-key registry platform, and any `[model.*]` block pointed at a
|
||||
/// third-party host → `None`, so the aux model's own key survives to the wire.
|
||||
///
|
||||
/// SECURITY: no token is logged.
|
||||
pub(crate) fn aux_bearer_resolver(
|
||||
stamped: Option<kigi_sampler::SharedBearerResolver>,
|
||||
platform: Option<kigi_models::PlatformId>,
|
||||
base_url: &str,
|
||||
) -> Option<kigi_sampler::SharedBearerResolver> {
|
||||
if let Some(oauth) = platform.and_then(kigi_models::PlatformId::oauth) {
|
||||
return Some(auth_manager_bearer_resolver(
|
||||
crate::auth::oauth_registry::global_manager_for(
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
oauth,
|
||||
),
|
||||
));
|
||||
}
|
||||
if crate::agent::auth_method::platform_takes_session_credential(platform, base_url) {
|
||||
return stamped;
|
||||
}
|
||||
None
|
||||
}
|
||||
impl SessionActor {
|
||||
pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
|
||||
let mcp_wait_start = std::time::Instant::now();
|
||||
@@ -206,7 +261,12 @@ impl SessionActor {
|
||||
fn auth_gate(&self, model_id: &str, base_url: &str) -> SessionTokenAuthGate {
|
||||
let byok = self.model_auth_facts(model_id).byok;
|
||||
let auth_method = self.auth_method_id.load();
|
||||
SessionTokenAuthGate::new(auth_method.as_deref(), byok, base_url)
|
||||
SessionTokenAuthGate::new(
|
||||
auth_method.as_deref(),
|
||||
byok,
|
||||
base_url,
|
||||
self.model_platform(model_id),
|
||||
)
|
||||
}
|
||||
/// The [`AuthManager`](crate::auth::AuthManager) that governs INFERENCE auth
|
||||
/// for the model whose routing slug is `model` (the sampling config's
|
||||
@@ -238,26 +298,30 @@ impl SessionActor {
|
||||
// and resolves to the primary.
|
||||
let managed_key = self.managed_key_for_model(model);
|
||||
crate::auth::oauth_registry::manager_for_model(
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
managed_key.as_deref().unwrap_or(model),
|
||||
self.auth_manager.as_ref(),
|
||||
)
|
||||
}
|
||||
/// Recover the managed catalog key (`{platform}/{model}`) for a routing slug
|
||||
/// from the live catalog. `None` for a bare / config / unlisted model.
|
||||
///
|
||||
/// H5: the catalog KEY the picker selected
|
||||
/// ([`crate::agent::models::ModelsManager::current_model_id`]) is
|
||||
/// authoritative — `model` is the ambiguous bare slug. See
|
||||
/// [`crate::agent::models::managed_key_for_slug`].
|
||||
fn managed_key_for_model(&self, model: &str) -> Option<String> {
|
||||
let models = self.models_manager.models();
|
||||
crate::agent::config::find_model_by_id(&models, model).and_then(|e| e.info().id.clone())
|
||||
let current = self.models_manager.current_model_id();
|
||||
crate::agent::models::managed_key_for_slug(&models, Some(current.0.as_ref()), model)
|
||||
}
|
||||
/// Whether the aux/session model `model` routes to a generic device-code
|
||||
/// OAuth platform (xai-grok). The gate for re-pointing an aux model's
|
||||
/// bearer_resolver away from the session (Kimi) resolver — a first-party /
|
||||
/// non-oauth model returns `false` and keeps the stamped session resolver.
|
||||
fn model_is_oauth_platform(&self, model: &str) -> bool {
|
||||
let managed_key = self.managed_key_for_model(model);
|
||||
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model))
|
||||
.and_then(|(platform, _)| platform.oauth())
|
||||
.is_some()
|
||||
/// The registry platform the routing slug `model` belongs to, from the SAME
|
||||
/// lookup [`Self::auth_manager_for_model`] routes on. `None` for a bare /
|
||||
/// `[model.*]` / unlisted model.
|
||||
fn model_platform(&self, model: &str) -> Option<kigi_models::PlatformId> {
|
||||
let models = self.models_manager.models();
|
||||
let current = self.models_manager.current_model_id();
|
||||
crate::agent::models::platform_for_slug(&models, Some(current.0.as_ref()), model)
|
||||
}
|
||||
/// Whether `model` routes to the Claude Pro/Max OAuth-Messages platform
|
||||
/// (claude-pro-max) — the gate for the sampler's OAuth Messages adaptation
|
||||
@@ -266,22 +330,18 @@ impl SessionActor {
|
||||
/// which is ChatCompletions) returns `false`, keeping the API-key Anthropic
|
||||
/// / MiniMax Messages requests byte-identical.
|
||||
fn model_is_anthropic_oauth(&self, model: &str) -> bool {
|
||||
let managed_key = self.managed_key_for_model(model);
|
||||
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model)).is_some_and(
|
||||
|(platform, _)| {
|
||||
platform.oauth().is_some()
|
||||
&& platform.wire_api() == kigi_models::PlatformWireApi::Messages
|
||||
},
|
||||
)
|
||||
self.model_platform(model).is_some_and(|platform| {
|
||||
platform.oauth().is_some()
|
||||
&& platform.wire_api() == kigi_models::PlatformWireApi::Messages
|
||||
})
|
||||
}
|
||||
/// Whether `model` routes to the GitHub Copilot ChatCompletions platform
|
||||
/// (github-copilot) — the gate for the sampler's editor-identity headers +
|
||||
/// `X-Initiator`. Every other model returns `false`, keeping the other
|
||||
/// ChatCompletions providers byte-identical.
|
||||
fn model_is_github_copilot(&self, model: &str) -> bool {
|
||||
let managed_key = self.managed_key_for_model(model);
|
||||
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model))
|
||||
.is_some_and(|(platform, _)| platform.sends_copilot_editor_headers())
|
||||
self.model_platform(model)
|
||||
.is_some_and(kigi_models::PlatformId::sends_copilot_editor_headers)
|
||||
}
|
||||
/// Whether `model` routes to the ChatGPT/Codex Responses platform
|
||||
/// (openai-codex) — the gate for the sampler's Codex identity headers
|
||||
@@ -289,28 +349,23 @@ impl SessionActor {
|
||||
/// returns `false`, keeping the API-key `openai` Responses request
|
||||
/// byte-identical.
|
||||
fn model_is_openai_codex(&self, model: &str) -> bool {
|
||||
let managed_key = self.managed_key_for_model(model);
|
||||
kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model))
|
||||
.is_some_and(|(platform, _)| platform.sends_codex_responses_headers())
|
||||
self.model_platform(model)
|
||||
.is_some_and(kigi_models::PlatformId::sends_codex_responses_headers)
|
||||
}
|
||||
/// LEAK guard for the stamped aux paths (auto-mode classifier, image
|
||||
/// describe). After [`crate::agent::config::stamp_session_local_sampler_fields`]
|
||||
/// has copied the SESSION model's `bearer_resolver` onto an aux
|
||||
/// `SamplerConfig`, re-point it at the AUX model's OWN platform manager when
|
||||
/// the aux model is an oauth platform (xai-grok) — so a grok aux model never
|
||||
/// inherits the live Kimi session bearer (→ api.x.ai). No-op for a
|
||||
/// first-party / non-oauth aux model (keeps the stamped session resolver →
|
||||
/// byte-identical). SECURITY: no token is logged.
|
||||
pub(super) fn repoint_aux_bearer_resolver_for_oauth(
|
||||
/// describe). Applies [`aux_bearer_resolver`] to the config
|
||||
/// [`crate::agent::config::stamp_session_local_sampler_fields`] just stamped
|
||||
/// the SESSION model's `bearer_resolver` onto.
|
||||
pub(super) fn repoint_aux_bearer_resolver(
|
||||
&self,
|
||||
cfg: &mut kigi_sampler::SamplerConfig,
|
||||
slug: &str,
|
||||
) {
|
||||
if self.model_is_oauth_platform(slug)
|
||||
&& let Some(manager) = self.auth_manager_for_model(slug)
|
||||
{
|
||||
cfg.bearer_resolver = Some(auth_manager_bearer_resolver(manager));
|
||||
}
|
||||
cfg.bearer_resolver = aux_bearer_resolver(
|
||||
cfg.bearer_resolver.take(),
|
||||
self.model_platform(slug),
|
||||
&cfg.base_url,
|
||||
);
|
||||
}
|
||||
/// Emit a unified-log breadcrumb whenever the session-token refresh gate is
|
||||
/// evaluated with an **`Unknown`** per-model BYOK status on a session-based
|
||||
@@ -330,8 +385,9 @@ impl SessionActor {
|
||||
let ctx = serde_json::json!(
|
||||
{ "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based"
|
||||
: gate.is_session_based, "endpoint_is_first_party" : gate
|
||||
.endpoint_is_first_party, "refresh_active" : refresh_active, "base_url" :
|
||||
base_url, }
|
||||
.endpoint_is_first_party, "endpoint_takes_session_credential" : gate
|
||||
.endpoint_takes_session_credential, "refresh_active" : refresh_active,
|
||||
"base_url" : base_url, }
|
||||
);
|
||||
let sid = Some(self.session_info.id.0.as_ref());
|
||||
if refresh_active {
|
||||
@@ -385,8 +441,12 @@ impl SessionActor {
|
||||
let creds = self.chat_state_handle.get_credentials().await;
|
||||
let model_facts = self.model_auth_facts(cfg.model.as_str());
|
||||
let auth_method = self.auth_method_id.load();
|
||||
let gate =
|
||||
SessionTokenAuthGate::new(auth_method.as_deref(), model_facts.byok, &cfg.base_url);
|
||||
let gate = SessionTokenAuthGate::new(
|
||||
auth_method.as_deref(),
|
||||
model_facts.byok,
|
||||
&cfg.base_url,
|
||||
self.model_platform(cfg.model.as_str()),
|
||||
);
|
||||
let use_bearer_resolver = gate.active();
|
||||
self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url);
|
||||
// Resolve the bearer from the ACTIVE model's OWN manager: a grok model
|
||||
@@ -589,17 +649,18 @@ impl SessionActor {
|
||||
slug: &str,
|
||||
) -> Option<kigi_sampler::SamplerConfig> {
|
||||
let creds = self.chat_state_handle.get_credentials().await;
|
||||
// Resolve the aux token by the aux model's OWN platform: a grok
|
||||
// (oauth-platform) aux model draws its pooled grok token or `None` —
|
||||
// NEVER the primary Kimi session token (which `resolve_credentials`
|
||||
// would otherwise stamp onto an api.x.ai request). A first-party /
|
||||
// non-oauth aux model still gets the primary (byte-identical).
|
||||
let session_key = crate::auth::oauth_registry::session_key_for_model(
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
let models = self.models_manager.models();
|
||||
// Resolve the aux token by the aux model's OWN platform AND endpoint: a
|
||||
// grok (oauth-platform) aux model draws its pooled grok token or `None`,
|
||||
// and an API-key registry platform draws NOTHING — NEVER the primary
|
||||
// Kimi session token (which `resolve_credentials` would otherwise stamp
|
||||
// onto an api.x.ai / api.deepseek.com request). The first-party
|
||||
// subscription channel still gets the primary (byte-identical).
|
||||
let session_key = crate::auth::oauth_registry::session_key_for_catalog_model(
|
||||
&models,
|
||||
slug,
|
||||
self.auth_manager.as_ref(),
|
||||
);
|
||||
let models = self.models_manager.models();
|
||||
let endpoints = self.models_manager.endpoints();
|
||||
crate::agent::config::resolve_aux_model_sampling_config(
|
||||
slug,
|
||||
@@ -625,10 +686,11 @@ impl SessionActor {
|
||||
&active_session_config,
|
||||
Some(self.max_retries),
|
||||
);
|
||||
// LEAK 1b: a grok aux classifier must not inherit the SESSION model's
|
||||
// (Kimi) bearer_resolver stamped above; re-point it at grok's own
|
||||
// manager (its pooled token, or None). No-op for a non-oauth aux.
|
||||
self.repoint_aux_bearer_resolver_for_oauth(&mut cfg, slug);
|
||||
// LEAK 1b: the aux classifier must not inherit the SESSION model's
|
||||
// (Kimi) bearer_resolver stamped above — re-point it at an OAuth aux
|
||||
// model's own manager, or clear it for an API-key-platform / third-party
|
||||
// aux endpoint. No-op for the first-party subscription channel.
|
||||
self.repoint_aux_bearer_resolver(&mut cfg, slug);
|
||||
let model = cfg.model.clone();
|
||||
let client = kigi_sampler::SamplingClient::new(cfg)
|
||||
.map_err(|e| {
|
||||
@@ -774,6 +836,8 @@ impl SessionActor {
|
||||
session_id = % self.session_info.id.0, is_session_based = gate
|
||||
.is_session_based, model_byok = gate.model_byok.as_str(),
|
||||
endpoint_is_first_party = gate.endpoint_is_first_party,
|
||||
endpoint_takes_session_credential = gate
|
||||
.endpoint_takes_session_credential,
|
||||
"auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401",
|
||||
);
|
||||
kigi_log::unified_log::warn(
|
||||
@@ -783,7 +847,9 @@ impl SessionActor {
|
||||
{ "kind" : error.kind.as_str(), "status_code" : error
|
||||
.status_code, "is_session_based" : gate.is_session_based,
|
||||
"model_byok" : gate.model_byok.as_str(),
|
||||
"endpoint_is_first_party" : gate.endpoint_is_first_party, }
|
||||
"endpoint_is_first_party" : gate.endpoint_is_first_party,
|
||||
"endpoint_takes_session_credential" : gate
|
||||
.endpoint_takes_session_credential, }
|
||||
)),
|
||||
);
|
||||
}
|
||||
@@ -1036,6 +1102,15 @@ impl SessionActor {
|
||||
.map(|c| c.model)
|
||||
.unwrap_or_default();
|
||||
let Some(ref key) = current_key else { return };
|
||||
// M7: a registry-platform model's key comes from that platform's
|
||||
// credential resolved into its catalog entry — it is NEVER a
|
||||
// `[model.*]` block. With the session gate now inactive for every
|
||||
// API-key platform, those turns all fell through to here and paid a
|
||||
// `load_effective_config()` disk read PER TURN, then logged a
|
||||
// permanently false "Model not found in config.toml [model.*]" warning.
|
||||
if self.model_platform(¤t_model_id).is_some() {
|
||||
return;
|
||||
}
|
||||
let Some(new_key) = self.reload_api_key_from_config(¤t_model_id) else {
|
||||
return;
|
||||
};
|
||||
@@ -1161,7 +1236,7 @@ mod bearer_resolver_tests {
|
||||
.oauth()
|
||||
.expect("xai-grok carries an OAuthConfig");
|
||||
let grok = crate::auth::oauth_registry::global_manager_for(
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
oauth,
|
||||
);
|
||||
assert_ne!(
|
||||
|
||||
+16
-10
@@ -487,21 +487,27 @@ fn session_token_auth_gate_truth_table() {
|
||||
use crate::agent::auth_method::{ModelByok, session_token_auth_gate as gate};
|
||||
// Non-session methods never refresh, regardless of BYOK status or endpoint.
|
||||
for fp in [false, true] {
|
||||
assert!(!gate(false, ModelByok::NotByok, fp));
|
||||
assert!(!gate(false, ModelByok::Byok, fp));
|
||||
assert!(!gate(false, ModelByok::Unknown, fp));
|
||||
// Session method: a definite classification ignores the endpoint —
|
||||
// NotByok always refreshes (only ever routes to the session endpoint),
|
||||
// a genuine per-model Byok never does.
|
||||
assert!(gate(true, ModelByok::NotByok, fp));
|
||||
assert!(!gate(true, ModelByok::Byok, fp));
|
||||
assert!(!gate(false, ModelByok::NotByok, fp, true));
|
||||
assert!(!gate(false, ModelByok::Byok, fp, true));
|
||||
assert!(!gate(false, ModelByok::Unknown, fp, true));
|
||||
// Session method on an endpoint that DOES take the session credential
|
||||
// (kimi-code, an OAuth platform's own pool, or a bare / [model.*]
|
||||
// model): a definite classification ignores the endpoint — NotByok
|
||||
// refreshes, a genuine per-model Byok never does.
|
||||
assert!(gate(true, ModelByok::NotByok, fp, true));
|
||||
assert!(!gate(true, ModelByok::Byok, fp, true));
|
||||
// …and an API-key registry platform endpoint is refused on every arm,
|
||||
// first-party flag included: that is the leak guard.
|
||||
assert!(!gate(true, ModelByok::NotByok, fp, false));
|
||||
assert!(!gate(true, ModelByok::Byok, fp, false));
|
||||
assert!(!gate(true, ModelByok::Unknown, fp, false));
|
||||
}
|
||||
// Session method + Unknown BYOK: refresh only against a first-party xAI
|
||||
// host, so a transiently-unclassifiable config can't demote a live session
|
||||
// (the stale-token 401 regression) yet the session token never leaks to a
|
||||
// third-party BYOK endpoint. This arm was unconditionally `false` pre-fix.
|
||||
assert!(gate(true, ModelByok::Unknown, true));
|
||||
assert!(!gate(true, ModelByok::Unknown, false));
|
||||
assert!(gate(true, ModelByok::Unknown, true, true));
|
||||
assert!(!gate(true, ModelByok::Unknown, false, true));
|
||||
}
|
||||
|
||||
/// Pre-fix, the gate read `auth_type` and skipped recovery here, 401'ing every
|
||||
|
||||
+289
@@ -0,0 +1,289 @@
|
||||
//! LEAK GUARD, part 2: the model→platform lookup (H5) and the AUX resolver
|
||||
//! decision (H3/H4). Shares the fixtures in
|
||||
//! [`super::session_bearer_leak_tests`]; see that module's header for the chain
|
||||
//! and the storage-discipline contract.
|
||||
|
||||
use super::session_bearer_leak_tests::{
|
||||
KIMI_TOKEN, actor_on_managed_model, actor_with_catalog, managed_entry,
|
||||
};
|
||||
use super::*;
|
||||
use kigi_sampler::BearerResolver;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn dual_credential_slug_collision_resolves_the_selected_oauth_platform() {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
// (api-key twin, oauth twin, shared slug, host)
|
||||
for (api_key_twin, oauth_twin, slug, base_url) in [
|
||||
(
|
||||
"xai/grok-4.5",
|
||||
"xai-grok/grok-4.5",
|
||||
"grok-4.5",
|
||||
"https://api.x.ai/v1",
|
||||
),
|
||||
(
|
||||
"anthropic/claude-opus-4-8",
|
||||
"claude-pro-max/claude-opus-4-8",
|
||||
"claude-opus-4-8",
|
||||
"https://api.anthropic.com/v1",
|
||||
),
|
||||
(
|
||||
"openai/gpt-5.5-codex",
|
||||
"openai-codex/gpt-5.5-codex",
|
||||
"gpt-5.5-codex",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
] {
|
||||
// API-key platform FIRST, exactly as `PlatformId::ALL` orders them.
|
||||
let catalog = vec![
|
||||
managed_entry(api_key_twin, slug, base_url),
|
||||
managed_entry(oauth_twin, slug, base_url),
|
||||
];
|
||||
let (_dir, actor, _rx) = actor_with_catalog(catalog, oauth_twin, "unused").await;
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
|
||||
let resolver = cfg.bearer_resolver.as_ref().unwrap_or_else(|| {
|
||||
panic!(
|
||||
"{oauth_twin}: selecting the OAuth twin must keep a LIVE bearer_resolver \
|
||||
(mid-session refresh); resolving {api_key_twin} instead drops it"
|
||||
)
|
||||
});
|
||||
assert_ne!(
|
||||
resolver.current_bearer(),
|
||||
Some(KIMI_TOKEN.to_string()),
|
||||
"{oauth_twin}: the resolver must read its OWN pool, never the Kimi primary"
|
||||
);
|
||||
|
||||
let platform = kigi_models::parse_managed_model_key(oauth_twin)
|
||||
.expect("managed key")
|
||||
.0;
|
||||
assert_eq!(
|
||||
cfg.anthropic_oauth,
|
||||
platform.wire_api() == kigi_models::PlatformWireApi::Messages,
|
||||
"{oauth_twin}: the Claude OAuth Messages adaptation must follow the \
|
||||
SELECTED platform"
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.openai_codex,
|
||||
platform.sends_codex_responses_headers(),
|
||||
"{oauth_twin}: the Codex identity headers must follow the SELECTED platform"
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.github_copilot,
|
||||
platform.sends_copilot_editor_headers(),
|
||||
"{oauth_twin}: the Copilot editor headers must follow the SELECTED platform"
|
||||
);
|
||||
}
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// The other half of H5: selecting the API-KEY twin of a colliding slug must
|
||||
/// still resolve the API-key platform — no bearer_resolver, no adaptations. The
|
||||
/// unified lookup must not simply prefer OAuth.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn dual_credential_slug_collision_resolves_the_selected_api_key_platform() {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let catalog = vec![
|
||||
managed_entry(
|
||||
"anthropic/claude-opus-4-8",
|
||||
"claude-opus-4-8",
|
||||
"https://api.anthropic.com/v1",
|
||||
),
|
||||
managed_entry(
|
||||
"claude-pro-max/claude-opus-4-8",
|
||||
"claude-opus-4-8",
|
||||
"https://api.anthropic.com/v1",
|
||||
),
|
||||
];
|
||||
let (_dir, actor, _rx) =
|
||||
actor_with_catalog(catalog, "anthropic/claude-opus-4-8", "sk-ant-byok").await;
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
assert!(
|
||||
cfg.bearer_resolver.is_none(),
|
||||
"selecting the API-key twin must get NO session bearer resolver"
|
||||
);
|
||||
assert!(
|
||||
!cfg.anthropic_oauth,
|
||||
"the API-key Anthropic Messages request must stay byte-identical"
|
||||
);
|
||||
assert_eq!(cfg.api_key.as_deref(), Some("sk-ant-byok"));
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// MANDATORY counterpart: the subscription-OAuth platforms have NON-first-party
|
||||
/// base URLs, so the fix must not disable their resolver. Each must still get a
|
||||
/// LIVE `bearer_resolver` — and it must read THAT platform's own pooled
|
||||
/// `AuthManager`, never the Kimi primary. The pooled managers are empty here (a
|
||||
/// TempDir pool home), which is what makes `current_bearer() == None` a proof
|
||||
/// that the Kimi bearer cannot be what they resolve.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn oauth_platform_models_keep_a_live_resolver_from_their_own_pool() {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
for (catalog_key, slug, base_url) in [
|
||||
(
|
||||
"claude-pro-max/claude-opus-4-8",
|
||||
"claude-opus-4-8",
|
||||
"https://api.anthropic.com/v1",
|
||||
),
|
||||
(
|
||||
"github-copilot/gpt-4.1",
|
||||
"gpt-4.1",
|
||||
"https://api.githubcopilot.com",
|
||||
),
|
||||
(
|
||||
"xai-grok/grok-4-latest",
|
||||
"grok-4-latest",
|
||||
"https://api.x.ai/v1",
|
||||
),
|
||||
(
|
||||
"openai-codex/gpt-5.5",
|
||||
"gpt-5.5",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
),
|
||||
] {
|
||||
let (_dir, actor, _rx) =
|
||||
actor_on_managed_model(catalog_key, slug, base_url, "unused").await;
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
let resolver = cfg.bearer_resolver.as_ref().unwrap_or_else(|| {
|
||||
panic!("{catalog_key}: must keep a live bearer_resolver for refresh")
|
||||
});
|
||||
assert_ne!(
|
||||
resolver.current_bearer(),
|
||||
Some(KIMI_TOKEN.to_string()),
|
||||
"{catalog_key}: the Kimi bearer must never be what it resolves"
|
||||
);
|
||||
|
||||
// The resolver is LIVE over that platform's pooled manager: a
|
||||
// token rotated inside the pool is observed by the
|
||||
// already-built resolver (this is what mid-session refresh
|
||||
// does). The pool is read here, never mutated.
|
||||
let pooled = crate::auth::oauth_registry::manager_for_model(
|
||||
&crate::auth::oauth_registry::pool_home(),
|
||||
catalog_key,
|
||||
actor.auth_manager.as_ref(),
|
||||
)
|
||||
.expect("an OAuth platform always resolves a manager");
|
||||
assert!(
|
||||
!Arc::ptr_eq(
|
||||
&pooled,
|
||||
actor.auth_manager.as_ref().expect("primary is present")
|
||||
),
|
||||
"{catalog_key}: must route to its OWN pooled manager, not the Kimi primary"
|
||||
);
|
||||
assert_eq!(
|
||||
resolver.current_bearer(),
|
||||
pooled.current_or_expired().map(|a| a.key),
|
||||
"{catalog_key}: the resolver must read THIS platform's pooled manager"
|
||||
);
|
||||
}
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// H3/H4 — the stamped AUX paths (image-describe, the auto-mode classifier and
|
||||
/// the session-summary client all funnel through `aux_bearer_resolver`).
|
||||
/// `stamp_session_local_sampler_fields` copies the SESSION model's resolver onto
|
||||
/// every aux config and `SamplingClient::post` REPLACES the request's auth
|
||||
/// header from it, so an API-key-platform aux model would have its own key
|
||||
/// overwritten by the Kimi bearer ON THE AUX HOST.
|
||||
///
|
||||
/// Revert-to-red: returning `stamped` unconditionally (the pre-fix
|
||||
/// "re-point only when the aux model is OAuth" shape) makes the deepseek /
|
||||
/// openai / `[model.*]`-on-openai.com rows resolve `KIMI_TOKEN`.
|
||||
#[test]
|
||||
fn aux_bearer_resolver_clears_the_session_resolver_off_a_third_party_aux_host() {
|
||||
#[derive(Debug)]
|
||||
struct Fixed(&'static str);
|
||||
impl BearerResolver for Fixed {
|
||||
fn current_bearer(&self) -> Option<String> {
|
||||
Some(self.0.to_string())
|
||||
}
|
||||
}
|
||||
let stamped: kigi_sampler::SharedBearerResolver = Arc::new(Fixed(KIMI_TOKEN));
|
||||
let platform =
|
||||
|key: &str| kigi_models::parse_managed_model_key(key).map(|(platform, _)| platform);
|
||||
|
||||
// Cleared: every API-key registry platform, and a `[model.*]` aux model
|
||||
// pointed at a third-party host.
|
||||
for (key, base_url) in [
|
||||
("deepseek/deepseek-chat", "https://api.deepseek.com/v1"),
|
||||
("openai/gpt-5-mini", "https://api.openai.com/v1"),
|
||||
(
|
||||
"moonshot-cn/kimi-k2-turbo-preview",
|
||||
"https://api.moonshot.cn/v1",
|
||||
),
|
||||
] {
|
||||
assert!(
|
||||
crate::session::acp_session::sampler_turn::aux_bearer_resolver(
|
||||
Some(stamped.clone()),
|
||||
platform(key),
|
||||
base_url,
|
||||
)
|
||||
.is_none(),
|
||||
"LEAK: an aux model on {base_url} must not inherit the session bearer resolver"
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
crate::session::acp_session::sampler_turn::aux_bearer_resolver(
|
||||
Some(stamped.clone()),
|
||||
None,
|
||||
"https://api.openai.com/v1",
|
||||
)
|
||||
.is_none(),
|
||||
"LEAK: a [model.*] aux model on a third-party host must not inherit it either"
|
||||
);
|
||||
|
||||
// Kept (byte-identical): the first-party subscription channel and a
|
||||
// platform-less aux model on the session's own endpoint.
|
||||
for (key, base_url) in [
|
||||
(
|
||||
Some("kimi-code/kimi-for-coding"),
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
),
|
||||
(None, kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url),
|
||||
(None, "http://127.0.0.1:4141/v1"),
|
||||
] {
|
||||
let resolved = crate::session::acp_session::sampler_turn::aux_bearer_resolver(
|
||||
Some(stamped.clone()),
|
||||
key.and_then(platform),
|
||||
base_url,
|
||||
)
|
||||
.unwrap_or_else(|| panic!("{key:?} @ {base_url} must keep the session resolver"));
|
||||
assert_eq!(resolved.current_bearer(), Some(KIMI_TOKEN.to_string()));
|
||||
}
|
||||
|
||||
// Re-pointed: an OAuth aux model gets a LIVE resolver over its OWN pool
|
||||
// (empty here), never the stamped Kimi one.
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("runtime for the pooled manager's refresh task");
|
||||
rt.block_on(async {
|
||||
for key in [
|
||||
"xai-grok/grok-4-latest",
|
||||
"claude-pro-max/claude-opus-4-8",
|
||||
"github-copilot/gpt-4.1",
|
||||
"openai-codex/gpt-5.5",
|
||||
] {
|
||||
let resolved = crate::session::acp_session::sampler_turn::aux_bearer_resolver(
|
||||
Some(stamped.clone()),
|
||||
platform(key),
|
||||
"https://example.invalid/v1",
|
||||
)
|
||||
.unwrap_or_else(|| panic!("{key} must keep a live resolver from its own pool"));
|
||||
assert_ne!(
|
||||
resolved.current_bearer(),
|
||||
Some(KIMI_TOKEN.to_string()),
|
||||
"{key}: the aux resolver must never resolve the Kimi session bearer"
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,436 @@
|
||||
//! LEAK GUARD (bearer_resolver channel): the primary (Kimi) subscription bearer
|
||||
//! must never be stamped on a request to a host that does not own it.
|
||||
//!
|
||||
//! Chain the guard closes: a session-based ACP method (`cached_token` /
|
||||
//! `kimi-code` / any OAuth platform) + a selected API-key-platform model
|
||||
//! classifies `ModelByok::NotByok` (the model carries no `[model.*]` key), the
|
||||
//! pre-fix `session_token_auth_gate` returned `true` unconditionally on that
|
||||
//! arm, `auth_manager_for_model` fell through to the primary Kimi manager for a
|
||||
//! non-OAuth platform, and `SamplingClient::post` then REPLACED the correctly
|
||||
//! resolved provider key with the Kimi bearer on the wire.
|
||||
//!
|
||||
//! The `api_key` half of the same defect (the config never even gets the
|
||||
//! provider key, because `resolve_credentials` stamps the session token) is
|
||||
//! pinned in `agent/mvp_agent/tests/api_key_channel_leak_tests.rs`, which drives
|
||||
//! the real `prepare_sampling_config_for_model` resolution path. These tests
|
||||
//! deliberately do NOT hand-stamp a provider key except where the assertion is
|
||||
//! about the resolver overwriting one that already resolved correctly.
|
||||
//!
|
||||
//! The counterpart contract these tests also pin: the four subscription-OAuth
|
||||
//! platforms have non-first-party base URLs but MUST keep a live
|
||||
//! `bearer_resolver` drawn from their OWN pooled `AuthManager`, or they lose
|
||||
//! mid-session token refresh.
|
||||
//!
|
||||
//! STORAGE DISCIPLINE (H6): nothing here touches the developer's real `~/.kigi`
|
||||
//! and nothing hot-swaps the process-global OAuth pool. Under `cfg(test)`
|
||||
//! `oauth_registry::pool_home()` is a process-lifetime `TempDir`, so every
|
||||
//! pooled manager is empty — which is exactly what the assertions need (a live
|
||||
//! resolver that is provably NOT the Kimi one).
|
||||
|
||||
use super::support::*;
|
||||
use super::*;
|
||||
use crate::agent::auth_method::ModelByok;
|
||||
use crate::agent::config::{ModelAuthFacts, ModelEntry, ModelInfo};
|
||||
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
|
||||
use kigi_sampler::BearerResolver;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// The primary session bearer. Any occurrence of this string in an outgoing
|
||||
/// request to a third-party host is the defect.
|
||||
pub(super) const KIMI_TOKEN: &str = "kimi-subscription-token-DO-NOT-LEAK";
|
||||
|
||||
/// `(tempdir, manager)` standing in for the session's primary Kimi
|
||||
/// `AuthManager`, holding a live (unexpired) OAuth session bearer.
|
||||
fn kimi_primary() -> (tempfile::TempDir, Arc<AuthManager>) {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
|
||||
am.hot_swap(KimiAuth {
|
||||
key: KIMI_TOKEN.to_string(),
|
||||
auth_mode: AuthMode::OAuth,
|
||||
refresh_token: Some("rt".into()),
|
||||
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
|
||||
..KimiAuth::test_default()
|
||||
});
|
||||
(dir, am)
|
||||
}
|
||||
|
||||
/// One catalog entry: catalog key `catalog_key`, routing slug `slug`, routed at
|
||||
/// `base_url`, carrying no credential of its own (the shape every fetched
|
||||
/// registry model has).
|
||||
pub(super) fn managed_entry(catalog_key: &str, slug: &str, base_url: &str) -> (String, ModelEntry) {
|
||||
let mut info = ModelInfo::fallback(slug);
|
||||
info.id = Some(catalog_key.to_string());
|
||||
info.base_url = base_url.to_string();
|
||||
(
|
||||
catalog_key.to_string(),
|
||||
ModelEntry {
|
||||
info,
|
||||
api_key: None,
|
||||
env_key: None,
|
||||
api_base_url: None,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// A `SessionActor` on a session-based ACP method with a live Kimi primary,
|
||||
/// whose live catalog holds `catalog` and whose SELECTED model is the catalog
|
||||
/// key `selected` (the picker's own notion of "current"). `wire_key` is the
|
||||
/// already-correctly-resolved provider credential sitting in chat state.
|
||||
///
|
||||
/// The per-model BYOK memo is pinned to `NotByok` on purpose: that is what a
|
||||
/// fetched registry model actually resolves to (`resolve_model_auth_facts` only
|
||||
/// ever sees `default_models.json` + `[model.*]`), and pinning it keeps the test
|
||||
/// independent of the developer's on-disk `~/.kigi/config.toml`.
|
||||
pub(super) async fn actor_with_catalog(
|
||||
catalog: Vec<(String, ModelEntry)>,
|
||||
selected: &str,
|
||||
wire_key: &str,
|
||||
) -> (
|
||||
tempfile::TempDir,
|
||||
Arc<SessionActor>,
|
||||
mpsc::UnboundedReceiver<PersistenceMsg>,
|
||||
) {
|
||||
let (dir, am) = kimi_primary();
|
||||
let (gateway_tx, _gateway_rx) = mpsc::unbounded_channel();
|
||||
let (persistence_tx, persistence_rx) = mpsc::unbounded_channel();
|
||||
let mut actor = create_test_actor(50_000, 200_000, 85, gateway_tx, persistence_tx).await;
|
||||
actor.auth_manager = Some(am);
|
||||
actor.auth_method_id = test_auth_method_id("cached_token");
|
||||
|
||||
let mut selected_entry = None;
|
||||
for (key, entry) in catalog {
|
||||
if key == selected {
|
||||
selected_entry = Some(entry.clone());
|
||||
}
|
||||
actor.models_manager.insert_test_entry(key, entry);
|
||||
}
|
||||
let selected_entry = selected_entry.expect("the selected key must be in the catalog");
|
||||
actor
|
||||
.models_manager
|
||||
.set_current_model_id(acp::ModelId::new(selected.to_string()));
|
||||
|
||||
let slug = selected_entry.info().model.clone();
|
||||
actor
|
||||
.chat_state_handle
|
||||
.update_sampling_config(kigi_sampling_types::SamplingConfig {
|
||||
base_url: selected_entry.info().base_url.clone(),
|
||||
model: slug.clone(),
|
||||
max_completion_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
chat_compat: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(200_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
stream_tool_calls: None,
|
||||
});
|
||||
actor
|
||||
.chat_state_handle
|
||||
.update_credentials(kigi_chat_state::Credentials {
|
||||
api_key: Some(wire_key.to_string()),
|
||||
auth_type: kigi_chat_state::AuthType::SessionToken,
|
||||
..Default::default()
|
||||
});
|
||||
actor.model_auth_facts.replace(Some((
|
||||
slug,
|
||||
ModelAuthFacts {
|
||||
byok: ModelByok::NotByok,
|
||||
auth_scheme: Default::default(),
|
||||
},
|
||||
)));
|
||||
(dir, Arc::new(actor), persistence_rx)
|
||||
}
|
||||
|
||||
/// Single-entry convenience over [`actor_with_catalog`].
|
||||
pub(super) async fn actor_on_managed_model(
|
||||
catalog_key: &str,
|
||||
slug: &str,
|
||||
base_url: &str,
|
||||
wire_key: &str,
|
||||
) -> (
|
||||
tempfile::TempDir,
|
||||
Arc<SessionActor>,
|
||||
mpsc::UnboundedReceiver<PersistenceMsg>,
|
||||
) {
|
||||
actor_with_catalog(
|
||||
vec![managed_entry(catalog_key, slug, base_url)],
|
||||
catalog_key,
|
||||
wire_key,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// THE leak test, at the wire. A `deepseek/deepseek-chat` turn on a session
|
||||
/// (`cached_token`) method with a live Kimi primary must send DeepSeek's own key
|
||||
/// — the Kimi subscription bearer must not appear anywhere in the request.
|
||||
///
|
||||
/// Revert-to-red: dropping `endpoint_takes_session_credential` from
|
||||
/// `session_token_auth_gate` puts `Bearer <KIMI_TOKEN>` on this request.
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn deepseek_turn_under_a_kimi_session_sends_no_kimi_bearer_on_the_wire() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::path("/chat/completions"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"id": "cmpl-1",
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": "deepseek-chat",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": { "role": "assistant", "content": "ok" },
|
||||
"finish_reason": "stop"
|
||||
}]
|
||||
})),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let uri = server.uri();
|
||||
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let (_dir, actor, _rx) = actor_on_managed_model(
|
||||
"deepseek/deepseek-chat",
|
||||
"deepseek-chat",
|
||||
&uri,
|
||||
"sk-deepseek-provider-key",
|
||||
)
|
||||
.await;
|
||||
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
assert!(
|
||||
cfg.bearer_resolver.is_none(),
|
||||
"an API-key platform model must get NO session bearer resolver"
|
||||
);
|
||||
let client =
|
||||
kigi_sampler::SamplingClient::new(cfg).expect("sampling client must construct");
|
||||
let _ = client
|
||||
.chat_completion(kigi_sampling_types::ChatCompletionRequest::new(
|
||||
"deepseek-chat",
|
||||
vec![kigi_sampling_types::ChatRequestMessage::user("hi")],
|
||||
))
|
||||
.await;
|
||||
})
|
||||
.await;
|
||||
|
||||
let requests = server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("wiremock records requests");
|
||||
assert_eq!(requests.len(), 1, "exactly one inference request was sent");
|
||||
let auth = requests[0]
|
||||
.headers
|
||||
.get("authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.expect("the request must carry an Authorization header")
|
||||
.to_string();
|
||||
assert!(
|
||||
!auth.contains(KIMI_TOKEN),
|
||||
"the Kimi subscription bearer must never reach a third-party inference host"
|
||||
);
|
||||
assert_eq!(
|
||||
auth, "Bearer sk-deepseek-provider-key",
|
||||
"the correctly-resolved provider key must survive to the wire"
|
||||
);
|
||||
}
|
||||
|
||||
/// The same guard for every other API-key registry platform shape: OpenAI
|
||||
/// (Responses), Anthropic (x-api-key/Messages), Groq, Together and Z.AI CN — all
|
||||
/// classify `NotByok`, all route to a non-first-party host, none may receive a
|
||||
/// session bearer resolver.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn api_key_platform_models_get_no_session_bearer_resolver() {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
for (catalog_key, slug, base_url) in [
|
||||
("openai/gpt-5.2", "gpt-5.2", "https://api.openai.com/v1"),
|
||||
(
|
||||
"anthropic/claude-opus-4-8",
|
||||
"claude-opus-4-8",
|
||||
"https://api.anthropic.com/v1",
|
||||
),
|
||||
("groq/llama-4", "llama-4", "https://api.groq.com/openai/v1"),
|
||||
("together/qwen-3", "qwen-3", "https://api.together.xyz/v1"),
|
||||
(
|
||||
"zai-coding-cn/glm-5",
|
||||
"glm-5",
|
||||
"https://open.bigmodel.cn/api/paas/v4",
|
||||
),
|
||||
] {
|
||||
let (_dir, actor, _rx) =
|
||||
actor_on_managed_model(catalog_key, slug, base_url, "sk-provider-key").await;
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
assert!(
|
||||
cfg.bearer_resolver.is_none(),
|
||||
"{catalog_key}: an API-key platform must get no session bearer resolver"
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.api_key.as_deref(),
|
||||
Some("sk-provider-key"),
|
||||
"{catalog_key}: the provider key must stay on the config"
|
||||
);
|
||||
}
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// C2 at the resolver channel: a `[model.*]` entry has NO platform
|
||||
/// (`info.id == None`), which used to be a blanket allow. Pointed at a
|
||||
/// third-party host it must get no session resolver; pointed at the session's
|
||||
/// own coding endpoint (a `KIGI_CODE_BASE_URL` deployment or a local dev proxy)
|
||||
/// it must keep one — that is why the predicate is not `is_first_party_url`.
|
||||
///
|
||||
/// Revert-to-red: making the `None` arm of `platform_takes_session_credential`
|
||||
/// return `true` again puts a Kimi resolver on the openai.com config.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn config_model_entry_takes_a_session_resolver_only_on_its_own_endpoint() {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
for base_url in ["https://api.openai.com/v1", "https://api.deepseek.com/v1"] {
|
||||
let mut info = ModelInfo::fallback("gpt-4o");
|
||||
info.id = None; // a `[model.gpt-4o]` block
|
||||
info.base_url = base_url.to_string();
|
||||
let entry = ModelEntry {
|
||||
info,
|
||||
api_key: None,
|
||||
// An env_key that is NOT set: `has_own_credentials()` probes
|
||||
// `std::env::var` at call time, so this classifies NotByok.
|
||||
env_key: None,
|
||||
api_base_url: None,
|
||||
};
|
||||
let (_dir, actor, _rx) =
|
||||
actor_with_catalog(vec![("gpt-4o".to_string(), entry)], "gpt-4o", "").await;
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
assert!(
|
||||
cfg.bearer_resolver.is_none(),
|
||||
"LEAK: a [model.*] block at {base_url} must get no session bearer resolver"
|
||||
);
|
||||
}
|
||||
|
||||
for base_url in [
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
"http://127.0.0.1:4141/v1",
|
||||
] {
|
||||
let mut info = ModelInfo::fallback("kigi-4.5");
|
||||
info.id = None;
|
||||
info.base_url = base_url.to_string();
|
||||
let entry = ModelEntry {
|
||||
info,
|
||||
api_key: None,
|
||||
env_key: None,
|
||||
api_base_url: None,
|
||||
};
|
||||
let (_dir, actor, _rx) =
|
||||
actor_with_catalog(vec![("kigi-4.5".to_string(), entry)], "kigi-4.5", "").await;
|
||||
let resolver = actor
|
||||
.reconstruct_full_config()
|
||||
.await
|
||||
.bearer_resolver
|
||||
.expect("the session's own endpoint keeps the session resolver");
|
||||
assert_eq!(
|
||||
resolver.current_bearer(),
|
||||
Some(KIMI_TOKEN.to_string()),
|
||||
"{base_url}: a custom deployment / dev proxy is unchanged"
|
||||
);
|
||||
}
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// H5 — the slug collision. A user holding BOTH an xAI API key and a Grok
|
||||
/// subscription has `xai/grok-4.5` AND `xai-grok/grok-4.5` in one catalog, in
|
||||
/// `PlatformId::ALL` order (`Xai`(15) before `XaiGrok`(25)) and with the SAME
|
||||
/// routing slug. `cfg.model` is that bare slug, so the auth layer used to
|
||||
/// first-match the API-key entry: no bearer_resolver, no live refresh (the
|
||||
/// session dies ~1h in with an unrecoverable 401), and — for the Anthropic and
|
||||
/// Codex twins — the OAuth Messages adaptation and the Codex identity headers
|
||||
/// silently dropped.
|
||||
///
|
||||
/// The catalog KEY the picker selected is now authoritative.
|
||||
///
|
||||
/// Revert-to-red: resolving the platform from `find_model_by_id(models, slug)`
|
||||
/// instead of `current_model_id` resolves `xai/grok-4.5` /
|
||||
/// `anthropic/claude-opus-4-8` / `openai/gpt-5.5-codex` and every assertion
|
||||
/// below fails.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn kimi_first_party_model_still_rides_the_primary_session_bearer() {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let (_dir, actor, _rx) = actor_on_managed_model(
|
||||
"kimi-code/kimi-for-coding",
|
||||
"kimi-for-coding",
|
||||
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
|
||||
"stale-buffered-token",
|
||||
)
|
||||
.await;
|
||||
|
||||
let cfg = actor.reconstruct_full_config().await;
|
||||
let resolver = cfg
|
||||
.bearer_resolver
|
||||
.as_ref()
|
||||
.expect("the subscription model must keep the live session resolver");
|
||||
assert_eq!(
|
||||
resolver.current_bearer(),
|
||||
Some(KIMI_TOKEN.to_string()),
|
||||
"the first-party model resolves the primary session bearer"
|
||||
);
|
||||
|
||||
actor.refresh_token_if_expired().await;
|
||||
assert_eq!(
|
||||
actor
|
||||
.chat_state_handle
|
||||
.get_credentials()
|
||||
.await
|
||||
.api_key
|
||||
.as_deref(),
|
||||
Some(KIMI_TOKEN),
|
||||
"the first-party pre-flight refresh must still heal the stale key"
|
||||
);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// The persistence half of the defect: `refresh_token_if_expired` used to write
|
||||
/// the Kimi session token into `chat_state` `creds.api_key` for ANY
|
||||
/// session-method turn, from where it propagated to subagents and aux configs.
|
||||
/// A deepseek turn must leave the provider key untouched.
|
||||
///
|
||||
/// M7 rides along: a registry-platform model must not fall into
|
||||
/// `reload_api_key_from_config` at all (a `load_effective_config()` disk read
|
||||
/// per turn plus a permanently false "not found in config.toml" warning), so
|
||||
/// the key is left exactly as resolved.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn preflight_refresh_never_writes_the_kimi_token_into_a_platform_credential() {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
let (_dir, actor, _rx) = actor_on_managed_model(
|
||||
"deepseek/deepseek-chat",
|
||||
"deepseek-chat",
|
||||
"https://api.deepseek.com/v1",
|
||||
"sk-deepseek-provider-key",
|
||||
)
|
||||
.await;
|
||||
|
||||
actor.refresh_token_if_expired().await;
|
||||
|
||||
assert_eq!(
|
||||
actor
|
||||
.chat_state_handle
|
||||
.get_credentials()
|
||||
.await
|
||||
.api_key
|
||||
.as_deref(),
|
||||
Some("sk-deepseek-provider-key"),
|
||||
"the Kimi session token must never overwrite a platform credential"
|
||||
);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
Reference in New Issue
Block a user