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

feat(providers): add xAI Grok subscription OAuth (device-code) + per-provider session auth

Browse Source 
First subscription-OAuth provider beyond Kimi Code (26th registry variant).
Log in with a Grok/SuperGrok/X subscription via RFC-8628 device-code OAuth
(auth.x.ai), then use it against api.x.ai/v1 — reusing the existing xai wire
(ChatCompletions + OpenAI listing + Passthrough + restrict + models_dev_id
xai). Sourced from Pi (earendil-works/pi auth/oauth/xai.ts): client
b1a00492..., scope 'openid profile email offline_access grok-cli:access
api:access', standard Bearer (no x-xai-token-auth).

Foundation (generalizes Kigi's Kimi-singleton OAuth to per-provider, root
cause, not a patch):
- Registry: OAuthConfig on PlatformSpec (client_id/host/device+token
  paths/scope/scope_key); XAI_OAUTH_CONFIG + XAI_GROK_SPEC (uses_oauth, method
  id 'xai-grok', an interactive login after kimi-code).
- Generic device-code wire (auth/oauth_device.rs) + GenericDeviceRefresher,
  sharing the RFC-8628 core with Kimi; Kimi's bespoke flow is byte-identical
  (X-Msh headers, KIMI_CODE_OAUTH_SCOPE, keyring gating unchanged).
- Per-provider AuthManager via a process-global pool (auth/oauth_registry.rs):
  build-on-demand with start_proactive_refresh, keyed by scope. The session
  resolves the AuthManager for the ACTIVE model's platform for bearer/refresh/
  401-recovery/api_key — an oauth-platform model always uses its OWN token,
  never the primary.
- Live /models under OAuth; base routes oauth().is_some() -> platform.base_url()
  (kimi-code stays on proxy_url).

Security: adversarial review + a systematic token-leak audit found and closed
FIVE channels where the primary Kimi token could reach api.x.ai (bearer
resolver, api_key stamping, aux summary/classifier/image-describe models, and
subagent model-override). Each fix routes through the platform-aware resolver
(the oauth model's pooled token or None, NEVER the primary) and is revert-to-red
verified. No access/refresh token is ever logged.

Registry at 26; picker updated (xai-grok interactive login row); TUI
context-window already auto-updates per model. Full gate green (234 suites,
fmt, clippy -D warnings, deny). GPT/Claude/Grok officially permit third-party
subscription use.
This commit is contained in:
ZacharyZhang-NY
2026-07-22 01:36:29 -04:00
parent 8a26460251
commit dbce6bf305
26 changed files with 2359 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-shell/src/agent/app.rs
+8 -1
View File
@@ -178,7 +178,13 @@ async fn prefetch_models(agent_config: &AgentConfig) -> Option<IndexMap<String,
if auth.is_some() || endpoints.has_custom_endpoint() || platform_keys.any() {
tokio::task::spawn_blocking(move || {
prefetch_models_blocking(&endpoints, auth.as_ref(), fetch_auth, &platform_keys)
prefetch_models_blocking(
&endpoints,
auth.as_ref(),
&Default::default(),
fetch_auth,
&platform_keys,
)
})
.await
.ok()
@@ -621,6 +627,7 @@ pub async fn run_leader(
crate::agent::models::prefetch_models_blocking(
&endpoints_for_prefetch,
auth_for_prefetch.as_ref(),
&Default::default(),
fetch_auth_for_prefetch,
&platform_keys_for_prefetch,
)
crates/codegen/kigi-shell/src/agent/auth_method.rs
+125 -7
View File
@@ -114,7 +114,9 @@ pub struct BuiltAuthMethods {
/// 1. `xai.api_key` (if `has_external_api_key`)
/// 2. `cached_token` (if `has_cached_token`)
/// 3. `kimi-code` (the Kimi Code device login)
/// 4. every API-key registry platform, in `PlatformId::ALL` order
/// 4. every generic device-code OAuth login (`xai-grok`, …), in
/// `PlatformId::ALL` order — interactive logins after `kimi-code`
/// 5. every API-key registry platform, in `PlatformId::ALL` order
/// (`moonshot-cn`, `moonshot-ai`, …), always advertised
///
/// The platform methods are for the INTERACTIVE login picker only: they come
@@ -162,6 +164,15 @@ pub fn build_auth_methods(inputs: AuthMethodsBuildInputs<'_>) -> BuiltAuthMethod
}
methods.push(kimi_code_auth_method(login_label));
// Generic device-code OAuth logins (xai-grok, …) are interactive logins
// too: advertise them right after kimi-code, BEFORE the API-key platforms,
// so they stay out of the `auth_methods.first()` startup-metadata slot yet
// ahead of the api-key picker rows.
for platform in kigi_models::PlatformId::ALL {
if platform.oauth().is_some() {
methods.push(oauth_platform_auth_method(platform));
}
}
for platform in kigi_models::PlatformId::ALL {
if !platform.uses_oauth() {
methods.push(platform_auth_method(platform));
@@ -182,6 +193,9 @@ pub enum AuthMethodKind {
KimiCode,
/// Registry API-key platform login (method id = the platform id).
ApiKeyPlatform(kigi_models::PlatformId),
/// Generic device-code OAuth platform login (method id = the platform id,
/// e.g. `xai-grok`). Interactive, like Kimi Code.
OAuthPlatform(kigi_models::PlatformId),
Unknown,
}
@@ -191,9 +205,12 @@ impl AuthMethodKind {
XAI_API_KEY_METHOD_ID => Self::XaiApiKey,
CACHED_TOKEN_AUTH_METHOD_ID => Self::CachedToken,
KIMI_CODE_METHOD_ID => Self::KimiCode,
other => match platform_for_method_id_str(other) {
Some(platform) => Self::ApiKeyPlatform(platform),
None => Self::Unknown,
other => match kigi_models::PlatformId::parse(other) {
// A generic device-code OAuth platform (xai-grok).
Some(p) if p.oauth().is_some() => Self::OAuthPlatform(p),
// A non-OAuth API-key registry platform.
Some(p) if !p.uses_oauth() => Self::ApiKeyPlatform(p),
_ => Self::Unknown,
},
}
}
@@ -206,13 +223,33 @@ impl AuthMethodKind {
}
/// `true` for session-based methods (cached_token, interactive login).
///
/// `OAuthPlatform` (xai-grok) qualifies: it mints a refreshable device-code
/// session, so the per-turn refresh / 401-recovery gate must be ACTIVE for
/// it. This is correct ONLY because the session routes a model's
/// bearer/refresh/recovery to that model's OWN scope-keyed `AuthManager`
/// (see `SessionActor::auth_manager_for_model`) — the gate being active
/// wraps the grok manager for a grok turn, never the Kimi one.
pub fn is_session_based(self) -> bool {
matches!(self, Self::CachedToken | Self::KimiCode)
matches!(
self,
Self::CachedToken | Self::KimiCode | Self::OAuthPlatform(_)
)
}
/// Requires user interaction (device-code login in the browser).
/// Requires user interaction (device-code login in the browser). Both the
/// Kimi Code login and every generic device-code OAuth platform qualify.
pub fn needs_interactive_login(self) -> bool {
matches!(self, Self::KimiCode)
matches!(self, Self::KimiCode | Self::OAuthPlatform(_))
}
/// The generic device-code OAuth platform behind this method, if any
/// (drives the `authenticate` dispatch to the generic device flow).
pub fn oauth_platform(self) -> Option<kigi_models::PlatformId> {
match self {
Self::OAuthPlatform(p) => Some(p),
_ => None,
}
}
pub fn auth_error_message(self) -> &'static str {
@@ -337,6 +374,16 @@ pub fn kimi_code_auth_method(label: Option<&str>) -> acp::AuthMethod {
)
}
/// A generic device-code OAuth platform's interactive login method (method id
/// = the platform id, e.g. `xai-grok`; label from the spec's `login_label`).
pub fn oauth_platform_auth_method(platform: kigi_models::PlatformId) -> acp::AuthMethod {
let name = platform.login_label();
acp::AuthMethod::Agent(
acp::AuthMethodAgent::new(acp::AuthMethodId::new(platform.as_str()), name.to_string())
.description(Some(format!("Sign in with {name}"))),
)
}
/// Interactive API-key login method ids equal
/// [`kigi_models::PlatformId::as_str`] (`moonshot-cn` / `moonshot-ai` / …),
/// which is also the `[platforms.<id>]` config-table name and the auth.json
@@ -514,6 +561,73 @@ mod tests {
);
}
/// xai-grok is a generic device-code OAuth login: it classifies as an
/// `OAuthPlatform`, needs an interactive (browser) login, is NOT api-key,
/// is NOT an API-key registry platform, and is advertised right after
/// kimi-code among the interactive logins.
#[test]
fn xai_grok_is_an_interactive_oauth_login() {
let id = acp::AuthMethodId::new("xai-grok");
let kind = AuthMethodKind::from_id(&id);
assert_eq!(
kind,
AuthMethodKind::OAuthPlatform(kigi_models::PlatformId::XaiGrok)
);
assert!(
kind.needs_interactive_login(),
"xai-grok needs a browser login"
);
assert!(!kind.is_api_key(), "xai-grok is not an api-key method");
assert_eq!(
kind.oauth_platform(),
Some(kigi_models::PlatformId::XaiGrok)
);
// Never an API-key picker target (keeps it out of the paste-box path).
assert_eq!(platform_for_method_id(&id), None);
// Placement: immediately after kimi-code, before the api-key rows.
let built = build_auth_methods(default_inputs());
let ids = method_ids(&built);
let kimi_pos = ids.iter().position(|m| *m == KIMI_CODE_METHOD_ID).unwrap();
assert_eq!(
ids[kimi_pos + 1],
"xai-grok",
"xai-grok must be the interactive login right after kimi-code"
);
assert_eq!(
ids[kimi_pos + 2],
MOONSHOT_CN_METHOD_ID,
"the api-key rows follow the generic oauth logins"
);
}
/// A generic device-code OAuth platform (xai-grok) is SESSION-BASED: its
/// device-code session is refreshable, so the per-turn refresh / 401 gate
/// must be active for it (routing then sends the model's OWN manager). It
/// stays an interactive login and is NOT api-key-shaped.
#[test]
fn oauth_platform_is_session_based_and_refreshable() {
let id = acp::AuthMethodId::new("xai-grok");
let kind = AuthMethodKind::from_id(&id);
assert_eq!(
kind,
AuthMethodKind::OAuthPlatform(kigi_models::PlatformId::XaiGrok)
);
assert!(
kind.is_session_based(),
"OAuthPlatform must be session-based so the refresh/401 gate is active"
);
assert!(
is_session_based_method(&id),
"is_session_based_method(xai-grok) must be true"
);
// Session-based, yet still an interactive browser login and never
// api-key-shaped.
assert!(kind.needs_interactive_login());
assert!(!kind.is_api_key());
// The session-expired copy (not the api-key copy) is the right error.
assert_eq!(kind.auth_error_message(), AUTH_ERROR_SESSION_EXPIRED);
}
/// The OAuth platform id must never resolve as an API-key platform
/// method — `platform_for_method_id`'s `uses_oauth` filter is what keeps
/// the generic `authenticate` arm from hijacking the device login.
@@ -606,6 +720,7 @@ mod tests {
vec![
XAI_API_KEY_METHOD_ID,
KIMI_CODE_METHOD_ID,
"xai-grok",
MOONSHOT_CN_METHOD_ID,
MOONSHOT_AI_METHOD_ID,
"openai",
@@ -654,6 +769,7 @@ mod tests {
XAI_API_KEY_METHOD_ID,
CACHED_TOKEN_AUTH_METHOD_ID,
KIMI_CODE_METHOD_ID,
"xai-grok",
MOONSHOT_CN_METHOD_ID,
MOONSHOT_AI_METHOD_ID,
"openai",
@@ -695,6 +811,7 @@ mod tests {
vec![
CACHED_TOKEN_AUTH_METHOD_ID,
KIMI_CODE_METHOD_ID,
"xai-grok",
MOONSHOT_CN_METHOD_ID,
MOONSHOT_AI_METHOD_ID,
"openai",
@@ -739,6 +856,7 @@ mod tests {
method_ids(&built),
vec![
KIMI_CODE_METHOD_ID,
"xai-grok",
MOONSHOT_CN_METHOD_ID,
MOONSHOT_AI_METHOD_ID,
"openai",
crates/codegen/kigi-shell/src/agent/config.rs
+74
View File
@@ -4599,6 +4599,80 @@ reasoning_effort = "low"
assert_eq!(resolved.base_url, "https://vendor.example/v1");
assert_eq!(resolved.api_key.as_deref(), Some("vendor-key"));
}
/// A primary Kimi manager holding a fixed in-memory bearer, standing in for
/// an aux caller's session. `TempDir` returned so the caller keeps it alive.
fn kimi_primary(key: &str) -> (tempfile::TempDir, std::sync::Arc<AuthManager>) {
let dir = tempfile::tempdir().unwrap();
let manager = std::sync::Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
manager.hot_swap(crate::auth::KimiAuth {
key: key.to_string(),
auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
(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).
#[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.
let mut grok = test_model_entry("grok-4-latest", "https://api.x.ai/v1", None, None, None);
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",
Some(&kimi),
);
let grok_cfg = resolve_aux_model_sampling_config(
"grok",
&grok_catalog,
&endpoints,
grok_key.as_deref(),
None,
);
assert_ne!(
grok_cfg.and_then(|c| c.api_key).as_deref(),
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.
let mut k2 = test_model_entry(
"kimi-k2-0905-preview",
"https://vendor/v1",
None,
None,
None,
);
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",
Some(&kimi),
);
let k2_cfg = resolve_aux_model_sampling_config(
"k2",
&k2_catalog,
&endpoints,
k2_key.as_deref(),
None,
)
.expect("non-oauth aux resolves via the primary session token");
assert_eq!(
k2_cfg.api_key.as_deref(),
Some("kimi-tok"),
"a non-oauth aux model must still receive the primary session token",
);
}
#[test]
fn parses_model_api_key() {
let raw_config: toml::Value = toml::from_str(
crates/codegen/kigi-shell/src/agent/models.rs
+69 -8
View File
@@ -321,6 +321,7 @@ impl ModelsManager {
&cfg.endpoints,
fetch_auth,
has_session,
&Default::default(),
&platform_keys,
),
)
@@ -1062,7 +1063,6 @@ impl ModelsManager {
/// Build a `SamplingConfig` from the current model + auth state.
pub fn sampling_config(&self) -> SamplingConfig {
let config = self.inner.cfg.read().clone();
let auth_manager = self.inner.auth_manager.as_ref();
let current_model_id = self.current_model_id();
let all_models = self.models();
let fallback;
@@ -1079,9 +1079,12 @@ impl ModelsManager {
}
};
let session_auth = auth_manager.current_or_expired();
let credentials =
resolve_credentials(current_model, session_auth.as_ref().map(|a| a.key.as_str()));
// Resolve the session bearer PER the current model's platform: a
// generic device-code OAuth platform (xai-grok) resolves from ITS OWN
// scope-keyed store — never the Kimi session. Kimi and every other
// model keep the primary manager's token (path byte-identical).
let session_key = self.session_key_for_catalog_key(current_model_id.0.as_ref());
let credentials = resolve_credentials(current_model, session_key.as_deref());
sampling_config_for_model(
current_model,
@@ -1090,6 +1093,23 @@ impl ModelsManager {
)
}
/// The session bearer for a model catalog key. Generic device-code OAuth
/// platforms (xai-grok) resolve from their own persisted scope; every other
/// key resolves from the primary (Kimi) `AuthManager`. Reads the persisted
/// token (refresh happens in the async catalog/refresh paths); never logs
/// it.
fn session_key_for_catalog_key(&self, catalog_key: &str) -> Option<String> {
if let Some((platform, _)) = kigi_models::parse_managed_model_key(catalog_key)
&& let Some(oauth) = platform.oauth()
{
let kigi_home = crate::util::kigi_home::kigi_home();
return AuthManager::new_oauth_provider(&kigi_home, oauth)
.current_or_expired()
.map(|a| a.key);
}
self.inner.auth_manager.current_or_expired().map(|a| a.key)
}
/// Disk-cache origin key for this manager's current endpoints/auth shape
/// (see [`ModelsCache::origin`]).
fn cache_origin(&self) -> String {
@@ -1100,10 +1120,15 @@ impl ModelsManager {
let fetch_auth = *self.inner.fetch_auth.read();
let has_oauth = self.inner.auth_manager.current_or_expired().is_some();
let platform_keys = PlatformApiKeys::resolve(&platforms);
// The origin key encodes only enabled-platform NAMES + URLs (never
// tokens). Generic-oauth presence is reflected by the post-login
// `on_auth_changed` re-fetch; an empty map here keeps this sync path
// cheap (no per-provider AuthManager construction on the hot path).
crate::agent::models_fetch::models_fetch_origin(
&endpoints,
fetch_auth,
has_oauth,
&Default::default(),
&platform_keys,
)
}
@@ -1155,8 +1180,21 @@ impl ModelsManager {
let fetch_auth = *self.inner.fetch_auth.read();
let platform_keys = PlatformApiKeys::resolve(&cfg.platforms);
let auth = self.inner.auth_manager.auth().await.ok();
let outcome =
fetch_models_async(endpoints.clone(), auth, fetch_auth, platform_keys.clone()).await;
// Resolve each generic device-code OAuth platform's OWN session token
// (refreshed on expiry) from its own scope — independent of the Kimi
// session above.
let oauth_tokens = crate::agent::models_fetch::resolve_generic_oauth_tokens(
&crate::util::kigi_home::kigi_home(),
)
.await;
let outcome = fetch_models_async(
endpoints.clone(),
auth,
oauth_tokens.clone(),
fetch_auth,
platform_keys.clone(),
)
.await;
if outcome.models.is_some() {
return outcome.models;
}
@@ -1173,7 +1211,12 @@ impl ModelsManager {
return None;
}
let auth = self.inner.auth_manager.auth().await.ok();
let retry = fetch_models_async(endpoints, auth, fetch_auth, platform_keys).await;
let oauth_tokens = crate::agent::models_fetch::resolve_generic_oauth_tokens(
&crate::util::kigi_home::kigi_home(),
)
.await;
let retry =
fetch_models_async(endpoints, auth, oauth_tokens, fetch_auth, platform_keys).await;
if retry.oauth_unauthorized {
tracing::warn!("model catalog: still unauthorized after token refresh");
}
@@ -1614,12 +1657,14 @@ impl ModelsFetchOutcome {
pub(crate) fn prefetch_models_blocking(
endpoints: &config::EndpointsConfig,
auth: Option<&KimiAuth>,
oauth_tokens: &crate::agent::models_fetch::OAuthSessionTokens,
fetch_auth: ModelFetchAuth,
platform_keys: &PlatformApiKeys,
) -> Option<IndexMap<String, ModelEntry>> {
prefetch_models_blocking_gated(
endpoints,
auth,
oauth_tokens,
fetch_auth,
platform_keys,
crate::util::config::resolve_remote_fetch_enabled(),
@@ -1632,6 +1677,7 @@ pub(crate) fn prefetch_models_blocking(
fn prefetch_models_blocking_gated(
endpoints: &config::EndpointsConfig,
auth: Option<&KimiAuth>,
oauth_tokens: &crate::agent::models_fetch::OAuthSessionTokens,
fetch_auth: ModelFetchAuth,
platform_keys: &PlatformApiKeys,
remote_fetch_enabled: bool,
@@ -1643,6 +1689,7 @@ fn prefetch_models_blocking_gated(
endpoints,
fetch_auth,
auth.is_some(),
oauth_tokens,
platform_keys,
);
let cache = ModelsCacheManager::new();
@@ -1666,7 +1713,7 @@ fn prefetch_models_blocking_gated(
}
let _timer = crate::instrumentation_timer!("startup.fetch_models_blocking");
match fetch_models_blocking(endpoints, auth, fetch_auth, platform_keys) {
match fetch_models_blocking(endpoints, auth, oauth_tokens, fetch_auth, platform_keys) {
Ok(FetchModelsResult {
models,
etag,
@@ -1826,6 +1873,10 @@ fn spawn_prefetch_thread(env: PrefetchEnv) -> EarlyPrefetchHandle {
let models = prefetch_models_blocking(
&env.endpoints,
env.auth.as_ref(),
// Startup prefetch does not resolve generic-oauth (xai-grok)
// tokens; those platforms join on the first async catalog refresh
// (post-login `on_auth_changed` / periodic `spawn_fetch`).
&Default::default(),
env.model_fetch_auth,
&env.platform_keys,
);
@@ -2215,6 +2266,7 @@ pub(crate) fn validate_selectable(
pub(crate) async fn fetch_models_async(
endpoints: config::EndpointsConfig,
auth: Option<KimiAuth>,
oauth_tokens: crate::agent::models_fetch::OAuthSessionTokens,
fetch_auth: ModelFetchAuth,
platform_keys: PlatformApiKeys,
) -> ModelsFetchOutcome {
@@ -2222,6 +2274,7 @@ pub(crate) async fn fetch_models_async(
prefetch_models_blocking_gated(
&endpoints,
auth.as_ref(),
&oauth_tokens,
fetch_auth,
&platform_keys,
crate::util::config::resolve_remote_fetch_enabled(),
@@ -3950,6 +4003,7 @@ mod tests {
crate::agent::models_fetch::fetch_models_blocking(
&endpoints,
Some(&auth),
&Default::default(),
ModelFetchAuth::Platforms,
&PlatformApiKeys::default(),
)
@@ -4019,6 +4073,7 @@ mod tests {
crate::agent::models_fetch::fetch_models_blocking(
&endpoints,
None,
&Default::default(),
ModelFetchAuth::Platforms,
&keys,
)
@@ -4170,6 +4225,7 @@ mod tests {
let outcome = prefetch_models_blocking_gated(
&endpoints,
Some(&auth),
&Default::default(),
ModelFetchAuth::Platforms,
&keys,
true,
@@ -4185,6 +4241,7 @@ mod tests {
&endpoints,
ModelFetchAuth::Platforms,
true,
&Default::default(),
&keys,
);
let cache = ModelsCacheManager::new();
@@ -4200,6 +4257,7 @@ mod tests {
let outcome = prefetch_models_blocking_gated(
&endpoints,
Some(&auth),
&Default::default(),
ModelFetchAuth::Platforms,
&keys,
true,
@@ -4218,6 +4276,7 @@ mod tests {
let outcome = prefetch_models_blocking_gated(
&endpoints,
Some(&auth),
&Default::default(),
ModelFetchAuth::Platforms,
&keys,
true,
@@ -4272,6 +4331,7 @@ mod tests {
&endpoints,
ModelFetchAuth::Platforms,
true,
&Default::default(),
&PlatformApiKeys::test_keys(Some("sk"), None),
);
let cache = ModelsCacheManager::new();
@@ -4291,6 +4351,7 @@ mod tests {
let outcome = prefetch_models_blocking_gated(
&endpoints,
Some(&auth),
&Default::default(),
ModelFetchAuth::Platforms,
&PlatformApiKeys::default(),
true,
crates/codegen/kigi-shell/src/agent/models_fetch.rs
+232 -39
View File
@@ -28,6 +28,44 @@ pub(crate) const DEFAULT_CONTEXT_WINDOW: u64 = 256_000;
struct ModelsResponse {
data: Vec<serde_json::Value>,
}
/// Session bearer tokens for the GENERIC device-code OAuth platforms
/// (`platform.oauth().is_some()` — xai-grok today), resolved per provider from
/// its own scope-keyed `AuthManager` (refreshed on expiry) before the blocking
/// fetch. kimi-code is NOT here — it still rides the single `auth: &KimiAuth`.
///
/// SECURITY: values are access tokens — never logged, never persisted here.
pub(crate) type OAuthSessionTokens = std::collections::BTreeMap<kigi_models::PlatformId, String>;
/// Resolve session bearers for every generic device-code OAuth platform whose
/// own `AuthManager` holds a usable (refreshed-on-expiry) session. Each
/// platform resolves from ITS OWN scope (`oauth/xai`, …) — independent of the
/// Kimi session. Providers without a stored session are simply absent. Only a
/// non-secret count is logged.
pub(crate) async fn resolve_generic_oauth_tokens(
kigi_home: &std::path::Path,
) -> OAuthSessionTokens {
let mut tokens = OAuthSessionTokens::new();
for platform in kigi_models::PlatformId::ALL {
let Some(oauth) = platform.oauth() else {
continue;
};
let manager = std::sync::Arc::new(crate::auth::AuthManager::new_oauth_provider(
kigi_home, oauth,
));
manager.configure_refresher();
if let Ok(auth) = manager.auth().await {
tokens.insert(platform, auth.key);
}
}
if !tokens.is_empty() {
tracing::info!(
count = tokens.len(),
"resolved generic oauth platform session tokens"
);
}
tokens
}
/// The models-fetch origin key for this endpoints/auth shape. Used as the
/// models disk-cache origin: cached entries embed absolute `base_url`s from
/// the backend(s) that served them, so a catalog fetched against one fetch
@@ -38,12 +76,13 @@ pub(crate) fn models_fetch_origin(
endpoints: &crate::agent::config::EndpointsConfig,
fetch_auth: crate::agent::models::ModelFetchAuth,
has_oauth: bool,
oauth_tokens: &OAuthSessionTokens,
platform_keys: &crate::agent::models::PlatformApiKeys,
) -> String {
match fetch_auth {
crate::agent::models::ModelFetchAuth::CustomEndpoint => endpoints.resolve_models_list_url(),
crate::agent::models::ModelFetchAuth::Platforms => {
let parts: Vec<String> = enabled_platforms(has_oauth, platform_keys)
let parts: Vec<String> = enabled_platforms(has_oauth, oauth_tokens, platform_keys)
.into_iter()
.map(|p| format!("{}={}", p.as_str(), platform_models_url(p, endpoints)))
.collect();
@@ -53,14 +92,23 @@ pub(crate) fn models_fetch_origin(
}
/// The platforms with usable credentials, in registry order (kimi-code first
/// so "default model = first list item" favors the subscription).
///
/// - kimi-code (`uses_oauth`, no `OAuthConfig`) is gated on the single Kimi
/// session (`has_oauth`);
/// - a generic device-code OAuth platform (`oauth().is_some()`, e.g. xai-grok)
/// is gated on ITS OWN resolved session token (`oauth_tokens`);
/// - an API-key platform is gated on a stored key.
fn enabled_platforms(
has_oauth: bool,
oauth_tokens: &OAuthSessionTokens,
platform_keys: &crate::agent::models::PlatformApiKeys,
) -> Vec<kigi_models::PlatformId> {
kigi_models::PlatformId::ALL
.into_iter()
.filter(|p| {
if p.uses_oauth() {
if p.oauth().is_some() {
oauth_tokens.contains_key(p)
} else if p.uses_oauth() {
has_oauth
} else {
platform_keys.key_for(*p).is_some()
@@ -68,20 +116,30 @@ fn enabled_platforms(
})
.collect()
}
/// `{base}/models` for one platform. The subscription platform resolves its
/// base through the endpoints config (`coding_api_base_url` override,
/// else `KIGI_CODE_BASE_URL` / production default via kigi-env); the open
/// platforms use their fixed bases.
/// The inference/listing base for one platform's `/models` fetch:
/// - kimi-code (`uses_oauth`, no `OAuthConfig`) → the Kimi subscription base
/// via the endpoints config (`proxy_url`);
/// - every other platform — API-key AND generic OAuth (xai-grok) — → its own
/// registry `base_url()` (e.g. api.x.ai/v1).
fn platform_fetch_base(
platform: kigi_models::PlatformId,
endpoints: &crate::agent::config::EndpointsConfig,
) -> String {
if platform.uses_oauth() && platform.oauth().is_none() {
endpoints.proxy_url()
} else {
platform.base_url()
}
}
/// `{base}/models` for one platform (see [`platform_fetch_base`]).
fn platform_models_url(
platform: kigi_models::PlatformId,
endpoints: &crate::agent::config::EndpointsConfig,
) -> String {
let base = if platform.uses_oauth() {
endpoints.proxy_url()
} else {
platform.base_url()
};
format!("{}/models", base.trim_end_matches('/'))
format!(
"{}/models",
platform_fetch_base(platform, endpoints).trim_end_matches('/')
)
}
/// Fetch result: model entries + optional etag from the subscription platform.
pub struct FetchModelsResult {
@@ -106,6 +164,7 @@ pub struct FetchModelsResult {
pub(crate) fn fetch_models_blocking(
endpoints: &crate::agent::config::EndpointsConfig,
auth: Option<&KimiAuth>,
oauth_tokens: &OAuthSessionTokens,
fetch_auth: crate::agent::models::ModelFetchAuth,
platform_keys: &crate::agent::models::PlatformApiKeys,
) -> Result<FetchModelsResult, BackendError> {
@@ -114,7 +173,7 @@ pub(crate) fn fetch_models_blocking(
fetch_custom_endpoint_models_blocking(endpoints, auth)
}
crate::agent::models::ModelFetchAuth::Platforms => {
fetch_platform_models_blocking(endpoints, auth, platform_keys)
fetch_platform_models_blocking(endpoints, auth, oauth_tokens, platform_keys)
}
}
}
@@ -173,9 +232,10 @@ fn fetch_custom_endpoint_models_blocking(
fn fetch_platform_models_blocking(
endpoints: &crate::agent::config::EndpointsConfig,
auth: Option<&KimiAuth>,
oauth_tokens: &OAuthSessionTokens,
platform_keys: &crate::agent::models::PlatformApiKeys,
) -> Result<FetchModelsResult, BackendError> {
let enabled = enabled_platforms(auth.is_some(), platform_keys);
let enabled = enabled_platforms(auth.is_some(), oauth_tokens, platform_keys);
if enabled.is_empty() {
return Err(BackendError::Auth(
"No platform credentials: log in with `kigi login`, paste a platform API key in \
@@ -194,7 +254,14 @@ fn fetch_platform_models_blocking(
// serves its own metadata (kimi/moonshot today).
let enrichment = crate::agent::enrichment_fetch::load_enrichment_catalog(&enabled);
for platform in &enabled {
let bearer = if platform.uses_oauth() {
let bearer = if platform.oauth().is_some() {
// Generic device-code OAuth platform (xai-grok): its OWN resolved
// session token — never the Kimi session.
oauth_tokens
.get(platform)
.expect("enabled_platforms gated on generic-oauth token presence")
.clone()
} else if platform.uses_oauth() {
auth.map(|a| a.key.clone())
.expect("enabled_platforms gated on auth presence")
} else {
@@ -211,7 +278,9 @@ fn fetch_platform_models_blocking(
"platform models fetch succeeded"
);
successes += 1;
if platform.uses_oauth() {
// The catalog etag tracks the Kimi subscription listing only
// (kimi-code: uses_oauth with no generic OAuthConfig).
if platform.uses_oauth() && platform.oauth().is_none() {
etag = platform_etag;
}
models.extend(platform_models);
@@ -383,11 +452,7 @@ fn fetch_one_platform_models(
);
}
}
let base_url = if platform.uses_oauth() {
endpoints.proxy_url()
} else {
platform.base_url()
};
let base_url = platform_fetch_base(platform, endpoints);
let models = filtered
.into_iter()
.map(|mut wire| {
@@ -835,7 +900,7 @@ mod tests {
"sk-oai",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -983,7 +1048,7 @@ mod tests {
"sk-ant",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -1095,7 +1160,7 @@ mod tests {
"sk-ds",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -1193,7 +1258,7 @@ mod tests {
"gsk-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -1277,7 +1342,7 @@ mod tests {
"msk-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -1363,7 +1428,7 @@ mod tests {
"fw-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -1471,7 +1536,7 @@ mod tests {
"gk-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -1555,7 +1620,7 @@ mod tests {
"or-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -1660,7 +1725,7 @@ mod tests {
"tg-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -1753,7 +1818,7 @@ mod tests {
"cb-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -1861,7 +1926,7 @@ mod tests {
"nvapi-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -1955,7 +2020,7 @@ mod tests {
"vg-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -2048,7 +2113,7 @@ mod tests {
"xai-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -2139,7 +2204,7 @@ mod tests {
"qtp-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -2221,7 +2286,7 @@ mod tests {
"kc-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -2312,7 +2377,7 @@ mod tests {
"zai-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -2402,7 +2467,7 @@ mod tests {
"xm-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -2493,7 +2558,7 @@ mod tests {
"mm-1",
);
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys)
fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
})
.await
.unwrap()
@@ -2991,6 +3056,125 @@ mod tests {
"https://registry.acme.com/api/list-models"
);
}
/// xai-grok OAuth-cycle e2e (mock wire): the generic device-code OAuth
/// bearer (`grok-oauth-tok`, resolved from a stored `oauth/xai` session —
/// mocked here as the token map) fetches `GET {base}/models` against the
/// platform's OWN base (api.x.ai/v1 via its env), enriches from models.dev
/// "xai", restricts to tool-calling chat models, and keys each entry under
/// `xai-grok/<id>` with the Passthrough dialect — NOT the API-key `xai/`.
#[tokio::test(flavor = "multi_thread")]
#[serial_test::serial]
async fn xai_grok_oauth_listing_is_enriched_restricted_and_keyed() {
let platform_server = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("GET"))
.and(wiremock::matchers::path("/models"))
.and(wiremock::matchers::header(
"Authorization",
"Bearer grok-oauth-tok",
))
.respond_with(wiremock::ResponseTemplate::new(200).set_body_json(
serde_json::json!({ "data": [
{ "id": "grok-4.5", "object": "model" },
// enrichment-known but NOT tool-calling → dropped by restrict.
{ "id": "grok-2-image", "object": "model" }
]}),
))
.expect(1)
.mount(&platform_server)
.await;
let modelsdev_server = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("GET"))
.and(wiremock::matchers::path("/api.json"))
.respond_with(wiremock::ResponseTemplate::new(200).set_body_json(
serde_json::json!({ "xai": { "models": {
"grok-4.5": {
"name": "Grok 4.5",
"reasoning": true,
"limit": {"context": 256000, "output": 64000},
"tool_call": true
},
"grok-2-image": { "limit": {"context": 8192} }
}}}),
))
.expect(1)
.mount(&modelsdev_server)
.await;
let cache_dir = tempfile::tempdir().unwrap();
let _base = kigi_test_support::EnvGuard::set(
kigi_models::XAI_GROK_BASE_URL_ENV,
platform_server.uri(),
);
let _mdev = kigi_test_support::EnvGuard::set(
crate::agent::enrichment_fetch::MODELS_DEV_URL_ENV,
format!("{}/api.json", modelsdev_server.uri()),
);
let _mdev_cache = kigi_test_support::EnvGuard::set(
crate::agent::enrichment_fetch::MODELS_DEV_CACHE_DIR_ENV,
cache_dir.path(),
);
let endpoints = crate::agent::config::EndpointsConfig::default();
let keys = crate::agent::models::PlatformApiKeys::default();
let mut tokens = OAuthSessionTokens::new();
tokens.insert(kigi_models::PlatformId::XaiGrok, "grok-oauth-tok".into());
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &tokens, &keys)
})
.await
.unwrap()
.expect("fetch must succeed");
assert_eq!(
result
.models
.iter()
.map(|m| m.id.as_deref().unwrap_or_default())
.collect::<Vec<_>>(),
vec!["xai-grok/grok-4.5"],
"grok-2-image (enrichment-known, not tool-calling) must be dropped; \
the surviving model is keyed under xai-grok/ (not xai/)"
);
let entry = &result.models[0];
assert_eq!(entry.model, "grok-4.5");
assert_eq!(
entry.base_url,
platform_server.uri(),
"xai-grok fetches against its OWN base (its base env), not proxy_url"
);
assert_eq!(
entry.context_window.get(),
256_000,
"context window must come from models.dev \"xai\" enrichment"
);
assert_eq!(entry.max_completion_tokens, Some(64_000));
assert_eq!(
entry.api_backend,
crate::sampling::ApiBackend::ChatCompletions
);
assert_eq!(entry.name.as_deref(), Some("Grok 4.5"));
assert!(
entry.env_key.is_none(),
"an OAuth channel carries no api-key env"
);
assert!(
!entry.supported_in_api,
"subscription models require the OAuth session (not the public API)"
);
// Passthrough dialect (identical to the API-key xai wire).
let model_entry = crate::agent::config::ModelEntry::from_config_entry(entry);
let creds = crate::agent::config::ResolvedCredentials {
api_key: Some("grok-oauth-tok".into()),
base_url: entry.base_url.clone(),
auth_type: kigi_chat_state::AuthType::SessionToken,
auth_scheme: Default::default(),
};
let cfg = crate::agent::config::sampling_config_for_model(&model_entry, creds, None);
assert_eq!(
cfg.chat_compat,
kigi_sampling_types::ChatCompat::Passthrough
);
}
/// INVARIANT: each platform's `/models` URL matches its registry base —
/// kimi-code → the subscription proxy (config override respected, else the
/// kigi-env default), moonshot platforms → their fixed bases — and the
@@ -3024,6 +3208,12 @@ mod tests {
platform_models_url(kigi_models::PlatformId::OpenAi, &cfg),
"https://api.openai.com/v1/models"
);
// xai-grok is uses_oauth but carries an OAuthConfig → its OWN base
// (api.x.ai/v1), NOT the Kimi subscription proxy.
assert_eq!(
platform_models_url(kigi_models::PlatformId::XaiGrok, &cfg),
"https://api.x.ai/v1/models"
);
// Proxy override re-points the subscription platform only.
let proxied = EndpointsConfig::from_config_value(
&toml::from_str(
@@ -3047,6 +3237,7 @@ mod tests {
&cfg,
ModelFetchAuth::Platforms,
true,
&Default::default(),
&PlatformApiKeys::default(),
);
assert_eq!(
@@ -3057,6 +3248,7 @@ mod tests {
&cfg,
ModelFetchAuth::Platforms,
true,
&Default::default(),
&crate::agent::models::PlatformApiKeys::test_keys(Some("sk-secret-cn"), None),
);
assert_ne!(
@@ -3082,6 +3274,7 @@ mod tests {
&custom,
ModelFetchAuth::CustomEndpoint,
false,
&Default::default(),
&PlatformApiKeys::default(),
),
"https://models.acme.com/v1/models"
crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs
+7 -1
View File
@@ -509,7 +509,13 @@ impl acp::Agent for MvpAgent {
Ok(self.auth_response_with_meta())
}
_ => {
if let Some(platform) =
if let Some(platform) = auth_method::AuthMethodKind::from_id(
&arguments.method_id,
)
.oauth_platform()
{
self.authenticate_oauth_platform(platform, arguments).await
} else if let Some(platform) =
auth_method::platform_for_method_id(&arguments.method_id)
{
self.authenticate_api_key_platform(
crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs
+116 -6
View File
@@ -25,7 +25,16 @@ impl MvpAgent {
primary: &SamplingConfig,
) -> Result<(OaiCompatClient, String), acp::Error> {
let slug = self.resolve_session_summary_model();
let session_key = self.auth_manager.current_or_expired().map(|a| a.key.clone());
// 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(),
&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();
@@ -514,6 +523,73 @@ impl MvpAgent {
.and_then(|v| v.as_object().cloned());
Ok(AuthenticateResponse::new().meta(meta))
}
/// `authenticate(<generic-oauth platform id>)`: interactive device-code
/// login for a `uses_oauth` platform carrying an `OAuthConfig` (xai-grok).
///
/// Uses a per-provider [`AuthManager`] scoped to the platform's `scope_key`
/// (NOT the primary Kimi manager) so the minted session is persisted under
/// its own `auth.json` scope, then triggers a catalog re-sync so the
/// platform's models appear (their bearer is resolved per-provider at fetch
/// / sampling time). The tokens are never logged.
pub(super) async fn authenticate_oauth_platform(
&self,
platform: kigi_models::PlatformId,
arguments: acp::AuthenticateRequest,
) -> Result<AuthenticateResponse, acp::Error> {
let method_id = arguments.method_id.clone();
let oauth = platform
.oauth()
.expect("oauth_platform() guarantees a device-code OAuthConfig");
let auth_meta = AuthRequestMeta::from_json(arguments.meta.as_ref());
tracing::info!(
method = method_id.0.as_ref(),
headless = auth_meta.headless,
reauth = auth_meta.reauth,
"auth: generic oauth device login",
);
let kigi_home = crate::util::kigi_home::kigi_home();
let auth_manager =
std::sync::Arc::new(crate::auth::AuthManager::new_oauth_provider(&kigi_home, oauth));
auth_manager.configure_refresher();
let flow_result = if !auth_meta.headless {
let (url_tx, url_rx) = tokio::sync::oneshot::channel();
let (code_tx, code_rx) = tokio::sync::mpsc::channel(1);
*self.auth_code_tx.borrow_mut() = Some(code_tx);
*self.auth_url_rx.borrow_mut() = Some(url_rx);
let result = crate::auth::run_oauth_provider_flow(
&auth_manager,
oauth,
auth_meta.reauth,
Some(crate::auth::AuthChannels {
url_tx: Some(url_tx),
code_rx,
}),
)
.await;
*self.auth_code_tx.borrow_mut() = None;
*self.auth_url_rx.borrow_mut() = None;
result
} else {
crate::auth::run_oauth_provider_flow(&auth_manager, oauth, auth_meta.reauth, None).await
};
let (_auth, _did_auth) = flow_result.map_err(|e| {
emit_login_span(false, method_id.0.as_ref(), None, Some("login_flow_failed"));
let mut err = acp::Error::auth_required();
err.message = e.to_string();
err
})?;
// Do NOT stamp this token onto the shared sampling_config: it authorizes
// ONLY this platform's models (api.x.ai/v1), not the primary session.
// The catalog re-sync below resolves it per-provider.
self.set_auth_method(method_id.clone());
self.models_manager.on_auth_changed().await;
emit_login_span(true, method_id.0.as_ref(), None, None);
Ok(self.auth_response_with_meta())
}
pub(crate) fn deployment_key(&self) -> Option<String> {
self.cfg.borrow().endpoints.deployment_key.clone()
}
@@ -593,16 +669,50 @@ 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.
///
/// 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.
fn session_token_for_model(&self, model: &ModelEntry) -> Option<crate::auth::KimiAuth> {
if let Some(oauth) = model
.info()
.id
.as_deref()
.and_then(kigi_models::parse_managed_model_key)
.and_then(|(platform, _)| platform.oauth())
{
return crate::auth::oauth_registry::global_manager_for(
&crate::util::kigi_home::kigi_home(),
oauth,
)
.current_or_expired();
}
if self.is_session_based_auth() {
self.auth_manager.current_or_expired()
} else {
None
}
}
pub(crate) fn prepare_sampling_config_for_model(
&self,
model: &ModelEntry,
origin_client: Option<crate::http::OriginClientInfo>,
) -> SamplingConfig {
let session = if self.is_session_based_auth() {
self.auth_manager.current_or_expired()
} else {
None
};
// Resolve the session token by the MODEL's platform, not the primary
// auth method: an oauth-platform model (xai-grok) uses its OWN
// pool-backed token (`None` — never the Kimi key — when the user has not
// logged into that provider), closing the api_key-channel leak where the
// primary Kimi session token was stamped onto a grok request. A
// first-party / Kimi model is unchanged. GUARANTEE: a grok model's
// api_key is its own grok token or `None`, never the primary Kimi key.
let session = self.session_token_for_model(model);
let has_session_key = session.is_some();
let mut credentials = resolve_credentials(
model,
crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs
+66
View File
@@ -1158,6 +1158,72 @@ fn build_agent_with_auth(auth: crate::auth::KimiAuth) -> MvpAgent {
let cfg = AgentConfig::default();
MvpAgent::new(gateway, &cfg, auth_manager, None).expect("valid test config")
}
/// Regression (token-leak, Facet B via the api_key channel): under a
/// session-based (Kimi) primary auth method, `prepare_sampling_config_for_model`
/// must resolve the session token by the MODEL's platform — so a grok
/// (oauth-platform) model NEVER carries the primary Kimi session key as its
/// `api_key`, while a first-party Kimi model still does. Relies on the
/// process-global OAuth pool (there is no per-session snapshot).
///
/// Reverting the fix (resolving the session token from the primary regardless of
/// the model's platform) fails the grok assertion below — it would stamp the
/// live Kimi key on a request bound for api.x.ai.
#[tokio::test]
#[serial_test::serial]
async fn prepare_sampling_config_never_stamps_kimi_key_on_grok_model() {
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::{EndpointsConfig, ModelEntry};
use kigi_test_support::EnvGuard;
const KIMI_KEY: &str = "kimi-session-secret-DO-NOT-LEAK";
// No ambient BYOK env key: a grok model with no stored oauth session then
// resolves to no api_key at all, rather than a global-key fallback that could
// mask the leak under test.
let _house = EnvGuard::unset(HOUSE_API_KEY_ENV_VAR);
let _xai = EnvGuard::unset(XAI_API_KEY_ENV_VAR);
let _legacy = EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR);
// Primary: a live Kimi session token under a session-based auth method.
let agent = build_agent_with_auth(crate::auth::KimiAuth {
key: KIMI_KEY.to_string(),
auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
agent.set_auth_method(acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID));
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
// 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());
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)"
);
// xai-grok model (oauth platform): the session token resolves from its OWN
// pool-backed manager, INDEPENDENT of the Kimi primary — so its api_key can
// NEVER be the Kimi session key.
let mut grok_model = ModelEntry::fallback("grok-4-latest", &endpoints);
grok_model.info.id = Some("xai-grok/grok-4-latest".to_string());
assert!(!grok_model.has_own_credentials());
let grok_cfg = agent.prepare_sampling_config_for_model(&grok_model, None);
assert_ne!(
grok_cfg.api_key.as_deref(),
Some(KIMI_KEY),
"LEAK: a grok model must never carry the primary Kimi session key as api_key"
);
}
/// Regression: boot-time plugin discovery is deferred past ACP
/// `initialize`, so the shared plugin registry starts empty.
/// `resolve_mcp_servers` reads that snapshot to merge plugin-contributed
crates/codegen/kigi-shell/src/agent/server.rs
+1
View File
@@ -176,6 +176,7 @@ async fn handle_connection(ws: WebSocket, state: Arc<ServerState>, peer_addr: So
prefetch_models_blocking(
&agent_config.endpoints,
auth.as_ref(),
&Default::default(),
fetch_auth,
&platform_keys,
)
crates/codegen/kigi-shell/src/agent/subagent/mod.rs
+13 -2
View File
@@ -982,9 +982,20 @@ fn resolve_model_override_to_config(
} else {
acp::ModelId::new(entry.info().model.clone())
};
let session_key = ctx.auth.as_ref().map(|a| a.key.as_str());
// 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,
Some(&ctx.auth_manager),
);
let has_session_key = session_key.is_some();
let mut credentials = resolve_credentials(&entry, session_key);
let mut credentials = resolve_credentials(&entry, session_key.as_deref());
credentials.auth_type = subagent_auth_type(Some(&entry), &ctx.auth_method_id);
let resolved_auth_type = credentials.auth_type;
let config = sampling_config_for_model(&entry, credentials, ctx.alpha_test_key.clone());
crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs
+70
View File
@@ -2467,6 +2467,76 @@ async fn resolve_subagent_config_override_unknown_model_falls_through_to_inherit
assert_eq!(config.model, "kigi-4.5");
assert_eq!(model_id.0.as_ref(), "kigi-4.5");
}
/// Build an `Arc<AuthManager>` (primary Kimi) holding `key` as its live bearer.
/// The `TempDir` is returned so the caller keeps it alive.
fn kimi_primary_with_token(key: &str) -> (tempfile::TempDir, std::sync::Arc<crate::auth::AuthManager>) {
let dir = tempfile::tempdir().unwrap();
let manager = std::sync::Arc::new(crate::auth::AuthManager::new(
dir.path(),
crate::auth::KimiCodeConfig::default(),
));
manager.hot_swap(crate::auth::KimiAuth {
key: key.to_string(),
auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
(dir, manager)
}
/// LEAK 2 (subagent model-override): a grok (oauth-platform) override with a
/// Kimi primary must NEVER receive the primary Kimi session token as its
/// `api_key` — it draws grok's own pooled token (or `None`). Revert-to-red: the
/// pre-fix code passed `ctx.auth` (Kimi) straight to `resolve_credentials`, so
/// `config.api_key == "kimi-secret"` and this assertion fails.
#[tokio::test]
async fn subagent_override_grok_model_never_leaks_kimi_session_token() {
let (_kd, manager) = kimi_primary_with_token("kimi-secret");
let mut grok = test_model_entry("grok-4-latest");
grok.info.id = Some("xai-grok/grok-4-latest".to_string());
grok.info.base_url = "https://api.x.ai/v1".to_string();
let mut models = indexmap::IndexMap::new();
models.insert("grok".to_string(), grok);
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("grok", &ctx).expect("grok override resolves to a config");
assert_ne!(
config.api_key.as_deref(),
Some("kimi-secret"),
"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).
#[tokio::test]
async fn subagent_override_non_oauth_model_still_gets_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());
let mut models = indexmap::IndexMap::new();
models.insert("k2".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("k2", &ctx).expect("non-oauth override resolves to a config");
assert_eq!(
config.api_key.as_deref(),
Some("kimi-secret"),
"a non-oauth override must still receive the primary session token",
);
}
/// An unresolvable `AgentDefinition.model` pin (model absent from
/// `available_models`) falls through to inherit the parent model.
#[tokio::test]