From dbce6bf30577d28220842c284f4124bad226b468 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Wed, 22 Jul 2026 01:36:29 -0400 Subject: [PATCH] feat(providers): add xAI Grok subscription OAuth (device-code) + per-provider session auth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- AGENTS.md | 15 + crates/codegen/kigi-models/src/lib.rs | 174 ++++++- crates/codegen/kigi-shell/src/agent/app.rs | 9 +- .../kigi-shell/src/agent/auth_method.rs | 132 +++++- crates/codegen/kigi-shell/src/agent/config.rs | 74 +++ crates/codegen/kigi-shell/src/agent/models.rs | 77 +++- .../kigi-shell/src/agent/models_fetch.rs | 271 +++++++++-- .../src/agent/mvp_agent/acp_agent.rs | 8 +- .../src/agent/mvp_agent/agent_ops.rs | 122 ++++- .../kigi-shell/src/agent/mvp_agent/tests.rs | 66 +++ crates/codegen/kigi-shell/src/agent/server.rs | 1 + .../kigi-shell/src/agent/subagent/mod.rs | 15 +- .../src/agent/subagent/tests/rest.rs | 70 +++ .../kigi-shell/src/auth/device_code.rs | 64 ++- crates/codegen/kigi-shell/src/auth/flow.rs | 31 ++ .../codegen/kigi-shell/src/auth/kimi_oauth.rs | 5 +- crates/codegen/kigi-shell/src/auth/manager.rs | 60 +++ crates/codegen/kigi-shell/src/auth/mod.rs | 5 +- .../kigi-shell/src/auth/oauth_device.rs | 423 ++++++++++++++++++ .../kigi-shell/src/auth/oauth_registry.rs | 257 +++++++++++ .../src/auth/refresh/generic_refresher.rs | 282 ++++++++++++ .../kigi-shell/src/auth/refresh/mod.rs | 20 +- .../session/acp_session_impl/model_switch.rs | 7 +- .../session/acp_session_impl/prompt_build.rs | 15 +- .../session/acp_session_impl/sampler_turn.rs | 250 +++++++++-- crates/codegen/kigi-tui/src/app/app_view.rs | 67 +-- 26 files changed, 2359 insertions(+), 161 deletions(-) create mode 100644 crates/codegen/kigi-shell/src/auth/oauth_device.rs create mode 100644 crates/codegen/kigi-shell/src/auth/oauth_registry.rs create mode 100644 crates/codegen/kigi-shell/src/auth/refresh/generic_refresher.rs diff --git a/AGENTS.md b/AGENTS.md index ddd7fcf..bb3b8db 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -168,6 +168,21 @@ edges stay deterministic Rust. The harness appends a terminal are built generically from advertised methods (`AuthMethodKind:: ApiKeyPlatform`), so new registry rows appear in the picker with no TUI changes. +- Refreshable-OAuth providers beyond Kimi Code use a GENERIC device-code + (RFC-8628) path, NOT Kimi's bespoke wire. A `uses_oauth` platform carrying + `oauth: Some(&OAuthConfig)` (client id / auth host / device+token paths / + scope / `scope_key` / optional extra device field) drives + `auth::oauth_device` (plain kigi UA, no X-Msh headers) + a scope-keyed + `AuthManager::new_oauth_provider` + `refresh::GenericDeviceRefresher` + (selected by `build_refresher` via `oauth_config_for_scope_key`). Kimi Code + keeps `oauth: None` and its bespoke path unchanged. First such provider: + `xai-grok` (`scope_key oauth/xai`, base `api.x.ai/v1`, same wire as the + API-key `xai` row) — an INTERACTIVE login row advertised right after + `kimi-code` (`AuthMethodKind::OAuthPlatform`). Its `authenticate` arm runs + the generic device flow under its own scope; the catalog fetch resolves each + such platform's OWN session token (`resolve_generic_oauth_tokens`, refreshed + on expiry) and routes `platform.oauth().is_some()` → `platform.base_url()` + (kimi-code alone → `proxy_url()`). Tokens are NEVER logged. - Model metadata (context window, thinking levels) comes from the provider wire when served; metadata-poor listings are enriched from models.dev (`kigi-models/src/enrichment.rs` — bundled raw snapshot regenerated by diff --git a/crates/codegen/kigi-models/src/lib.rs b/crates/codegen/kigi-models/src/lib.rs index a340664..b583b9a 100644 --- a/crates/codegen/kigi-models/src/lib.rs +++ b/crates/codegen/kigi-models/src/lib.rs @@ -96,6 +96,53 @@ enum BaseUrlSource { }, } +/// Generic RFC-8628 device-code OAuth configuration carried by a `uses_oauth` +/// platform whose login is the GENERIC device-code path (xai-grok today). +/// +/// Kimi Code keeps its bespoke device flow (client id, `/api/oauth/*` paths, +/// X-Msh device headers, `kigi_env::oauth_host()`); its `oauth` field stays +/// `None`. All fields here are non-secret wire constants — the access/refresh +/// tokens they mint are NEVER stored in this struct and never logged. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct OAuthConfig { + /// OAuth client id sent on every device-authorization / token call. + pub client_id: &'static str, + /// Authorization-server origin (no trailing slash), e.g. `https://auth.x.ai`. + pub auth_host: &'static str, + /// Device-authorization path (POST), relative to `auth_host`. + pub device_path: &'static str, + /// Token path (POST) — used for BOTH the device grant and refresh. + pub token_path: &'static str, + /// OAuth scope string requested at device authorization. + pub scope: &'static str, + /// auth.json map key + keyring entry name for this provider's persisted + /// session (e.g. `oauth/xai`). + pub scope_key: &'static str, + /// A non-standard extra form field sent ONLY on the device-authorization + /// request (e.g. `("referrer", "kigi")`). `None` = no extra field. + pub extra_device_field: Option<(&'static str, &'static str)>, +} + +/// xAI / Grok subscription device-code OAuth (ported from Pi +/// `earendil-works/pi` `auth/oauth/xai.ts`). +pub const XAI_OAUTH_CONFIG: OAuthConfig = OAuthConfig { + client_id: "b1a00492-073a-47ea-816f-4c329264a828", + auth_host: "https://auth.x.ai", + device_path: "/oauth2/device/code", + token_path: "/oauth2/token", + scope: "openid profile email offline_access grok-cli:access api:access", + scope_key: "oauth/xai", + extra_device_field: Some(("referrer", "kigi")), +}; + +/// The generic device-code OAuth config for a platform, or `None` for API-key +/// platforms and for Kimi Code (whose bespoke flow uses no generic config). +pub fn oauth_config_for_scope_key(scope_key: &str) -> Option<&'static OAuthConfig> { + PlatformId::ALL + .into_iter() + .find_map(|p| p.oauth().filter(|c| c.scope_key == scope_key)) +} + /// One row of the platform registry. All per-platform data lives in these /// rows; [`PlatformId`] methods only read fields. Adding a platform touches /// exactly four sites in this file: the enum variant, the `ALL` entry, the @@ -111,6 +158,11 @@ struct PlatformSpec { base_url: BaseUrlSource, /// True for OAuth-bearer subscription channels. uses_oauth: bool, + /// Generic device-code OAuth config, for a `uses_oauth` platform whose + /// login is the RFC-8628 device-code path (xai-grok). `None` for API-key + /// platforms AND for Kimi Code (whose bespoke device flow lives in + /// kigi-shell, not this generic config). + oauth: Option<&'static OAuthConfig>, /// Model-id prefixes admitted from this platform's `/models` listing. /// `None` = no filtering (listing served pre-filtered). allowed_model_prefixes: Option<&'static [&'static str]>, @@ -166,6 +218,9 @@ const KIMI_CODE_SPEC: PlatformSpec = PlatformSpec { display_name: "Kimi Code", base_url: BaseUrlSource::KigiEnvCoding, uses_oauth: true, + // Kimi Code keeps its bespoke device flow (kigi-shell), not the generic + // OAuthConfig path — see architecture decision #1. + oauth: None, allowed_model_prefixes: None, api_key_envs: &[], vendor: "Kimi", @@ -190,6 +245,7 @@ const MOONSHOT_CN_SPEC: PlatformSpec = PlatformSpec { default: "https://api.moonshot.cn/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: Some(&["kimi-k"]), api_key_envs: &[MOONSHOT_CN_API_KEY_ENV, MOONSHOT_API_KEY_ENV], vendor: "Moonshot", @@ -214,6 +270,7 @@ const MOONSHOT_AI_SPEC: PlatformSpec = PlatformSpec { default: "https://api.moonshot.ai/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: Some(&["kimi-k"]), api_key_envs: &[MOONSHOT_AI_API_KEY_ENV, MOONSHOT_API_KEY_ENV], vendor: "Moonshot", @@ -241,6 +298,7 @@ const OPENAI_SPEC: PlatformSpec = PlatformSpec { default: "https://api.openai.com/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["OPENAI_API_KEY"], vendor: "OpenAI", @@ -270,6 +328,7 @@ const ANTHROPIC_SPEC: PlatformSpec = PlatformSpec { default: "https://api.anthropic.com/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["ANTHROPIC_API_KEY"], vendor: "Anthropic", @@ -301,6 +360,7 @@ const DEEPSEEK_SPEC: PlatformSpec = PlatformSpec { default: "https://api.deepseek.com", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["DEEPSEEK_API_KEY"], vendor: "DeepSeek", @@ -328,6 +388,7 @@ const GROQ_SPEC: PlatformSpec = PlatformSpec { default: "https://api.groq.com/openai/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["GROQ_API_KEY"], vendor: "Groq", @@ -357,6 +418,7 @@ const MISTRAL_SPEC: PlatformSpec = PlatformSpec { default: "https://api.mistral.ai/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["MISTRAL_API_KEY"], vendor: "Mistral", @@ -390,6 +452,7 @@ const FIREWORKS_SPEC: PlatformSpec = PlatformSpec { default: "https://api.fireworks.ai/inference/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["FIREWORKS_API_KEY"], vendor: "Fireworks", @@ -420,6 +483,7 @@ const GOOGLE_SPEC: PlatformSpec = PlatformSpec { default: "https://generativelanguage.googleapis.com/v1beta/openai", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["GEMINI_API_KEY"], vendor: "Google", @@ -450,6 +514,7 @@ const OPENROUTER_SPEC: PlatformSpec = PlatformSpec { default: "https://openrouter.ai/api/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["OPENROUTER_API_KEY"], vendor: "OpenRouter", @@ -480,6 +545,7 @@ const TOGETHER_SPEC: PlatformSpec = PlatformSpec { default: "https://api.together.xyz/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["TOGETHER_API_KEY"], vendor: "Together", @@ -510,6 +576,7 @@ const CEREBRAS_SPEC: PlatformSpec = PlatformSpec { default: "https://api.cerebras.ai/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["CEREBRAS_API_KEY"], vendor: "Cerebras", @@ -542,6 +609,7 @@ const NVIDIA_SPEC: PlatformSpec = PlatformSpec { default: "https://integrate.api.nvidia.com/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["NVIDIA_API_KEY"], vendor: "NVIDIA", @@ -574,6 +642,7 @@ const VERCEL_SPEC: PlatformSpec = PlatformSpec { default: "https://ai-gateway.vercel.sh/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["AI_GATEWAY_API_KEY"], vendor: "Vercel", @@ -606,6 +675,7 @@ const XAI_SPEC: PlatformSpec = PlatformSpec { default: "https://api.x.ai/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, // NOTE: `XAI_API_KEY` is ALSO read as a legacy fallback by the house BYOK // path (`read_xai_api_key_env`), whose primary env is now `KIGI_API_KEY`. @@ -632,6 +702,41 @@ const XAI_SPEC: PlatformSpec = PlatformSpec { restrict_to_enriched: true, }; +/// Base-URL override for the xAI Grok subscription (OAuth) platform. Its own +/// env (distinct from `KIGI_XAI_BASE_URL`) so a mock can target it +/// independently of the API-key `xai` row. +pub const XAI_GROK_BASE_URL_ENV: &str = "KIGI_XAI_GROK_BASE_URL"; +const XAI_GROK_SPEC: PlatformSpec = PlatformSpec { + id: "xai-grok", + display_name: "xAI (Grok subscription)", + // Same wire as the API-key `xai` row (api.x.ai/v1), but reached with an + // OAuth subscription bearer instead of an API key. + base_url: BaseUrlSource::EnvOr { + env: XAI_GROK_BASE_URL_ENV, + default: "https://api.x.ai/v1", + }, + uses_oauth: true, + oauth: Some(&XAI_OAUTH_CONFIG), + allowed_model_prefixes: None, + // OAuth channel: no API key envs (the device-flow session is the bearer). + api_key_envs: &[], + vendor: "xAI", + console_host: Some("x.ai"), + login_label: Some("xAI Grok (subscription)"), + models_dev_id: Some("xai"), + // Same as XAI_SPEC: bare-id /v1/models, enrich from models.dev "xai", + // restrict to tool-calling chat models (drops the grok-imagine-* media + // generators). + wire_serves_metadata: false, + wire_api: PlatformWireApi::ChatCompletions, + listing: ListingDialect::OpenAi, + chat_compat: PlatformChatCompat::Passthrough, + key_header: PlatformKeyHeader::Bearer, + key_validation_path: None, + strip_listing_id_prefix: None, + restrict_to_enriched: true, +}; + pub const QWEN_TOKEN_PLAN_BASE_URL_ENV: &str = "KIGI_QWEN_TOKEN_PLAN_BASE_URL"; const QWEN_TOKEN_PLAN_SPEC: PlatformSpec = PlatformSpec { id: "qwen-token-plan", @@ -641,6 +746,7 @@ const QWEN_TOKEN_PLAN_SPEC: PlatformSpec = PlatformSpec { default: "https://token-plan.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["QWEN_TOKEN_PLAN_API_KEY"], vendor: "Alibaba", @@ -672,6 +778,7 @@ const QWEN_TOKEN_PLAN_CN_SPEC: PlatformSpec = PlatformSpec { default: "https://token-plan.cn-beijing.maas.aliyuncs.com/compatible-mode/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["QWEN_TOKEN_PLAN_CN_API_KEY"], vendor: "Alibaba", @@ -698,6 +805,7 @@ const KIMI_CODING_SPEC: PlatformSpec = PlatformSpec { // Kimi thinking dialect applies verbatim. base_url: BaseUrlSource::KigiEnvCoding, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["KIMI_API_KEY"], vendor: "Kimi", @@ -727,6 +835,7 @@ const ZAI_SPEC: PlatformSpec = PlatformSpec { default: "https://api.z.ai/api/coding/paas/v4", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["ZAI_API_KEY"], vendor: "Z.AI", @@ -756,6 +865,7 @@ const ZAI_CODING_CN_SPEC: PlatformSpec = PlatformSpec { default: "https://open.bigmodel.cn/api/coding/paas/v4", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["ZAI_CODING_CN_API_KEY"], vendor: "Z.AI", @@ -782,6 +892,7 @@ const XIAOMI_SPEC: PlatformSpec = PlatformSpec { default: "https://api.xiaomimimo.com/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["XIAOMI_API_KEY"], vendor: "Xiaomi", @@ -810,6 +921,7 @@ const XIAOMI_TOKEN_PLAN_CN_SPEC: PlatformSpec = PlatformSpec { default: "https://token-plan-cn.xiaomimimo.com/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["XIAOMI_TOKEN_PLAN_CN_API_KEY"], vendor: "Xiaomi", @@ -841,6 +953,7 @@ const MINIMAX_SPEC: PlatformSpec = PlatformSpec { default: "https://api.minimax.io/anthropic/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["MINIMAX_API_KEY"], vendor: "MiniMax", @@ -871,6 +984,7 @@ const MINIMAX_CN_SPEC: PlatformSpec = PlatformSpec { default: "https://api.minimaxi.com/anthropic/v1", }, uses_oauth: false, + oauth: None, allowed_model_prefixes: None, api_key_envs: &["MINIMAX_CN_API_KEY"], vendor: "MiniMax", @@ -941,12 +1055,14 @@ pub enum PlatformId { Minimax, /// MiniMax, China (API key, Anthropic-compatible Messages). MinimaxCn, + /// xAI Grok subscription via device-code OAuth (same wire as `Xai`). + XaiGrok, } impl PlatformId { /// All platforms, in catalog precedence order: the subscription channel /// first so "default model = first list item" favors it when present. - pub const ALL: [PlatformId; 25] = [ + pub const ALL: [PlatformId; 26] = [ Self::KimiCode, Self::MoonshotCn, Self::MoonshotAi, @@ -972,6 +1088,7 @@ impl PlatformId { Self::XiaomiTokenPlanCn, Self::Minimax, Self::MinimaxCn, + Self::XaiGrok, ]; /// The registry row backing this platform (single source of per-platform @@ -1003,6 +1120,7 @@ impl PlatformId { Self::XiaomiTokenPlanCn => &XIAOMI_TOKEN_PLAN_CN_SPEC, Self::Minimax => &MINIMAX_SPEC, Self::MinimaxCn => &MINIMAX_CN_SPEC, + Self::XaiGrok => &XAI_GROK_SPEC, } } @@ -1034,6 +1152,14 @@ impl PlatformId { self.spec().uses_oauth } + /// Generic device-code OAuth config for this platform, or `None` for + /// API-key platforms and for Kimi Code (bespoke flow). A `Some` here means + /// `uses_oauth()` is also true; the fetch/base-url routing treats these + /// providers as their own `base_url()` (not the Kimi `proxy_url()`). + pub fn oauth(self) -> Option<&'static OAuthConfig> { + self.spec().oauth + } + /// Model-id prefixes admitted from this platform's `/models` listing. /// `None` = no filtering (subscription listing is served pre-filtered). pub fn allowed_model_prefixes(self) -> Option<&'static [&'static str]> { @@ -1764,6 +1890,49 @@ mod tests { assert_eq!(PlatformId::parse("not-a-platform"), None); } + /// xai-grok is the second refreshable-OAuth platform: it carries a generic + /// device-code `OAuthConfig` (Kimi Code carries `None`), reuses the xai + /// wire (models.dev "xai", restrict-to-enriched, Bearer, ChatCompletions), + /// and keys its models under `xai-grok/` — distinct from the API-key `xai`. + #[test] + fn xai_grok_is_a_generic_oauth_platform() { + let g = PlatformId::XaiGrok; + assert_eq!(g.as_str(), "xai-grok"); + assert!(g.uses_oauth()); + // The only two uses_oauth platforms; only xai-grok carries a config. + assert!(PlatformId::KimiCode.uses_oauth()); + assert_eq!(PlatformId::KimiCode.oauth(), None); + let cfg = g + .oauth() + .expect("xai-grok carries a device-code OAuthConfig"); + assert_eq!(cfg, &XAI_OAUTH_CONFIG); + assert_eq!(cfg.client_id, "b1a00492-073a-47ea-816f-4c329264a828"); + assert_eq!(cfg.auth_host, "https://auth.x.ai"); + assert_eq!(cfg.device_path, "/oauth2/device/code"); + assert_eq!(cfg.token_path, "/oauth2/token"); + assert_eq!( + cfg.scope, + "openid profile email offline_access grok-cli:access api:access" + ); + assert_eq!(cfg.scope_key, "oauth/xai"); + assert_eq!(cfg.extra_device_field, Some(("referrer", "kigi"))); + // Scope-key lookup resolves the config (drives the generic refresher). + assert_eq!( + oauth_config_for_scope_key("oauth/xai"), + Some(&XAI_OAUTH_CONFIG) + ); + assert_eq!(oauth_config_for_scope_key("oauth/kimi-code"), None); + // Reuses the xai wire, its OWN base env, and the models.dev "xai" keys. + assert_eq!(g.models_dev_id(), Some("xai")); + assert!(g.restrict_to_enriched()); + assert_eq!(g.key_header(), PlatformKeyHeader::Bearer); + assert_eq!(g.wire_api(), PlatformWireApi::ChatCompletions); + assert_eq!(g.api_key_env_names(), &[] as &[&str]); + assert_eq!(g.managed_model_key("grok-4.5"), "xai-grok/grok-4.5"); + let _guard = kigi_env::EnvVarGuard::set(XAI_GROK_BASE_URL_ENV, "https://mock.grok/v1"); + assert_eq!(g.base_url(), "https://mock.grok/v1"); + } + /// A variant missing from `ALL` compiles fine (`ALL`'s length is a plain /// literal) but is silently unparseable and excluded from model sync. /// The exhaustive match below fails compilation when a variant is added, @@ -1797,9 +1966,10 @@ mod tests { PlatformId::XiaomiTokenPlanCn => 22, PlatformId::Minimax => 23, PlatformId::MinimaxCn => 24, + PlatformId::XaiGrok => 25, } } - const VARIANT_COUNT: usize = 25; // update together with `ordinal` + const VARIANT_COUNT: usize = 26; // update together with `ordinal` let mut seen: Vec = PlatformId::ALL.iter().map(|&p| ordinal(p)).collect(); seen.sort_unstable(); seen.dedup(); diff --git a/crates/codegen/kigi-shell/src/agent/app.rs b/crates/codegen/kigi-shell/src/agent/app.rs index 9fc917d..a689d1a 100644 --- a/crates/codegen/kigi-shell/src/agent/app.rs +++ b/crates/codegen/kigi-shell/src/agent/app.rs @@ -178,7 +178,13 @@ async fn prefetch_models(agent_config: &AgentConfig) -> Option) -> 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 { + 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.]` 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", diff --git a/crates/codegen/kigi-shell/src/agent/config.rs b/crates/codegen/kigi-shell/src/agent/config.rs index 2830ee5..2ef7fdb 100644 --- a/crates/codegen/kigi-shell/src/agent/config.rs +++ b/crates/codegen/kigi-shell/src/agent/config.rs @@ -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) { + 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( diff --git a/crates/codegen/kigi-shell/src/agent/models.rs b/crates/codegen/kigi-shell/src/agent/models.rs index 0dc4f87..8c06ebb 100644 --- a/crates/codegen/kigi-shell/src/agent/models.rs +++ b/crates/codegen/kigi-shell/src/agent/models.rs @@ -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 { + 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> { 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, + 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, diff --git a/crates/codegen/kigi-shell/src/agent/models_fetch.rs b/crates/codegen/kigi-shell/src/agent/models_fetch.rs index 846936b..acdd9de 100644 --- a/crates/codegen/kigi-shell/src/agent/models_fetch.rs +++ b/crates/codegen/kigi-shell/src/agent/models_fetch.rs @@ -28,6 +28,44 @@ pub(crate) const DEFAULT_CONTEXT_WINDOW: u64 = 256_000; struct ModelsResponse { data: Vec, } + +/// 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; + +/// 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 = enabled_platforms(has_oauth, platform_keys) + let parts: Vec = 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::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 { @@ -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 { - 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/` 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!["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" diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs index 2ee944e..dada96e 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs @@ -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( diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs index f17ff46..e683f0a 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs @@ -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()`: 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 { + 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 { 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 { + 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, ) -> 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, diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs index 65c6cde..42fa85e 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs @@ -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 diff --git a/crates/codegen/kigi-shell/src/agent/server.rs b/crates/codegen/kigi-shell/src/agent/server.rs index 98554c3..283fca4 100644 --- a/crates/codegen/kigi-shell/src/agent/server.rs +++ b/crates/codegen/kigi-shell/src/agent/server.rs @@ -176,6 +176,7 @@ async fn handle_connection(ws: WebSocket, state: Arc, peer_addr: So prefetch_models_blocking( &agent_config.endpoints, auth.as_ref(), + &Default::default(), fetch_auth, &platform_keys, ) diff --git a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs index a79dcc3..db5fa31 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs @@ -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()); diff --git a/crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs b/crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs index c9c246e..1ff4a61 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/tests/rest.rs @@ -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` (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) { + 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] diff --git a/crates/codegen/kigi-shell/src/auth/device_code.rs b/crates/codegen/kigi-shell/src/auth/device_code.rs index e2c1f97..e093163 100644 --- a/crates/codegen/kigi-shell/src/auth/device_code.rs +++ b/crates/codegen/kigi-shell/src/auth/device_code.rs @@ -13,15 +13,47 @@ use std::sync::Arc; -use crate::auth::kimi_oauth::{ - DeviceAuthorization, DevicePollResult, poll_device_token, request_device_authorization, -}; +use kigi_models::OAuthConfig; + +use crate::auth::kimi_oauth::{DeviceAuthorization, DevicePollResult}; use crate::auth::{AuthChannels, AuthManager, AuthUrlInfo, AuthUrlMode, KimiAuth}; /// Extra wait added to the poll interval when the server answers `slow_down` /// (OAuth-standard device-flow backpressure). const SLOW_DOWN_INCREMENT_SECS: u64 = 5; +/// The wire behind a device-code login. The `Kimi` arm calls the bespoke Kimi +/// Code wire (X-Msh headers, `/api/oauth/*`) verbatim — byte-identical to the +/// pre-generalization path; the `Generic` arm drives a registry +/// [`OAuthConfig`] provider (xai-grok) through [`crate::auth::oauth_device`]. +enum DeviceFlowBackend<'a> { + Kimi { host: &'a str }, + Generic(&'a OAuthConfig), +} + +impl DeviceFlowBackend<'_> { + async fn request(&self) -> anyhow::Result { + match self { + Self::Kimi { host } => { + crate::auth::kimi_oauth::request_device_authorization(host).await + } + Self::Generic(cfg) => { + crate::auth::oauth_device::request_device_authorization(cfg).await + } + } + } + async fn poll(&self, device_code: &str) -> anyhow::Result { + match self { + Self::Kimi { host } => { + crate::auth::kimi_oauth::poll_device_token(host, device_code).await + } + Self::Generic(cfg) => { + crate::auth::oauth_device::poll_device_token(cfg, device_code).await + } + } + } +} + /// Outcome of one full poll loop over a single device authorization. enum PollLoopOutcome { /// Access token issued. @@ -40,11 +72,29 @@ pub async fn run_device_code_login_channels( host: &str, auth_manager: &Arc, channels: &mut Option, +) -> anyhow::Result<(KimiAuth, bool)> { + run_device_code_login_backend(DeviceFlowBackend::Kimi { host }, auth_manager, channels).await +} + +/// Device-code login for a GENERIC [`OAuthConfig`] provider (xai-grok). Same +/// TUI/CLI presentation as the Kimi login; only the wire differs. +pub async fn run_device_code_login_generic( + oauth: &OAuthConfig, + auth_manager: &Arc, + channels: &mut Option, +) -> anyhow::Result<(KimiAuth, bool)> { + run_device_code_login_backend(DeviceFlowBackend::Generic(oauth), auth_manager, channels).await +} + +async fn run_device_code_login_backend( + backend: DeviceFlowBackend<'_>, + auth_manager: &Arc, + channels: &mut Option, ) -> anyhow::Result<(KimiAuth, bool)> { let interactive_tui = channels.is_some(); let mut channels = channels.take(); loop { - let device_auth = request_device_authorization(host).await?; + let device_auth = backend.request().await?; let display_uri = device_auth.verification_uri_complete.clone(); if interactive_tui { @@ -62,7 +112,7 @@ pub async fn run_device_code_login_channels( prompt_on_stderr(&device_auth).await; } - match complete_device_code_login(host, &device_auth).await? { + match complete_device_code_login(&backend, &device_auth).await? { PollLoopOutcome::Done(auth) => { let auth = auth_manager .update(*auth) @@ -112,7 +162,7 @@ async fn prompt_on_stderr(device_auth: &DeviceAuthorization) { /// Poll the token endpoint until the user approves, the device code expires /// (→ [`PollLoopOutcome::Restart`]), or the wire fails. async fn complete_device_code_login( - host: &str, + backend: &DeviceFlowBackend<'_>, device_auth: &DeviceAuthorization, ) -> anyhow::Result { let mut poll_interval = std::time::Duration::from_secs(device_auth.interval.max(1) as u64); @@ -120,7 +170,7 @@ async fn complete_device_code_login( // Sleep first: an immediate poll on a fresh code only returns // authorization_pending (and risks slow_down). tokio::time::sleep(poll_interval).await; - match poll_device_token(host, &device_auth.device_code).await? { + match backend.poll(&device_auth.device_code).await? { DevicePollResult::Success(auth) => { tracing::info!("auth: device login authorized"); return Ok(PollLoopOutcome::Done(auth)); diff --git a/crates/codegen/kigi-shell/src/auth/flow.rs b/crates/codegen/kigi-shell/src/auth/flow.rs index a65e95d..e6bb001 100644 --- a/crates/codegen/kigi-shell/src/auth/flow.rs +++ b/crates/codegen/kigi-shell/src/auth/flow.rs @@ -67,6 +67,37 @@ pub async fn run_auth_flow( run_auth_flow_inner(auth_manager, kimi_code_config, reauth, false, channels).await } +/// Login flow for a GENERIC device-code OAuth provider (xai-grok): use a valid +/// cached session unless re-authing, otherwise run the generic device flow +/// (persisting under the provider's own scope via `auth_manager`). Unlike the +/// Kimi flow this does not run the silent-refresh dance — the device flow's +/// `AuthManager::update` persists a fresh token set directly. +pub async fn run_oauth_provider_flow( + auth_manager: &Arc, + oauth: &'static kigi_models::OAuthConfig, + reauth: bool, + channels: Option, +) -> anyhow::Result<(KimiAuth, bool)> { + tracing::info!( + scope_key = oauth.scope_key, + reauth, + "auth: starting generic oauth login" + ); + if reauth { + auth_manager.clear()?; + } + if !reauth && let Some(auth) = auth_manager.current() { + tracing::info!( + scope_key = oauth.scope_key, + "auth: using cached oauth session" + ); + return Ok((auth, false)); + } + let mut channels = channels; + crate::auth::device_code::run_device_code_login_generic(oauth, auth_manager, &mut channels) + .await +} + async fn run_auth_flow_inner( auth_manager: &Arc, _kimi_code_config: &KimiCodeConfig, diff --git a/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs b/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs index cb0dd33..30cb4cf 100644 --- a/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs +++ b/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs @@ -130,8 +130,9 @@ fn with_device_headers( } /// Defend against control characters / non-https redirects from a -/// compromised or mis-configured OAuth host. -fn validate_verification_uri(uri: &str) -> anyhow::Result<()> { +/// compromised or mis-configured OAuth host. Shared with the generic +/// device-code wire ([`super::oauth_device`]). +pub(crate) fn validate_verification_uri(uri: &str) -> anyhow::Result<()> { if uri.chars().any(|c| c.is_ascii_control()) { anyhow::bail!("Server returned invalid verification URI"); } diff --git a/crates/codegen/kigi-shell/src/auth/manager.rs b/crates/codegen/kigi-shell/src/auth/manager.rs index cfdfc46..5bb7bb4 100644 --- a/crates/codegen/kigi-shell/src/auth/manager.rs +++ b/crates/codegen/kigi-shell/src/auth/manager.rs @@ -385,6 +385,59 @@ impl AuthManager { ) } + /// Build a manager for a GENERIC device-code OAuth provider (xai-grok), + /// scoped to `oauth.scope_key`. Unlike [`Self::new`] this path is + /// file-store only (no keyring — that is gated to the default Kimi install) + /// and ignores the Kimi-specific `KIGI_AUTH` inline-credential env; it + /// otherwise shares the same multi-scope `auth.json` (honoring + /// `KIGI_AUTH_PATH`). The refresher is selected from the scope by + /// [`super::refresh::build_refresher`]. + pub(crate) fn new_oauth_provider( + kigi_home: &Path, + oauth: &'static kigi_models::OAuthConfig, + ) -> Self { + let scope = oauth.scope_key.to_owned(); + let path = std::env::var("KIGI_AUTH_PATH") + .map(PathBuf::from) + .unwrap_or_else(|_| kigi_home.join("auth.json")); + + let (auth, disk_state) = match read_auth_json(&path) { + Ok(map) => { + let found = lookup_auth(&map, &scope); + let state = if found.is_some() { + DiskAuthState::Ok + } else { + DiskAuthState::EntryMissing + }; + (found, state) + } + Err(e) => { + let state = if e.kind() == std::io::ErrorKind::NotFound { + DiskAuthState::FileMissing + } else { + DiskAuthState::Unreadable + }; + (None, state) + } + }; + kigi_log::unified_log::info( + "AuthManager::new_oauth_provider", + None, + Some(serde_json::json!({ + "scope": &scope, + "found": auth.is_some(), + "is_expired": auth.as_ref().map(is_expired), + })), + ); + Self::assemble( + auth, + path, + scope, + KimiCodeConfig::default(), + Some(disk_state), + ) + } + /// Single field-assembly point for [`Self::new`]'s two construction paths /// (inline `KIGI_AUTH` vs. on-disk `auth.json`), which differ only in the /// threaded fields. One literal means a newly added field can't be silently @@ -789,6 +842,13 @@ impl AuthManager { &self.kimi_code_config } + /// The auth.json / keyring scope key this manager persists under + /// (`oauth/kimi-code` for Kimi, `oauth/xai` for xai-grok, …). Drives the + /// refresher selection in [`super::refresh::build_refresher`]. + pub(crate) fn scope(&self) -> &str { + &self.scope + } + /// Handle notified after every successful token refresh. /// /// Used by [`ModelsManager`] to trigger model catalog recovery diff --git a/crates/codegen/kigi-shell/src/auth/mod.rs b/crates/codegen/kigi-shell/src/auth/mod.rs index 9615afa..4068984 100644 --- a/crates/codegen/kigi-shell/src/auth/mod.rs +++ b/crates/codegen/kigi-shell/src/auth/mod.rs @@ -8,6 +8,8 @@ mod flow; pub(crate) mod kimi_oauth; pub(crate) mod manager; mod model; +pub(crate) mod oauth_device; +pub(crate) mod oauth_registry; pub(crate) mod recovery; pub(crate) mod refresh; mod storage; @@ -17,7 +19,8 @@ pub(crate) use flow::try_ensure_session_noninteractive; pub use flow::{ AuthChannels, AuthUrlInfo, AuthUrlMode, LogoutResult, ensure_authenticated, ensure_authenticated_or_noninteractive, perform_logout, run_auth_flow, - run_auth_flow_with_stderr_bridge, run_cli_login, run_cli_logout, try_ensure_fresh_auth, + run_auth_flow_with_stderr_bridge, run_cli_login, run_cli_logout, run_oauth_provider_flow, + try_ensure_fresh_auth, }; mod meta; pub use device::device_headers; diff --git a/crates/codegen/kigi-shell/src/auth/oauth_device.rs b/crates/codegen/kigi-shell/src/auth/oauth_device.rs new file mode 100644 index 0000000..42c735c --- /dev/null +++ b/crates/codegen/kigi-shell/src/auth/oauth_device.rs @@ -0,0 +1,423 @@ +//! Generic RFC-8628 device-code OAuth wire, driven by a registry +//! [`kigi_models::OAuthConfig`] (xai-grok today; Copilot/Claude later). +//! +//! Three `application/x-www-form-urlencoded` POSTs against `{auth_host}`: +//! +//! - `POST {device_path}` — form `client_id` + `scope` + the optional +//! `extra_device_field` (e.g. `referrer=kigi`) +//! - `POST {token_path}` (poll) — form `client_id` + `device_code` + +//! `grant_type=urn:ietf:params:oauth:grant-type:device_code` +//! - `POST {token_path}` (refresh) — form `client_id` + +//! `grant_type=refresh_token` + `refresh_token`, with the same exponential +//! backoff / status handling as the Kimi wire. +//! +//! Unlike [`super::kimi_oauth`] this sends NO X-Msh device headers — just the +//! shared kigi `User-Agent` and `Accept: application/json`. Access/refresh +//! tokens are NEVER logged (only non-secret events: requested, poll succeeded, +//! refreshed). + +use kigi_models::OAuthConfig; +use serde::Deserialize; + +use super::kimi_oauth::{ + DeviceAuthorization, DevicePollResult, RefreshError, TokenResponse, validate_verification_uri, +}; + +const DEVICE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code"; +const REFRESH_GRANT_TYPE: &str = "refresh_token"; +/// Refresh retry budget over the retryable statuses / network blips. +const MAX_REFRESH_RETRIES: u32 = 3; +/// HTTP statuses worth retrying a refresh for (kimi-cli parity). +const RETRYABLE_REFRESH_STATUSES: [u16; 5] = [429, 500, 502, 503, 504]; + +#[derive(Deserialize)] +struct DeviceAuthorizationResponse { + user_code: String, + device_code: String, + #[serde(default)] + verification_uri: Option, + /// Optional here (the Kimi wire requires it): Pi's xAI response may omit + /// `verification_uri_complete` and carry only `verification_uri`. + #[serde(default)] + verification_uri_complete: Option, + #[serde(default)] + expires_in: Option, + #[serde(default)] + interval: Option, +} + +#[derive(Deserialize, Default)] +struct OAuthErrorBody { + #[serde(default)] + error: Option, + #[serde(default)] + error_description: Option, +} + +fn oauth_url(host: &str, path: &str) -> String { + format!("{}{path}", host.trim_end_matches('/')) +} + +/// The device-authorization form fields: `client_id`, `scope`, and the +/// optional non-standard `extra_device_field`. +fn device_form(cfg: &OAuthConfig) -> Vec<(&'static str, &'static str)> { + let mut form = vec![("client_id", cfg.client_id), ("scope", cfg.scope)]; + if let Some((name, value)) = cfg.extra_device_field { + form.push((name, value)); + } + form +} + +/// `POST {auth_host}{device_path}` — start a device login. +pub(crate) async fn request_device_authorization( + cfg: &OAuthConfig, +) -> anyhow::Result { + let url = oauth_url(cfg.auth_host, cfg.device_path); + tracing::info!(url = %url, "auth: requesting device authorization (generic oauth)"); + let resp = crate::http::shared_client() + .post(&url) + .header("Accept", "application/json") + .form(&device_form(cfg)) + .send() + .await?; + + let status = resp.status(); + if !status.is_success() { + let body = resp.text().await.unwrap_or_default(); + tracing::warn!(%status, "auth: device authorization failed (generic oauth)"); + anyhow::bail!("Device authorization failed (HTTP {status}): {body}"); + } + let parsed: DeviceAuthorizationResponse = resp.json().await?; + + if !parsed + .user_code + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '-') + { + anyhow::bail!("Server returned invalid user_code format (expected [A-Z0-9-])"); + } + // Pi forces the displayed URI to https; we require a valid https (or + // localhost) verification target, preferring the pre-filled complete form. + let verification_uri_complete = parsed + .verification_uri_complete + .clone() + .or_else(|| parsed.verification_uri.clone()) + .ok_or_else(|| anyhow::anyhow!("Server returned no verification URI"))?; + validate_verification_uri(&verification_uri_complete)?; + if let Some(ref uri) = parsed.verification_uri { + validate_verification_uri(uri)?; + } + + tracing::info!( + user_code = %parsed.user_code, + interval = parsed.interval.unwrap_or(5), + expires_in = ?parsed.expires_in, + "auth: device authorization issued (generic oauth)" + ); + Ok(DeviceAuthorization { + user_code: parsed.user_code, + device_code: parsed.device_code, + verification_uri: parsed.verification_uri.filter(|u| !u.is_empty()), + verification_uri_complete, + expires_in: parsed.expires_in.filter(|&e| e > 0), + interval: parsed.interval.unwrap_or(5), + }) +} + +/// One poll of `POST {auth_host}{token_path}` with the device grant. +pub(crate) async fn poll_device_token( + cfg: &OAuthConfig, + device_code: &str, +) -> anyhow::Result { + let url = oauth_url(cfg.auth_host, cfg.token_path); + let resp = crate::http::shared_client() + .post(&url) + .header("Accept", "application/json") + .form(&[ + ("client_id", cfg.client_id), + ("device_code", device_code), + ("grant_type", DEVICE_GRANT_TYPE), + ]) + .send() + .await + .map_err(|e| anyhow::anyhow!("Token polling request failed: {e}"))?; + + let status = resp.status(); + if status.is_server_error() { + anyhow::bail!("Token polling server error: {status}"); + } + let body = resp.bytes().await?; + if status.is_success() { + if let Ok(tokens) = serde_json::from_slice::(&body) { + tracing::info!("auth: device poll succeeded, access token issued (generic oauth)"); + return Ok(DevicePollResult::Success(Box::new(tokens.into_auth()))); + } + tracing::warn!( + "auth: device poll returned 200 without access_token; continuing (generic oauth)" + ); + return Ok(DevicePollResult::Pending { + error: "missing_access_token".to_owned(), + description: None, + }); + } + let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default(); + let error = err.error.unwrap_or_else(|| "unknown_error".to_owned()); + if error == "expired_token" { + tracing::info!( + "auth: device code expired; restarting device authorization (generic oauth)" + ); + return Ok(DevicePollResult::Expired); + } + tracing::debug!(error = %error, "auth: device poll pending (generic oauth)"); + Ok(DevicePollResult::Pending { + error, + description: err.error_description, + }) +} + +/// `POST {auth_host}{token_path}` with `grant_type=refresh_token`. Retries the +/// retryable statuses / network errors with exponential backoff; 401/403 +/// returns immediately as [`RefreshError::Unauthorized`]. +pub(crate) async fn refresh_token( + cfg: &OAuthConfig, + refresh_token: &str, +) -> Result { + let url = oauth_url(cfg.auth_host, cfg.token_path); + let mut last_error = String::from("no attempt made"); + for attempt in 0..MAX_REFRESH_RETRIES { + if attempt > 0 { + let backoff = std::time::Duration::from_secs(1 << (attempt - 1)); + tracing::warn!( + attempt, + backoff_secs = backoff.as_secs(), + last_error = %last_error, + "auth: retrying token refresh (generic oauth)" + ); + tokio::time::sleep(backoff).await; + } + tracing::info!(attempt, "auth: token refresh attempt (generic oauth)"); + let send_result = crate::http::shared_client() + .post(&url) + .header("Accept", "application/json") + .form(&[ + ("client_id", cfg.client_id), + ("grant_type", REFRESH_GRANT_TYPE), + ("refresh_token", refresh_token), + ]) + .send() + .await; + + let resp = match send_result { + Ok(resp) => resp, + Err(e) => { + last_error = format!("network error: {e}"); + continue; + } + }; + let status = resp.status().as_u16(); + let body = resp.bytes().await.unwrap_or_default(); + if status == 401 || status == 403 { + let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default(); + return Err(RefreshError::Unauthorized { + status, + description: err + .error_description + .unwrap_or_else(|| "Token refresh unauthorized.".to_owned()), + }); + } + if status == 200 { + return match serde_json::from_slice::(&body) { + Ok(tokens) => Ok(tokens.into_auth()), + Err(e) => Err(RefreshError::Fatal { + status, + description: format!("malformed token payload: {e}"), + }), + }; + } + let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default(); + let description = err + .error_description + .unwrap_or_else(|| format!("Token refresh failed (HTTP {status}).")); + if RETRYABLE_REFRESH_STATUSES.contains(&status) { + last_error = description; + continue; + } + return Err(RefreshError::Fatal { + status, + description, + }); + } + Err(RefreshError::Exhausted { last_error }) +} + +#[cfg(test)] +mod tests { + use super::*; + use kigi_models::XAI_OAUTH_CONFIG; + use wiremock::matchers::{body_string_contains, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + /// An OAuthConfig pointed at a mock server (copies XAI's client_id/scope/ + /// paths but overrides the host). + fn mock_cfg(host: &'static str) -> OAuthConfig { + OAuthConfig { + auth_host: host, + ..XAI_OAUTH_CONFIG + } + } + + fn token_json(access: &str, refresh: &str) -> serde_json::Value { + serde_json::json!({ + "access_token": access, + "refresh_token": refresh, + "expires_in": 3600, + "scope": "grok-cli:access", + "token_type": "bearer", + }) + } + + #[tokio::test] + async fn device_authorization_sends_client_scope_and_referrer() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/oauth2/device/code")) + .and(body_string_contains( + "client_id=b1a00492-073a-47ea-816f-4c329264a828", + )) + .and(body_string_contains("scope=openid")) + .and(body_string_contains("referrer=kigi")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "user_code": "GROK-1234", + "device_code": "dev-xai-1", + "verification_uri": "https://x.ai/device", + "verification_uri_complete": "https://x.ai/device?user_code=GROK-1234", + "expires_in": 900, + "interval": 5, + }))) + .expect(1) + .mount(&server) + .await; + let auth = request_device_authorization(&mock_cfg(host)).await.unwrap(); + assert_eq!(auth.user_code, "GROK-1234"); + assert_eq!(auth.device_code, "dev-xai-1"); + assert_eq!( + auth.verification_uri_complete, + "https://x.ai/device?user_code=GROK-1234" + ); + assert_eq!(auth.expires_in, Some(900)); + } + + /// A response with only `verification_uri` (no `_complete`) still yields a + /// valid display URI. + #[tokio::test] + async fn device_authorization_falls_back_to_verification_uri() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/oauth2/device/code")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "user_code": "GROK-9", + "device_code": "d", + "verification_uri": "https://x.ai/device", + }))) + .mount(&server) + .await; + let auth = request_device_authorization(&mock_cfg(host)).await.unwrap(); + assert_eq!(auth.verification_uri_complete, "https://x.ai/device"); + assert_eq!(auth.interval, 5, "default interval"); + } + + #[tokio::test] + async fn poll_success_builds_auth_with_expiry() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/oauth2/token")) + .and(body_string_contains("grant_type=urn")) + .and(body_string_contains("device_code=dev-xai-1")) + .respond_with( + ResponseTemplate::new(200).set_body_json(token_json("grok-at", "grok-rt")), + ) + .expect(1) + .mount(&server) + .await; + let result = poll_device_token(&mock_cfg(host), "dev-xai-1") + .await + .unwrap(); + let DevicePollResult::Success(auth) = result else { + panic!("expected success, got {result:?}"); + }; + assert_eq!(auth.key, "grok-at"); + assert_eq!(auth.refresh_token.as_deref(), Some("grok-rt")); + assert_eq!(auth.expires_in, Some(3600)); + } + + #[tokio::test] + async fn poll_maps_authorization_pending_to_pending() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/oauth2/token")) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(serde_json::json!({ "error": "authorization_pending" })), + ) + .mount(&server) + .await; + let result = poll_device_token(&mock_cfg(host), "dev-xai-1") + .await + .unwrap(); + match result { + DevicePollResult::Pending { error, .. } => assert_eq!(error, "authorization_pending"), + other => panic!("expected pending, got {other:?}"), + } + } + + #[tokio::test] + async fn refresh_success_round_trip() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/oauth2/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=grok-rt-old")) + .respond_with( + ResponseTemplate::new(200).set_body_json(token_json("grok-at-new", "grok-rt-new")), + ) + .expect(1) + .mount(&server) + .await; + let auth = refresh_token(&mock_cfg(host), "grok-rt-old").await.unwrap(); + assert_eq!(auth.key, "grok-at-new"); + assert_eq!(auth.refresh_token.as_deref(), Some("grok-rt-new")); + } + + #[tokio::test] + async fn refresh_401_maps_to_unauthorized() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/oauth2/token")) + .respond_with( + ResponseTemplate::new(401) + .set_body_json(serde_json::json!({ "error_description": "refresh revoked" })), + ) + .expect(1) + .mount(&server) + .await; + let err = refresh_token(&mock_cfg(host), "grok-rt-dead") + .await + .unwrap_err(); + match err { + RefreshError::Unauthorized { + status, + description, + } => { + assert_eq!(status, 401); + assert_eq!(description, "refresh revoked"); + } + other => panic!("expected Unauthorized, got {other:?}"), + } + } +} diff --git a/crates/codegen/kigi-shell/src/auth/oauth_registry.rs b/crates/codegen/kigi-shell/src/auth/oauth_registry.rs new file mode 100644 index 0000000..aa7b8bc --- /dev/null +++ b/crates/codegen/kigi-shell/src/auth/oauth_registry.rs @@ -0,0 +1,257 @@ +//! Process-global per-provider OAuth `AuthManager` pool for INFERENCE-time auth. +//! +//! A session binds its primary (Kimi / first-party) [`AuthManager`] for the +//! subscription path, but a `uses_oauth` platform that carries an +//! [`kigi_models::OAuthConfig`] (xai-grok today) needs its OWN scope-keyed +//! manager for every per-turn decision — bearer resolution, proactive / +//! on-expiry refresh, and 401 recovery. Reusing the Kimi manager for a grok +//! turn would transmit the Kimi subscription bearer to `api.x.ai` (a +//! cross-provider leak, guaranteed 401) and, without proactive refresh, would +//! 401 every turn once the ~1h grok token expired until a process restart. +//! +//! The pool is the SINGLE SOURCE OF TRUTH: one long-lived `AuthManager` per +//! 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). +//! +//! SECURITY: access/refresh tokens and resolved bearers are NEVER logged here. + +use std::collections::HashMap; +use std::path::Path; +use std::sync::{Arc, OnceLock}; + +use parking_lot::Mutex; + +use crate::auth::AuthManager; + +/// Process-wide pool of live per-scope OAuth managers. +/// +/// Auth is process-global (one user), so a single manager per scope is correct +/// and lets the proactive-refresh task start exactly once per scope no matter +/// how many sessions spawn. Keyed by the OAuth `scope_key` (`oauth/xai`, …). +fn oauth_manager_pool() -> &'static Mutex>> { + static POOL: OnceLock>>> = OnceLock::new(); + POOL.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// 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 +/// construction (thereafter kept fresh by the proactive-refresh loop), so a +/// grok login that lands after this scope was first built is adopted on the +/// manager's own refresh tick — no session ever needs re-spawning. +/// +/// MUST be called from within a Tokio runtime (the proactive-refresh loop +/// spawns a task, mirroring the primary). +pub(crate) fn global_manager_for( + kigi_home: &Path, + oauth: &'static kigi_models::OAuthConfig, +) -> Arc { + let mut pool = oauth_manager_pool().lock(); + if let Some(existing) = pool.get(oauth.scope_key) { + return existing.clone(); + } + let manager = Arc::new(AuthManager::new_oauth_provider(kigi_home, oauth)); + manager.configure_refresher(); + // Never-cancelled token = process-lifetime, matching the api-server / + // per-session eager-refresh sites that pass a fresh token. + manager.start_proactive_refresh(tokio_util::sync::CancellationToken::new()); + pool.insert(oauth.scope_key, manager.clone()); + manager +} + +/// The `AuthManager` that governs INFERENCE auth for `managed_key` +/// (`{platform}/{model}`, e.g. `xai-grok/grok-4-latest`). +/// +/// A generic device-code OAuth platform routes to ITS OWN scope-keyed manager +/// from the process-global pool ([`global_manager_for`], built on demand from +/// the on-disk token); every other key (Kimi, API-key platforms, `[model.*]` +/// entries, or an unprefixed bare id) routes to `primary`. +/// +/// The pool is the single source of truth — there is no per-session snapshot to +/// freeze at spawn, so a grok login that happens AFTER a session spawned is +/// resolved correctly on the next grok turn. A grok key NEVER resolves to +/// `primary`: even before the user logs into grok the pooled manager simply +/// holds no token (its bearer / api_key is then `None`), so the Kimi +/// subscription bearer can never reach a third-party host — fail-fast, never a +/// silent fallback to the Kimi manager. +pub(crate) fn manager_for_model( + kigi_home: &Path, + managed_key: &str, + primary: Option<&Arc>, +) -> Option> { + if let Some((platform, _)) = kigi_models::parse_managed_model_key(managed_key) + && let Some(oauth) = platform.oauth() + { + return Some(global_manager_for(kigi_home, oauth)); + } + primary.cloned() +} + +/// The SESSION token (the raw bearer/key string) that 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. +/// +/// 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, + primary: Option<&Arc>, +) -> Option { + manager_for_model(kigi_home, managed_key, primary) + .and_then(|am| am.current_or_expired()) + .map(|a| a.key) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::auth::KimiCodeConfig; + use crate::auth::{AuthMode, KimiAuth}; + + /// A Kimi manager holding a fixed in-memory bearer, standing in for a + /// session's primary. The `TempDir` is returned so the caller keeps it + /// alive; the token is read from memory (`current_or_expired`), so disk + /// contents are irrelevant to the assertion. + fn primary_with_token(key: &str) -> (tempfile::TempDir, Arc) { + let dir = tempfile::tempdir().unwrap(); + let manager = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); + manager.hot_swap(KimiAuth { + key: key.to_string(), + auth_mode: AuthMode::OAuth, + ..KimiAuth::test_default() + }); + (dir, manager) + } + + fn xai_oauth() -> &'static kigi_models::OAuthConfig { + kigi_models::PlatformId::XaiGrok + .oauth() + .expect("xai-grok carries an OAuthConfig") + } + + /// A non-OAuth managed key (moonshot-cn/…) and an unprefixed bare id both + /// route to the primary Kimi manager — the Kimi / first-party path is + /// untouched and never consults the pool (no runtime needed). + #[test] + fn non_oauth_and_bare_models_route_to_primary() { + let (_kd, kimi) = primary_with_token("kimi-tok"); + let home = tempfile::tempdir().unwrap(); + for key in ["moonshot-cn/kimi-k2", "kimi-k2-0905-preview"] { + let resolved = manager_for_model(home.path(), key, Some(&kimi)) + .expect("non-oauth key routes to the primary"); + assert!( + Arc::ptr_eq(&resolved, &kimi), + "{key} must resolve to the primary manager" + ); + assert_eq!(resolved.current_or_expired().unwrap().key, "kimi-tok"); + } + } + + /// The primary being `None` (test / BYOK sessions) still yields `None` for a + /// non-oauth key, never a panic — and without touching the pool. + #[test] + fn none_primary_is_passed_through_for_non_oauth() { + let home = tempfile::tempdir().unwrap(); + assert!(manager_for_model(home.path(), "kimi-k2", None).is_none()); + } + + /// An `xai-grok/` turn resolves to the process-global pooled xai + /// manager, NEVER the primary Kimi manager — the pool is the single source. + #[tokio::test] + async fn grok_model_resolves_to_pooled_xai_manager_not_kimi() { + let (_kd, kimi) = primary_with_token("kimi-tok"); + let home = tempfile::tempdir().unwrap(); + let resolved = manager_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi)) + .expect("grok model resolves to the pooled xai manager"); + assert!( + !Arc::ptr_eq(&resolved, &kimi), + "grok model must NOT resolve to the Kimi manager" + ); + assert!( + Arc::ptr_eq(&resolved, &global_manager_for(home.path(), xai_oauth())), + "grok model must resolve to the process-global pooled xai manager" + ); + } + + /// Facet B guard: the resolver routes purely by the model's platform, with + /// no auth-method input — so even when the session's primary is a Kimi + /// (session) manager holding "kimi-tok", a grok model never resolves that + /// Kimi token. + #[tokio::test] + async fn grok_model_under_kimi_primary_never_yields_kimi_token() { + let (_kd, kimi) = primary_with_token("kimi-tok"); + let home = tempfile::tempdir().unwrap(); + let resolved = manager_for_model(home.path(), "xai-grok/grok-4-fast", Some(&kimi)) + .expect("grok model resolves to its own pooled manager regardless of primary"); + assert!(!Arc::ptr_eq(&resolved, &kimi)); + assert_ne!( + resolved.current_or_expired().map(|a| a.key), + Some("kimi-tok".to_string()), + "the Kimi bearer must never be what a grok turn resolves" + ); + } + + /// Fail-fast: a grok key resolves to the pooled xai manager (never the Kimi + /// primary) even with no stored grok session in the pool — the pooled + /// manager then simply holds no token, so nothing (least of all the Kimi + /// bearer) is sent to api.x.ai. + #[tokio::test] + async fn grok_never_falls_back_to_kimi_primary() { + let (_kd, kimi) = primary_with_token("kimi-tok"); + let home = tempfile::tempdir().unwrap(); + let resolved = manager_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi)) + .expect("grok routes to the pooled xai manager, not None"); + assert!( + !Arc::ptr_eq(&resolved, &kimi), + "an OAuth platform must never fall back to the primary Kimi manager" + ); + } + + /// `session_key_for_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). + #[test] + fn session_key_for_non_oauth_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_model(home.path(), key, Some(&kimi)), + Some("kimi-tok".to_string()), + "{key} (non-oauth) must yield the primary token unchanged" + ); + } + } + + /// LEAK guard (aux-model + subagent-override token routing): a grok key with + /// a Kimi primary NEVER yields the primary Kimi token — it draws from the + /// pooled xai manager (its own token, or `None`). This is the exact source + /// the aux `session_key` and the override `session_key` now use. + #[tokio::test] + async fn session_key_for_grok_is_never_the_kimi_primary() { + let (_kd, kimi) = primary_with_token("kimi-tok"); + let home = tempfile::tempdir().unwrap(); + assert_ne!( + session_key_for_model(home.path(), "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), + Some("kimi-tok".to_string()), + ); + } +} diff --git a/crates/codegen/kigi-shell/src/auth/refresh/generic_refresher.rs b/crates/codegen/kigi-shell/src/auth/refresh/generic_refresher.rs new file mode 100644 index 0000000..eff9be0 --- /dev/null +++ b/crates/codegen/kigi-shell/src/auth/refresh/generic_refresher.rs @@ -0,0 +1,282 @@ +//! Generic device-code token refresher: drives `POST {token_path}` with +//! `grant_type=refresh_token` for any [`kigi_models::OAuthConfig`] provider +//! (xai-grok today) through the [`TokenRefresher`] seam. +//! +//! Structurally identical to [`super::kimi_refresher::KimiRefresher`] — same +//! sibling-adoption + post-401 grace — but the wire call goes through +//! [`crate::auth::oauth_device`] (no X-Msh headers) instead of the Kimi wire. +//! Access/refresh tokens are NEVER logged. + +use std::sync::Arc; + +use kigi_models::OAuthConfig; + +use crate::auth::error::RefreshTokenFailedReason; +use crate::auth::kimi_oauth::RefreshError; +use crate::auth::manager::RefreshReason; +use crate::auth::oauth_device::{self}; + +use super::{AuthSnapshot, RefreshOutcome, TokenRefresher}; + +/// Grace period after a 401/403 before concluding the refresh token is dead: +/// a concurrent instance may still be persisting its rotated token. +const POST_UNAUTHORIZED_GRACE: std::time::Duration = std::time::Duration::from_secs(1); + +pub(crate) struct GenericDeviceRefresher { + auth: Arc, + cfg: &'static OAuthConfig, +} + +impl GenericDeviceRefresher { + pub(crate) fn new(auth: Arc, cfg: &'static OAuthConfig) -> Self { + Self { auth, cfg } + } + + /// Post-401 sibling check: wait a beat, re-read the persisted credential, + /// and adopt it when its refresh token differs from the rejected one. + async fn adopt_rotation_after_unauthorized(&self, tried_rt: &str) -> Option { + tokio::time::sleep(POST_UNAUTHORIZED_GRACE).await; + let latest = self.auth.read_disk_auth()?; + let latest_rt = latest.refresh_token.as_deref()?; + if latest_rt == tried_rt { + return None; + } + kigi_log::unified_log::info( + "auth.refresh.adopted_rotation_after_401", + None, + Some(serde_json::json!({ + "scope_key": self.cfg.scope_key, + "adopted_rt_prefix": crate::auth::token_suffix(latest_rt), + "rejected_rt_prefix": crate::auth::token_suffix(tried_rt), + })), + ); + Some(RefreshOutcome::success(latest)) + } +} + +#[async_trait::async_trait] +impl TokenRefresher for GenericDeviceRefresher { + async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome { + tracing::info!( + ?reason, + scope_key = self.cfg.scope_key, + "auth: generic refresh attempt" + ); + + let disk_auth = self.auth.read_disk_auth(); + + // Sibling short-circuit: a valid persisted token whose key differs from + // in-memory means another process refreshed already — adopt directly. + if let Some(ref d) = disk_auth + && !crate::auth::is_expired(d) + && self.auth.current().map(|a| a.key).as_deref() != Some(&d.key) + { + kigi_log::unified_log::info( + "auth.refresh.adopted_sibling_token", + None, + Some(serde_json::json!({ + "scope_key": self.cfg.scope_key, + "disk_key_prefix": crate::auth::token_suffix(&d.key), + })), + ); + return RefreshOutcome::success(d.clone()); + } + + let Some(auth) = super::resolve_refresh_credential(self.auth.as_ref(), disk_auth, reason) + else { + tracing::warn!( + ?reason, + "auth: no credential available for refresh (generic)" + ); + return RefreshOutcome::transient("no token with refresh_token available"); + }; + let Some(refresh_token) = auth.refresh_token.clone() else { + tracing::warn!( + ?reason, + "auth: resolved credential has no refresh token (generic)" + ); + return RefreshOutcome::transient("credential has no refresh token"); + }; + + tracing::info!( + rt_prefix = crate::auth::token_suffix(&refresh_token), + expires_at = ?auth.expires_at, + "auth: sending refresh_token grant (generic oauth)" + ); + + match oauth_device::refresh_token(self.cfg, &refresh_token).await { + Ok(new_auth) => { + kigi_log::unified_log::info( + "auth.refresh.token_rotated", + None, + Some(serde_json::json!({ + "scope_key": self.cfg.scope_key, + "new_key_prefix": crate::auth::token_suffix(&new_auth.key), + "expires_at": new_auth.expires_at.map(|e| e.to_rfc3339()), + })), + ); + RefreshOutcome::success(new_auth) + } + Err(RefreshError::Unauthorized { + status, + description, + }) => { + tracing::warn!(status, %description, "auth: refresh token rejected (generic)"); + if let Some(adopted) = self.adopt_rotation_after_unauthorized(&refresh_token).await + { + return adopted; + } + kigi_log::unified_log::warn( + "auth.refresh.unauthorized", + None, + Some(serde_json::json!({ + "scope_key": self.cfg.scope_key, + "status": status, + "description": description, + "rt_prefix": crate::auth::token_suffix(&refresh_token), + })), + ); + RefreshOutcome::permanent( + RefreshTokenFailedReason::RefreshTokenRejected, + Some(refresh_token), + ) + } + Err( + e @ (RefreshError::Exhausted { .. } + | RefreshError::Fatal { .. } + | RefreshError::Local(_)), + ) => { + tracing::warn!(error = %e, "auth: refresh attempt failed (transient, generic)"); + kigi_log::unified_log::warn( + "auth.refresh.transient_wire_failure", + None, + Some(serde_json::json!({ + "scope_key": self.cfg.scope_key, + "error": format!("{e}"), + })), + ); + RefreshOutcome::transient(format!("token refresh failed: {e}")) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::auth::model::KimiAuth; + use chrono::{Duration, Utc}; + use kigi_models::XAI_OAUTH_CONFIG; + use parking_lot::Mutex; + use wiremock::matchers::{body_string_contains, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + struct FakeSnapshot { + current: Mutex>, + disk: Mutex>, + } + impl FakeSnapshot { + fn new(current: Option, disk: Option) -> Arc { + Arc::new(Self { + current: Mutex::new(current), + disk: Mutex::new(disk), + }) + } + } + impl AuthSnapshot for FakeSnapshot { + fn current(&self) -> Option { + self.current + .lock() + .clone() + .filter(|a| !crate::auth::is_expired(a)) + } + fn expired_auth(&self) -> Option { + self.current.lock().clone().filter(crate::auth::is_expired) + } + fn read_disk_auth(&self) -> Option { + self.disk.lock().clone() + } + fn is_expired(&self) -> bool { + self.current + .lock() + .as_ref() + .is_some_and(crate::auth::is_expired) + } + } + + fn expired_session(key: &str, rt: &str) -> KimiAuth { + KimiAuth { + key: key.into(), + refresh_token: Some(rt.into()), + expires_at: Some(Utc::now() - Duration::hours(1)), + expires_in: Some(3600), + ..KimiAuth::test_default() + } + } + + fn mock_cfg(host: &'static str) -> OAuthConfig { + OAuthConfig { + auth_host: host, + ..XAI_OAUTH_CONFIG + } + } + + /// A successful refresh rotates the token via the generic wire. + #[tokio::test] + async fn refresh_success_returns_rotated_token() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/oauth2/token")) + .and(body_string_contains("refresh_token=grok-rt-old")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "grok-at-new", + "refresh_token": "grok-rt-new", + "expires_in": 3600, + }))) + .expect(1) + .mount(&server) + .await; + let stale = expired_session("grok-at-old", "grok-rt-old"); + let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale)); + let cfg: &'static OAuthConfig = Box::leak(Box::new(mock_cfg(host))); + let refresher = GenericDeviceRefresher::new(snap, cfg); + let outcome = refresher.refresh(RefreshReason::PreRequest).await; + let RefreshOutcome::Success(new_auth) = outcome else { + panic!("expected success, got {outcome:?}"); + }; + assert_eq!(new_auth.key, "grok-at-new"); + assert_eq!(new_auth.refresh_token.as_deref(), Some("grok-rt-new")); + } + + /// A 401 on refresh (with no sibling rotation) tombstones the rejected + /// refresh token as a permanent failure. + #[tokio::test] + async fn unauthorized_is_permanent_failure() { + let server = MockServer::start().await; + let host: &'static str = Box::leak(server.uri().into_boxed_str()); + Mock::given(method("POST")) + .and(path("/oauth2/token")) + .respond_with( + ResponseTemplate::new(401) + .set_body_json(serde_json::json!({ "error_description": "revoked" })), + ) + .expect(1) + .mount(&server) + .await; + let stale = expired_session("grok-at-old", "grok-rt-dead"); + let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale)); + let cfg: &'static OAuthConfig = Box::leak(Box::new(mock_cfg(host))); + let refresher = GenericDeviceRefresher::new(snap, cfg); + let outcome = refresher.refresh(RefreshReason::PreRequest).await; + let RefreshOutcome::PermanentFailure { + error, + rejected_refresh_token, + } = outcome + else { + panic!("expected permanent failure, got {outcome:?}"); + }; + assert_eq!(error.reason, RefreshTokenFailedReason::RefreshTokenRejected); + assert_eq!(rejected_refresh_token.as_deref(), Some("grok-rt-dead")); + } +} diff --git a/crates/codegen/kigi-shell/src/auth/refresh/mod.rs b/crates/codegen/kigi-shell/src/auth/refresh/mod.rs index 26f9e1c..330fee3 100644 --- a/crates/codegen/kigi-shell/src/auth/refresh/mod.rs +++ b/crates/codegen/kigi-shell/src/auth/refresh/mod.rs @@ -1,3 +1,4 @@ +mod generic_refresher; mod kimi_refresher; use std::sync::Arc; @@ -6,6 +7,7 @@ use crate::auth::manager::AuthManager; pub(crate) use crate::auth::manager::RefreshReason; use crate::auth::model::KimiAuth; +pub(crate) use generic_refresher::GenericDeviceRefresher; pub(crate) use kimi_refresher::KimiRefresher; /// Read-only view of `AuthManager` for refreshers. Enforces the @@ -120,8 +122,20 @@ pub(crate) trait TokenRefresher: Send + Sync { async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome; } -/// Build the production refresher against `kigi_env::oauth_host()`. +/// Build the production refresher for this manager's scope. A scope that maps +/// to a generic device-code [`kigi_models::OAuthConfig`] (xai-grok) gets the +/// [`GenericDeviceRefresher`]; every other scope — Kimi Code, whose registry +/// `oauth` field is `None` by design — gets the bespoke [`KimiRefresher`] +/// against `kigi_env::oauth_host()`. pub(crate) fn build_refresher(auth_manager: Arc) -> Arc { - let snapshot: Arc = auth_manager; - Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host())) + match kigi_models::oauth_config_for_scope_key(auth_manager.scope()) { + Some(cfg) => { + let snapshot: Arc = auth_manager; + Arc::new(GenericDeviceRefresher::new(snapshot, cfg)) + } + None => { + let snapshot: Arc = auth_manager; + Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host())) + } + } } diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs index 20d2b80..8523d28 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs @@ -60,9 +60,12 @@ impl SessionActor { stream_tool_calls: Some(sampling_config.stream_tool_calls), }); let existing = self.chat_state_handle.get_credentials().await; + // Read the session bearer from the switched-to model's OWN manager: a + // grok model reads the xai-grok token (used only to classify the + // credential's auth_type here), never the Kimi one. Kimi / non-oauth + // models resolve to the primary — byte-identical. let session_key = self - .auth_manager - .as_ref() + .auth_manager_for_model(&sampling_config.model) .and_then(|am| am.current_or_expired().map(|a| a.key)); self.chat_state_handle .update_credentials(kigi_chat_state::Credentials { diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs index 74d1f2a..b7351b8 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs @@ -624,12 +624,25 @@ impl SessionActor { let resolved_describe = self .resolve_aux_sampler_config(&self.image_description_model) .await; - let (describe_model, sampler_config) = + // LEAK 1b: only re-point the aux bearer_resolver when the aux model + // actually resolved (Some) — the `None` fallback yields the SESSION + // config, whose Kimi resolver must stay as-is. + let aux_resolved = resolved_describe.is_some(); + let (describe_model, mut sampler_config) = crate::agent::config::finalize_image_describe_sampler_config( resolved_describe, &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. + if aux_resolved { + self.repoint_aux_bearer_resolver_for_oauth( + &mut sampler_config, + &self.image_description_model, + ); + } let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| { acp::Error::internal_error().data(format!( "failed to build image-describe sampling client: {e}" diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs index b2a4b4c..60cc8ff 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs @@ -103,6 +103,30 @@ where result } } +/// Wraps an [`AuthManager`](crate::auth::AuthManager) as a sampler +/// [`BearerResolver`](kigi_sampler::BearerResolver), resolving the live +/// (current-or-expired) bearer at request time. Shared by +/// [`SessionActor::reconstruct_full_config`] (the session model) and the +/// aux-model bearer repoint ([`SessionActor::repoint_aux_bearer_resolver_for_oauth`]) +/// so both wrap ONE definition. SECURITY: the bearer is resolved per request +/// and never logged. +pub(crate) struct AuthManagerBearerResolver(pub(crate) std::sync::Arc); +impl std::fmt::Debug for AuthManagerBearerResolver { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AuthManagerBearerResolver").finish() + } +} +impl kigi_sampler::BearerResolver for AuthManagerBearerResolver { + fn current_bearer(&self) -> Option { + self.0.current_or_expired().map(|a| a.key) + } +} +/// Wrap `am` as a shared sampler bearer resolver. +fn auth_manager_bearer_resolver( + am: std::sync::Arc, +) -> kigi_sampler::SharedBearerResolver { + std::sync::Arc::new(AuthManagerBearerResolver(am)) +} impl SessionActor { pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec, u64) { let mcp_wait_start = std::time::Instant::now(); @@ -184,6 +208,76 @@ impl SessionActor { let auth_method = self.auth_method_id.load(); SessionTokenAuthGate::new(auth_method.as_deref(), byok, base_url) } + /// The [`AuthManager`](crate::auth::AuthManager) that governs INFERENCE auth + /// for the model whose routing slug is `model` (the sampling config's + /// `model`). A generic device-code OAuth platform (xai-grok) routes to its + /// OWN scope-keyed manager from the process-global OAuth pool + /// ([`crate::auth::oauth_registry::manager_for_model`], built on demand from + /// the on-disk token); every other model routes to the primary Kimi + /// `auth_manager`. + /// + /// This is the single chokepoint that keeps a grok turn from ever sending + /// the Kimi bearer (Facet B) and gives it its own proactive-refresh + 401 + /// recovery manager (Facet A). The pool is the single source of truth, so a + /// grok login that lands AFTER this session spawned is resolved on the next + /// grok turn (no frozen per-session snapshot). Cheap `Arc` clone. Kimi / + /// first-party path is byte-identical: a non-oauth model always resolves to + /// the primary. + /// + /// `None` when there is no governing manager: a BYOK / test session with no + /// primary. A grok model always resolves to its pooled manager (never the + /// Kimi primary); when the user has not logged into grok that manager simply + /// holds no token, so no Kimi bearer can leak. + pub(super) fn auth_manager_for_model( + &self, + model: &str, + ) -> Option> { + // `model` is the bare routing slug; recover the managed catalog key + // (`{platform}/{model}`) so the platform — and thus its OAuth scope — is + // unambiguous. A bare / config / unlisted model yields no managed key + // and resolves to the primary. + let managed_key = self.managed_key_for_model(model); + crate::auth::oauth_registry::manager_for_model( + &crate::util::kigi_home::kigi_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. + fn managed_key_for_model(&self, model: &str) -> Option { + let models = self.models_manager.models(); + crate::agent::config::find_model_by_id(&models, model).and_then(|e| e.info().id.clone()) + } + /// 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() + } + /// 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( + &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)); + } + } /// Emit a unified-log breadcrumb whenever the session-token refresh gate is /// evaluated with an **`Unknown`** per-model BYOK status on a session-based /// method — the condition that (pre-fix) silently demoted live sessions to @@ -237,18 +331,6 @@ impl SessionActor { } } } - #[allow(clippy::items_after_statements)] - struct AuthManagerBearerResolver(std::sync::Arc); - impl std::fmt::Debug for AuthManagerBearerResolver { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AuthManagerBearerResolver").finish() - } - } - impl kigi_sampler::BearerResolver for AuthManagerBearerResolver { - fn current_bearer(&self) -> Option { - self.0.current_or_expired().map(|a| a.key) - } - } let cfg = self .chat_state_handle .get_sampling_config() @@ -273,6 +355,16 @@ impl SessionActor { SessionTokenAuthGate::new(auth_method.as_deref(), model_facts.byok, &cfg.base_url); 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 + // wraps the xai-grok manager, never the Kimi one (captured before + // `cfg.model` is moved into the struct below). `None` when the gate is + // inactive or the oauth provider has no manager (fail-fast, no Kimi + // fallback). + let inference_auth_manager = if use_bearer_resolver { + self.auth_manager_for_model(&cfg.model) + } else { + None + }; let auth_scheme = model_facts.auth_scheme; let mut extra_headers = cfg.extra_headers; crate::agent::config::inject_url_derived_headers( @@ -323,15 +415,7 @@ impl SessionActor { idle_timeout_secs: None, origin_client: self.origin_client.clone(), attribution_callback: self.attribution_callback.clone(), - bearer_resolver: if use_bearer_resolver { - self.auth_manager - .as_ref() - .map(|am| -> kigi_sampler::SharedBearerResolver { - std::sync::Arc::new(AuthManagerBearerResolver(am.clone())) - }) - } else { - None - }, + bearer_resolver: inference_auth_manager.map(auth_manager_bearer_resolver), supports_backend_search: self.supports_backend_search.get(), compactions_remaining: self.compactions_remaining.get(), compaction_at_tokens: self.compaction_at_tokens.get(), @@ -461,10 +545,16 @@ impl SessionActor { slug: &str, ) -> Option { let creds = self.chat_state_handle.get_credentials().await; - let session_key = self - .auth_manager - .as_ref() - .and_then(|am| am.current_or_expired().map(|a| a.key.clone())); + // 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(), + 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( @@ -491,6 +581,10 @@ 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); let model = cfg.model.clone(); let client = kigi_sampler::SamplingClient::new(cfg) .map_err(|e| { @@ -661,29 +755,40 @@ impl SessionActor { )), ); } - if auth_recovery_eligible && let Some(ref am) = self.auth_manager { - if am.try_recover_unauthorized().await { - tracing::info!( + // Recover via the ACTIVE model's OWN manager: a grok 401 recovers the + // xai-grok session via the xai-grok manager, never the Kimi one. For a + // Kimi / non-oauth model this resolves to the primary — byte-identical. + if auth_recovery_eligible { + let recovery_model = self + .chat_state_handle + .get_sampling_config() + .await + .map(|c| c.model) + .unwrap_or_default(); + if let Some(am) = self.auth_manager_for_model(&recovery_model) { + if am.try_recover_unauthorized().await { + tracing::info!( + session_id = % self.session_info.id.0, + "auth recovery: sampler 401, recovered, retrying" + ); + kigi_log::unified_log::info( + "auth recovery: sampler 401, recovered, retrying", + Some(self.session_info.id.0.as_ref()), + None, + ); + self.prepare_sampler_for_turn().await; + return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit); + } + tracing::warn!( session_id = % self.session_info.id.0, - "auth recovery: sampler 401, recovered, retrying" + "auth recovery: sampler 401, refresh failed" ); - kigi_log::unified_log::info( - "auth recovery: sampler 401, recovered, retrying", + kigi_log::unified_log::warn( + "auth recovery: sampler 401, refresh failed", Some(self.session_info.id.0.as_ref()), None, ); - self.prepare_sampler_for_turn().await; - return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit); } - tracing::warn!( - session_id = % self.session_info.id.0, - "auth recovery: sampler 401, refresh failed" - ); - kigi_log::unified_log::warn( - "auth recovery: sampler 401, refresh failed", - Some(self.session_info.id.0.as_ref()), - None, - ); } if matches!(error.kind, SamplingErrorKind::IdleTimeout) { self.signals_handle().record_idle_timeout(); @@ -846,14 +951,17 @@ impl SessionActor { } /// Proactively refresh the auth token if near expiry. pub(super) async fn refresh_token_if_expired(&self) { - if let Some(ref am) = self.auth_manager { + let (model_id, base_url) = self + .chat_state_handle + .get_sampling_config() + .await + .map(|c| (c.model, c.base_url)) + .unwrap_or_default(); + // Refresh the ACTIVE model's OWN manager: a grok model refreshes the + // xai-grok token via the xai-grok manager, never the Kimi one. For a + // Kimi / non-oauth model this resolves to the primary — byte-identical. + if let Some(am) = self.auth_manager_for_model(&model_id) { let creds = self.chat_state_handle.get_credentials().await; - let (model_id, base_url) = self - .chat_state_handle - .get_sampling_config() - .await - .map(|c| (c.model, c.base_url)) - .unwrap_or_default(); if self.auth_gate(&model_id, &base_url).active() && let Ok(key) = am.get_valid_token().await { @@ -975,3 +1083,47 @@ impl SessionActor { .push_assistant_response(assistant_item); } } +#[cfg(test)] +mod bearer_resolver_tests { + use super::AuthManagerBearerResolver; + use kigi_sampler::BearerResolver; + + /// LEAK 1b: the shared `AuthManagerBearerResolver` resolves the LIVE bearer + /// of the manager it wraps. So an aux bearer_resolver built over grok's OWN + /// (oauth) pooled manager yields grok's token (or `None`) — NEVER the Kimi + /// session token that a Kimi-manager resolver would. The + /// `repoint_aux_bearer_resolver_for_oauth` fix wraps exactly this grok + /// manager for a grok aux model. + #[tokio::test] + async fn resolver_resolves_the_wrapped_manager_never_kimi() { + let dir = tempfile::tempdir().unwrap(); + let kimi = std::sync::Arc::new(crate::auth::AuthManager::new( + dir.path(), + crate::auth::KimiCodeConfig::default(), + )); + kimi.hot_swap(crate::auth::KimiAuth { + key: "kimi-tok".to_string(), + auth_mode: crate::auth::AuthMode::OAuth, + ..crate::auth::KimiAuth::test_default() + }); + // The Kimi-manager resolver yields the Kimi bearer. + assert_eq!( + AuthManagerBearerResolver(kimi.clone()).current_bearer(), + Some("kimi-tok".to_string()), + ); + // The grok (oauth) pooled manager is distinct — its resolver never + // yields the Kimi bearer (grok's own token, or None). + let oauth = kigi_models::PlatformId::XaiGrok + .oauth() + .expect("xai-grok carries an OAuthConfig"); + let grok = crate::auth::oauth_registry::global_manager_for( + &crate::util::kigi_home::kigi_home(), + oauth, + ); + assert_ne!( + AuthManagerBearerResolver(grok).current_bearer(), + Some("kimi-tok".to_string()), + "a grok aux bearer_resolver must never resolve the Kimi session token", + ); + } +} diff --git a/crates/codegen/kigi-tui/src/app/app_view.rs b/crates/codegen/kigi-tui/src/app/app_view.rs index 8985834..6ae6233 100644 --- a/crates/codegen/kigi-tui/src/app/app_view.rs +++ b/crates/codegen/kigi-tui/src/app/app_view.rs @@ -6894,28 +6894,33 @@ pub(crate) mod tests { #[test] fn pending_menu_items_lists_interactive_methods_plus_quit() { let items = pending_menu_items(&fresh_user_auth_methods(), None); - assert_eq!(items.len(), 26, "25 login rows + Quit, got {items:?}"); + assert_eq!(items.len(), 27, "26 login rows + Quit, got {items:?}"); assert!( matches!(&items[0], PendingMenuItem::Login { label } if label == "Kimi Code (OAuth)"), "row 0 must be the OAuth login, got {:?}", items[0] ); + assert!( + matches!(&items[1], PendingMenuItem::Login { label } if label == "xAI Grok (subscription) (OAuth)"), + "row 1 must be the xai-grok OAuth login (interactive, after kimi-code), got {:?}", + items[1] + ); assert_eq!( - items[1], + items[2], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::MoonshotCn), label: "Moonshot Open Platform (API key \u{b7} moonshot.cn)".into(), } ); assert_eq!( - items[2], + items[3], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::MoonshotAi), label: "Moonshot Open Platform (API key \u{b7} moonshot.ai)".into(), } ); assert_eq!( - items[3], + items[4], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::OpenAi), label: "OpenAI (API key)".into(), @@ -6923,153 +6928,153 @@ pub(crate) mod tests { "new registry rows must appear in the picker with zero TUI changes" ); assert_eq!( - items[4], + items[5], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Anthropic), label: "Anthropic (API key)".into(), } ); assert_eq!( - items[5], + items[6], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::DeepSeek), label: "DeepSeek (API key)".into(), } ); assert_eq!( - items[6], + items[7], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Groq), label: "Groq (API key)".into(), } ); assert_eq!( - items[7], + items[8], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Mistral), label: "Mistral (API key)".into(), } ); assert_eq!( - items[8], + items[9], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Fireworks), label: "Fireworks AI (API key)".into(), } ); assert_eq!( - items[9], + items[10], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Google), label: "Google Gemini (API key)".into(), } ); assert_eq!( - items[10], + items[11], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::OpenRouter), label: "OpenRouter (API key)".into(), } ); assert_eq!( - items[11], + items[12], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Together), label: "Together AI (API key)".into(), } ); assert_eq!( - items[12], + items[13], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Cerebras), label: "Cerebras (API key)".into(), } ); assert_eq!( - items[13], + items[14], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Nvidia), label: "NVIDIA NIM (API key)".into(), } ); assert_eq!( - items[14], + items[15], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Vercel), label: "Vercel AI Gateway (API key)".into(), } ); assert_eq!( - items[15], + items[16], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Xai), label: "xAI (Grok) (API key)".into(), } ); assert_eq!( - items[16], + items[17], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::QwenTokenPlan), label: "Qwen Token Plan (API key)".into(), } ); assert_eq!( - items[17], + items[18], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::QwenTokenPlanCn), label: "Qwen Token Plan China (API key)".into(), } ); assert_eq!( - items[18], + items[19], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::KimiCoding), label: "Kimi For Coding (API key)".into(), } ); assert_eq!( - items[19], + items[20], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Zai), label: "Z.AI (API key)".into(), } ); assert_eq!( - items[20], + items[21], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::ZaiCodingCn), label: "Z.AI Coding China (API key)".into(), } ); assert_eq!( - items[21], + items[22], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Xiaomi), label: "Xiaomi MiMo (API key)".into(), } ); assert_eq!( - items[22], + items[23], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::XiaomiTokenPlanCn), label: "Xiaomi Token Plan China (API key)".into(), } ); assert_eq!( - items[23], + items[24], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::Minimax), label: "MiniMax (API key)".into(), } ); assert_eq!( - items[24], + items[25], PendingMenuItem::ApiKey { target: PlatformLogin(kigi_shell::models::PlatformId::MinimaxCn), label: "MiniMax China (API key)".into(), } ); - assert_eq!(items[25], PendingMenuItem::Quit); + assert_eq!(items[26], PendingMenuItem::Quit); // The non-interactive methods must never appear as rows. let byok = kigi_shell::agent::auth_method::build_auth_methods( kigi_shell::agent::auth_method::AuthMethodsBuildInputs { @@ -7080,8 +7085,8 @@ pub(crate) mod tests { ); assert_eq!( pending_menu_items(&byok.methods, None).len(), - 26, - "xai.api_key / cached_token must not add rows" + 27, + "xai.api_key / cached_token must not add rows (26 login rows + Quit)" ); } /// Startup lands on the picker only when there is a real choice: the @@ -7094,7 +7099,7 @@ pub(crate) mod tests { assert!(!login_picker_has_choice(&kimi_only)); assert!(!login_picker_has_choice(&[])); } - /// With all three methods advertised, arrows+Enter select a Moonshot row + /// With all methods advertised, arrows+Enter select a Moonshot row /// and 'l' keeps selecting the first (OAuth) row. #[test] fn welcome_pending_arrows_select_moonshot_row() { @@ -7102,7 +7107,9 @@ pub(crate) mod tests { app.auth_methods = fresh_user_auth_methods(); app.auth_state = AuthState::Pending { error: None }; app.welcome_prompt_focused = false; - // Down → row 0 (OAuth), Down → row 1 (moonshot-cn). + // Interactive OAuth logins come first: row 0 (kimi-code), row 1 + // (xai-grok); the first API-key row (moonshot-cn) is now row 2. + app.handle_input(&key_event(KeyCode::Down, KeyModifiers::NONE)); app.handle_input(&key_event(KeyCode::Down, KeyModifiers::NONE)); app.handle_input(&key_event(KeyCode::Down, KeyModifiers::NONE)); let outcome = app.handle_input(&key_event(KeyCode::Enter, KeyModifiers::NONE));