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

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

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

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

Registry at 26; picker updated (xai-grok interactive login row); TUI
context-window already auto-updates per model. Full gate green (234 suites,
fmt, clippy -D warnings, deny). GPT/Claude/Grok officially permit third-party
subscription use.
This commit is contained in:
2026-07-22 01:36:29 -04:00
parent 8a26460251
commit dbce6bf305
26 changed files with 2359 additions and 161 deletions
+15
View File
@@ -168,6 +168,21 @@ edges stay deterministic Rust. The harness appends a terminal
are built generically from advertised methods (`AuthMethodKind:: are built generically from advertised methods (`AuthMethodKind::
ApiKeyPlatform`), so new registry rows appear in the picker with no TUI ApiKeyPlatform`), so new registry rows appear in the picker with no TUI
changes. 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 - Model metadata (context window, thinking levels) comes from the provider
wire when served; metadata-poor listings are enriched from models.dev wire when served; metadata-poor listings are enriched from models.dev
(`kigi-models/src/enrichment.rs` — bundled raw snapshot regenerated by (`kigi-models/src/enrichment.rs` — bundled raw snapshot regenerated by
+172 -2
View File
@@ -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 /// One row of the platform registry. All per-platform data lives in these
/// rows; [`PlatformId`] methods only read fields. Adding a platform touches /// rows; [`PlatformId`] methods only read fields. Adding a platform touches
/// exactly four sites in this file: the enum variant, the `ALL` entry, the /// exactly four sites in this file: the enum variant, the `ALL` entry, the
@@ -111,6 +158,11 @@ struct PlatformSpec {
base_url: BaseUrlSource, base_url: BaseUrlSource,
/// True for OAuth-bearer subscription channels. /// True for OAuth-bearer subscription channels.
uses_oauth: bool, 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. /// Model-id prefixes admitted from this platform's `/models` listing.
/// `None` = no filtering (listing served pre-filtered). /// `None` = no filtering (listing served pre-filtered).
allowed_model_prefixes: Option<&'static [&'static str]>, allowed_model_prefixes: Option<&'static [&'static str]>,
@@ -166,6 +218,9 @@ const KIMI_CODE_SPEC: PlatformSpec = PlatformSpec {
display_name: "Kimi Code", display_name: "Kimi Code",
base_url: BaseUrlSource::KigiEnvCoding, base_url: BaseUrlSource::KigiEnvCoding,
uses_oauth: true, 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, allowed_model_prefixes: None,
api_key_envs: &[], api_key_envs: &[],
vendor: "Kimi", vendor: "Kimi",
@@ -190,6 +245,7 @@ const MOONSHOT_CN_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.moonshot.cn/v1", default: "https://api.moonshot.cn/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: Some(&["kimi-k"]), allowed_model_prefixes: Some(&["kimi-k"]),
api_key_envs: &[MOONSHOT_CN_API_KEY_ENV, MOONSHOT_API_KEY_ENV], api_key_envs: &[MOONSHOT_CN_API_KEY_ENV, MOONSHOT_API_KEY_ENV],
vendor: "Moonshot", vendor: "Moonshot",
@@ -214,6 +270,7 @@ const MOONSHOT_AI_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.moonshot.ai/v1", default: "https://api.moonshot.ai/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: Some(&["kimi-k"]), allowed_model_prefixes: Some(&["kimi-k"]),
api_key_envs: &[MOONSHOT_AI_API_KEY_ENV, MOONSHOT_API_KEY_ENV], api_key_envs: &[MOONSHOT_AI_API_KEY_ENV, MOONSHOT_API_KEY_ENV],
vendor: "Moonshot", vendor: "Moonshot",
@@ -241,6 +298,7 @@ const OPENAI_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.openai.com/v1", default: "https://api.openai.com/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["OPENAI_API_KEY"], api_key_envs: &["OPENAI_API_KEY"],
vendor: "OpenAI", vendor: "OpenAI",
@@ -270,6 +328,7 @@ const ANTHROPIC_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.anthropic.com/v1", default: "https://api.anthropic.com/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["ANTHROPIC_API_KEY"], api_key_envs: &["ANTHROPIC_API_KEY"],
vendor: "Anthropic", vendor: "Anthropic",
@@ -301,6 +360,7 @@ const DEEPSEEK_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.deepseek.com", default: "https://api.deepseek.com",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["DEEPSEEK_API_KEY"], api_key_envs: &["DEEPSEEK_API_KEY"],
vendor: "DeepSeek", vendor: "DeepSeek",
@@ -328,6 +388,7 @@ const GROQ_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.groq.com/openai/v1", default: "https://api.groq.com/openai/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["GROQ_API_KEY"], api_key_envs: &["GROQ_API_KEY"],
vendor: "Groq", vendor: "Groq",
@@ -357,6 +418,7 @@ const MISTRAL_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.mistral.ai/v1", default: "https://api.mistral.ai/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["MISTRAL_API_KEY"], api_key_envs: &["MISTRAL_API_KEY"],
vendor: "Mistral", vendor: "Mistral",
@@ -390,6 +452,7 @@ const FIREWORKS_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.fireworks.ai/inference/v1", default: "https://api.fireworks.ai/inference/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["FIREWORKS_API_KEY"], api_key_envs: &["FIREWORKS_API_KEY"],
vendor: "Fireworks", vendor: "Fireworks",
@@ -420,6 +483,7 @@ const GOOGLE_SPEC: PlatformSpec = PlatformSpec {
default: "https://generativelanguage.googleapis.com/v1beta/openai", default: "https://generativelanguage.googleapis.com/v1beta/openai",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["GEMINI_API_KEY"], api_key_envs: &["GEMINI_API_KEY"],
vendor: "Google", vendor: "Google",
@@ -450,6 +514,7 @@ const OPENROUTER_SPEC: PlatformSpec = PlatformSpec {
default: "https://openrouter.ai/api/v1", default: "https://openrouter.ai/api/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["OPENROUTER_API_KEY"], api_key_envs: &["OPENROUTER_API_KEY"],
vendor: "OpenRouter", vendor: "OpenRouter",
@@ -480,6 +545,7 @@ const TOGETHER_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.together.xyz/v1", default: "https://api.together.xyz/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["TOGETHER_API_KEY"], api_key_envs: &["TOGETHER_API_KEY"],
vendor: "Together", vendor: "Together",
@@ -510,6 +576,7 @@ const CEREBRAS_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.cerebras.ai/v1", default: "https://api.cerebras.ai/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["CEREBRAS_API_KEY"], api_key_envs: &["CEREBRAS_API_KEY"],
vendor: "Cerebras", vendor: "Cerebras",
@@ -542,6 +609,7 @@ const NVIDIA_SPEC: PlatformSpec = PlatformSpec {
default: "https://integrate.api.nvidia.com/v1", default: "https://integrate.api.nvidia.com/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["NVIDIA_API_KEY"], api_key_envs: &["NVIDIA_API_KEY"],
vendor: "NVIDIA", vendor: "NVIDIA",
@@ -574,6 +642,7 @@ const VERCEL_SPEC: PlatformSpec = PlatformSpec {
default: "https://ai-gateway.vercel.sh/v1", default: "https://ai-gateway.vercel.sh/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["AI_GATEWAY_API_KEY"], api_key_envs: &["AI_GATEWAY_API_KEY"],
vendor: "Vercel", vendor: "Vercel",
@@ -606,6 +675,7 @@ const XAI_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.x.ai/v1", default: "https://api.x.ai/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
// NOTE: `XAI_API_KEY` is ALSO read as a legacy fallback by the house BYOK // 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`. // 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, 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"; pub const QWEN_TOKEN_PLAN_BASE_URL_ENV: &str = "KIGI_QWEN_TOKEN_PLAN_BASE_URL";
const QWEN_TOKEN_PLAN_SPEC: PlatformSpec = PlatformSpec { const QWEN_TOKEN_PLAN_SPEC: PlatformSpec = PlatformSpec {
id: "qwen-token-plan", 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", default: "https://token-plan.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["QWEN_TOKEN_PLAN_API_KEY"], api_key_envs: &["QWEN_TOKEN_PLAN_API_KEY"],
vendor: "Alibaba", 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", default: "https://token-plan.cn-beijing.maas.aliyuncs.com/compatible-mode/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["QWEN_TOKEN_PLAN_CN_API_KEY"], api_key_envs: &["QWEN_TOKEN_PLAN_CN_API_KEY"],
vendor: "Alibaba", vendor: "Alibaba",
@@ -698,6 +805,7 @@ const KIMI_CODING_SPEC: PlatformSpec = PlatformSpec {
// Kimi thinking dialect applies verbatim. // Kimi thinking dialect applies verbatim.
base_url: BaseUrlSource::KigiEnvCoding, base_url: BaseUrlSource::KigiEnvCoding,
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["KIMI_API_KEY"], api_key_envs: &["KIMI_API_KEY"],
vendor: "Kimi", vendor: "Kimi",
@@ -727,6 +835,7 @@ const ZAI_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.z.ai/api/coding/paas/v4", default: "https://api.z.ai/api/coding/paas/v4",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["ZAI_API_KEY"], api_key_envs: &["ZAI_API_KEY"],
vendor: "Z.AI", vendor: "Z.AI",
@@ -756,6 +865,7 @@ const ZAI_CODING_CN_SPEC: PlatformSpec = PlatformSpec {
default: "https://open.bigmodel.cn/api/coding/paas/v4", default: "https://open.bigmodel.cn/api/coding/paas/v4",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["ZAI_CODING_CN_API_KEY"], api_key_envs: &["ZAI_CODING_CN_API_KEY"],
vendor: "Z.AI", vendor: "Z.AI",
@@ -782,6 +892,7 @@ const XIAOMI_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.xiaomimimo.com/v1", default: "https://api.xiaomimimo.com/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["XIAOMI_API_KEY"], api_key_envs: &["XIAOMI_API_KEY"],
vendor: "Xiaomi", vendor: "Xiaomi",
@@ -810,6 +921,7 @@ const XIAOMI_TOKEN_PLAN_CN_SPEC: PlatformSpec = PlatformSpec {
default: "https://token-plan-cn.xiaomimimo.com/v1", default: "https://token-plan-cn.xiaomimimo.com/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["XIAOMI_TOKEN_PLAN_CN_API_KEY"], api_key_envs: &["XIAOMI_TOKEN_PLAN_CN_API_KEY"],
vendor: "Xiaomi", vendor: "Xiaomi",
@@ -841,6 +953,7 @@ const MINIMAX_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.minimax.io/anthropic/v1", default: "https://api.minimax.io/anthropic/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["MINIMAX_API_KEY"], api_key_envs: &["MINIMAX_API_KEY"],
vendor: "MiniMax", vendor: "MiniMax",
@@ -871,6 +984,7 @@ const MINIMAX_CN_SPEC: PlatformSpec = PlatformSpec {
default: "https://api.minimaxi.com/anthropic/v1", default: "https://api.minimaxi.com/anthropic/v1",
}, },
uses_oauth: false, uses_oauth: false,
oauth: None,
allowed_model_prefixes: None, allowed_model_prefixes: None,
api_key_envs: &["MINIMAX_CN_API_KEY"], api_key_envs: &["MINIMAX_CN_API_KEY"],
vendor: "MiniMax", vendor: "MiniMax",
@@ -941,12 +1055,14 @@ pub enum PlatformId {
Minimax, Minimax,
/// MiniMax, China (API key, Anthropic-compatible Messages). /// MiniMax, China (API key, Anthropic-compatible Messages).
MinimaxCn, MinimaxCn,
/// xAI Grok subscription via device-code OAuth (same wire as `Xai`).
XaiGrok,
} }
impl PlatformId { impl PlatformId {
/// All platforms, in catalog precedence order: the subscription channel /// All platforms, in catalog precedence order: the subscription channel
/// first so "default model = first list item" favors it when present. /// first so "default model = first list item" favors it when present.
pub const ALL: [PlatformId; 25] = [ pub const ALL: [PlatformId; 26] = [
Self::KimiCode, Self::KimiCode,
Self::MoonshotCn, Self::MoonshotCn,
Self::MoonshotAi, Self::MoonshotAi,
@@ -972,6 +1088,7 @@ impl PlatformId {
Self::XiaomiTokenPlanCn, Self::XiaomiTokenPlanCn,
Self::Minimax, Self::Minimax,
Self::MinimaxCn, Self::MinimaxCn,
Self::XaiGrok,
]; ];
/// The registry row backing this platform (single source of per-platform /// 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::XiaomiTokenPlanCn => &XIAOMI_TOKEN_PLAN_CN_SPEC,
Self::Minimax => &MINIMAX_SPEC, Self::Minimax => &MINIMAX_SPEC,
Self::MinimaxCn => &MINIMAX_CN_SPEC, Self::MinimaxCn => &MINIMAX_CN_SPEC,
Self::XaiGrok => &XAI_GROK_SPEC,
} }
} }
@@ -1034,6 +1152,14 @@ impl PlatformId {
self.spec().uses_oauth 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. /// Model-id prefixes admitted from this platform's `/models` listing.
/// `None` = no filtering (subscription listing is served pre-filtered). /// `None` = no filtering (subscription listing is served pre-filtered).
pub fn allowed_model_prefixes(self) -> Option<&'static [&'static str]> { pub fn allowed_model_prefixes(self) -> Option<&'static [&'static str]> {
@@ -1764,6 +1890,49 @@ mod tests {
assert_eq!(PlatformId::parse("not-a-platform"), None); 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 /// A variant missing from `ALL` compiles fine (`ALL`'s length is a plain
/// literal) but is silently unparseable and excluded from model sync. /// literal) but is silently unparseable and excluded from model sync.
/// The exhaustive match below fails compilation when a variant is added, /// The exhaustive match below fails compilation when a variant is added,
@@ -1797,9 +1966,10 @@ mod tests {
PlatformId::XiaomiTokenPlanCn => 22, PlatformId::XiaomiTokenPlanCn => 22,
PlatformId::Minimax => 23, PlatformId::Minimax => 23,
PlatformId::MinimaxCn => 24, 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<usize> = PlatformId::ALL.iter().map(|&p| ordinal(p)).collect(); let mut seen: Vec<usize> = PlatformId::ALL.iter().map(|&p| ordinal(p)).collect();
seen.sort_unstable(); seen.sort_unstable();
seen.dedup(); seen.dedup();
+8 -1
View File
@@ -178,7 +178,13 @@ async fn prefetch_models(agent_config: &AgentConfig) -> Option<IndexMap<String,
if auth.is_some() || endpoints.has_custom_endpoint() || platform_keys.any() { if auth.is_some() || endpoints.has_custom_endpoint() || platform_keys.any() {
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
prefetch_models_blocking(&endpoints, auth.as_ref(), fetch_auth, &platform_keys) prefetch_models_blocking(
&endpoints,
auth.as_ref(),
&Default::default(),
fetch_auth,
&platform_keys,
)
}) })
.await .await
.ok() .ok()
@@ -621,6 +627,7 @@ pub async fn run_leader(
crate::agent::models::prefetch_models_blocking( crate::agent::models::prefetch_models_blocking(
&endpoints_for_prefetch, &endpoints_for_prefetch,
auth_for_prefetch.as_ref(), auth_for_prefetch.as_ref(),
&Default::default(),
fetch_auth_for_prefetch, fetch_auth_for_prefetch,
&platform_keys_for_prefetch, &platform_keys_for_prefetch,
) )
@@ -114,7 +114,9 @@ pub struct BuiltAuthMethods {
/// 1. `xai.api_key` (if `has_external_api_key`) /// 1. `xai.api_key` (if `has_external_api_key`)
/// 2. `cached_token` (if `has_cached_token`) /// 2. `cached_token` (if `has_cached_token`)
/// 3. `kimi-code` (the Kimi Code device login) /// 3. `kimi-code` (the Kimi Code device login)
/// 4. every API-key registry platform, in `PlatformId::ALL` order /// 4. every generic device-code OAuth login (`xai-grok`, …), in
/// `PlatformId::ALL` order — interactive logins after `kimi-code`
/// 5. every API-key registry platform, in `PlatformId::ALL` order
/// (`moonshot-cn`, `moonshot-ai`, …), always advertised /// (`moonshot-cn`, `moonshot-ai`, …), always advertised
/// ///
/// The platform methods are for the INTERACTIVE login picker only: they come /// The platform methods are for the INTERACTIVE login picker only: they come
@@ -162,6 +164,15 @@ pub fn build_auth_methods(inputs: AuthMethodsBuildInputs<'_>) -> BuiltAuthMethod
} }
methods.push(kimi_code_auth_method(login_label)); 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 { for platform in kigi_models::PlatformId::ALL {
if !platform.uses_oauth() { if !platform.uses_oauth() {
methods.push(platform_auth_method(platform)); methods.push(platform_auth_method(platform));
@@ -182,6 +193,9 @@ pub enum AuthMethodKind {
KimiCode, KimiCode,
/// Registry API-key platform login (method id = the platform id). /// Registry API-key platform login (method id = the platform id).
ApiKeyPlatform(kigi_models::PlatformId), 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, Unknown,
} }
@@ -191,9 +205,12 @@ impl AuthMethodKind {
XAI_API_KEY_METHOD_ID => Self::XaiApiKey, XAI_API_KEY_METHOD_ID => Self::XaiApiKey,
CACHED_TOKEN_AUTH_METHOD_ID => Self::CachedToken, CACHED_TOKEN_AUTH_METHOD_ID => Self::CachedToken,
KIMI_CODE_METHOD_ID => Self::KimiCode, KIMI_CODE_METHOD_ID => Self::KimiCode,
other => match platform_for_method_id_str(other) { other => match kigi_models::PlatformId::parse(other) {
Some(platform) => Self::ApiKeyPlatform(platform), // A generic device-code OAuth platform (xai-grok).
None => Self::Unknown, 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). /// `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 { 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 { pub fn needs_interactive_login(self) -> bool {
matches!(self, Self::KimiCode) matches!(self, Self::KimiCode | Self::OAuthPlatform(_))
}
/// The generic device-code OAuth platform behind this method, if any
/// (drives the `authenticate` dispatch to the generic device flow).
pub fn oauth_platform(self) -> Option<kigi_models::PlatformId> {
match self {
Self::OAuthPlatform(p) => Some(p),
_ => None,
}
} }
pub fn auth_error_message(self) -> &'static str { 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 /// Interactive API-key login method ids equal
/// [`kigi_models::PlatformId::as_str`] (`moonshot-cn` / `moonshot-ai` / …), /// [`kigi_models::PlatformId::as_str`] (`moonshot-cn` / `moonshot-ai` / …),
/// which is also the `[platforms.<id>]` config-table name and the auth.json /// which is also the `[platforms.<id>]` config-table name and the auth.json
@@ -514,6 +561,73 @@ mod tests {
); );
} }
/// xai-grok is a generic device-code OAuth login: it classifies as an
/// `OAuthPlatform`, needs an interactive (browser) login, is NOT api-key,
/// is NOT an API-key registry platform, and is advertised right after
/// kimi-code among the interactive logins.
#[test]
fn xai_grok_is_an_interactive_oauth_login() {
let id = acp::AuthMethodId::new("xai-grok");
let kind = AuthMethodKind::from_id(&id);
assert_eq!(
kind,
AuthMethodKind::OAuthPlatform(kigi_models::PlatformId::XaiGrok)
);
assert!(
kind.needs_interactive_login(),
"xai-grok needs a browser login"
);
assert!(!kind.is_api_key(), "xai-grok is not an api-key method");
assert_eq!(
kind.oauth_platform(),
Some(kigi_models::PlatformId::XaiGrok)
);
// Never an API-key picker target (keeps it out of the paste-box path).
assert_eq!(platform_for_method_id(&id), None);
// Placement: immediately after kimi-code, before the api-key rows.
let built = build_auth_methods(default_inputs());
let ids = method_ids(&built);
let kimi_pos = ids.iter().position(|m| *m == KIMI_CODE_METHOD_ID).unwrap();
assert_eq!(
ids[kimi_pos + 1],
"xai-grok",
"xai-grok must be the interactive login right after kimi-code"
);
assert_eq!(
ids[kimi_pos + 2],
MOONSHOT_CN_METHOD_ID,
"the api-key rows follow the generic oauth logins"
);
}
/// A generic device-code OAuth platform (xai-grok) is SESSION-BASED: its
/// device-code session is refreshable, so the per-turn refresh / 401 gate
/// must be active for it (routing then sends the model's OWN manager). It
/// stays an interactive login and is NOT api-key-shaped.
#[test]
fn oauth_platform_is_session_based_and_refreshable() {
let id = acp::AuthMethodId::new("xai-grok");
let kind = AuthMethodKind::from_id(&id);
assert_eq!(
kind,
AuthMethodKind::OAuthPlatform(kigi_models::PlatformId::XaiGrok)
);
assert!(
kind.is_session_based(),
"OAuthPlatform must be session-based so the refresh/401 gate is active"
);
assert!(
is_session_based_method(&id),
"is_session_based_method(xai-grok) must be true"
);
// Session-based, yet still an interactive browser login and never
// api-key-shaped.
assert!(kind.needs_interactive_login());
assert!(!kind.is_api_key());
// The session-expired copy (not the api-key copy) is the right error.
assert_eq!(kind.auth_error_message(), AUTH_ERROR_SESSION_EXPIRED);
}
/// The OAuth platform id must never resolve as an API-key platform /// The OAuth platform id must never resolve as an API-key platform
/// method — `platform_for_method_id`'s `uses_oauth` filter is what keeps /// method — `platform_for_method_id`'s `uses_oauth` filter is what keeps
/// the generic `authenticate` arm from hijacking the device login. /// the generic `authenticate` arm from hijacking the device login.
@@ -606,6 +720,7 @@ mod tests {
vec![ vec![
XAI_API_KEY_METHOD_ID, XAI_API_KEY_METHOD_ID,
KIMI_CODE_METHOD_ID, KIMI_CODE_METHOD_ID,
"xai-grok",
MOONSHOT_CN_METHOD_ID, MOONSHOT_CN_METHOD_ID,
MOONSHOT_AI_METHOD_ID, MOONSHOT_AI_METHOD_ID,
"openai", "openai",
@@ -654,6 +769,7 @@ mod tests {
XAI_API_KEY_METHOD_ID, XAI_API_KEY_METHOD_ID,
CACHED_TOKEN_AUTH_METHOD_ID, CACHED_TOKEN_AUTH_METHOD_ID,
KIMI_CODE_METHOD_ID, KIMI_CODE_METHOD_ID,
"xai-grok",
MOONSHOT_CN_METHOD_ID, MOONSHOT_CN_METHOD_ID,
MOONSHOT_AI_METHOD_ID, MOONSHOT_AI_METHOD_ID,
"openai", "openai",
@@ -695,6 +811,7 @@ mod tests {
vec![ vec![
CACHED_TOKEN_AUTH_METHOD_ID, CACHED_TOKEN_AUTH_METHOD_ID,
KIMI_CODE_METHOD_ID, KIMI_CODE_METHOD_ID,
"xai-grok",
MOONSHOT_CN_METHOD_ID, MOONSHOT_CN_METHOD_ID,
MOONSHOT_AI_METHOD_ID, MOONSHOT_AI_METHOD_ID,
"openai", "openai",
@@ -739,6 +856,7 @@ mod tests {
method_ids(&built), method_ids(&built),
vec![ vec![
KIMI_CODE_METHOD_ID, KIMI_CODE_METHOD_ID,
"xai-grok",
MOONSHOT_CN_METHOD_ID, MOONSHOT_CN_METHOD_ID,
MOONSHOT_AI_METHOD_ID, MOONSHOT_AI_METHOD_ID,
"openai", "openai",
@@ -4599,6 +4599,80 @@ reasoning_effort = "low"
assert_eq!(resolved.base_url, "https://vendor.example/v1"); assert_eq!(resolved.base_url, "https://vendor.example/v1");
assert_eq!(resolved.api_key.as_deref(), Some("vendor-key")); assert_eq!(resolved.api_key.as_deref(), Some("vendor-key"));
} }
/// A primary Kimi manager holding a fixed in-memory bearer, standing in for
/// an aux caller's session. `TempDir` returned so the caller keeps it alive.
fn kimi_primary(key: &str) -> (tempfile::TempDir, std::sync::Arc<AuthManager>) {
let dir = tempfile::tempdir().unwrap();
let manager = std::sync::Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
manager.hot_swap(crate::auth::KimiAuth {
key: key.to_string(),
auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
(dir, manager)
}
/// LEAK 1a (aux/summary model): the aux `session_key` is resolved by the aux
/// model's OWN platform (as `build_summary_client` / `resolve_aux_sampler_config`
/// now do). A grok (oauth-platform) aux model's resolved sampler `api_key` is
/// therefore grok's own token or `None` — NEVER the primary Kimi key — while a
/// first-party / non-oauth aux model still gets the primary (byte-identical).
#[tokio::test]
async fn aux_model_session_key_is_platform_scoped_never_leaking_kimi() {
let home = tempfile::tempdir().unwrap();
let (_kd, kimi) = kimi_primary("kimi-tok");
let endpoints = EndpointsConfig::default();
// A grok aux catalog entry (managed id → oauth platform), no own key.
let mut grok = test_model_entry("grok-4-latest", "https://api.x.ai/v1", None, None, None);
grok.info.id = Some("xai-grok/grok-4-latest".to_string());
let mut grok_catalog = IndexMap::new();
grok_catalog.insert("grok".to_string(), grok);
let grok_key = crate::auth::oauth_registry::session_key_for_model(
home.path(),
"xai-grok/grok-4-latest",
Some(&kimi),
);
let grok_cfg = resolve_aux_model_sampling_config(
"grok",
&grok_catalog,
&endpoints,
grok_key.as_deref(),
None,
);
assert_ne!(
grok_cfg.and_then(|c| c.api_key).as_deref(),
Some("kimi-tok"),
"a grok aux model must never receive the primary Kimi session token",
);
// A non-oauth aux catalog entry still resolves to the primary token.
let mut k2 = test_model_entry(
"kimi-k2-0905-preview",
"https://vendor/v1",
None,
None,
None,
);
k2.info.id = Some("moonshot-cn/kimi-k2".to_string());
let mut k2_catalog = IndexMap::new();
k2_catalog.insert("k2".to_string(), k2);
let k2_key = crate::auth::oauth_registry::session_key_for_model(
home.path(),
"moonshot-cn/kimi-k2",
Some(&kimi),
);
let k2_cfg = resolve_aux_model_sampling_config(
"k2",
&k2_catalog,
&endpoints,
k2_key.as_deref(),
None,
)
.expect("non-oauth aux resolves via the primary session token");
assert_eq!(
k2_cfg.api_key.as_deref(),
Some("kimi-tok"),
"a non-oauth aux model must still receive the primary session token",
);
}
#[test] #[test]
fn parses_model_api_key() { fn parses_model_api_key() {
let raw_config: toml::Value = toml::from_str( let raw_config: toml::Value = toml::from_str(
+69 -8
View File
@@ -321,6 +321,7 @@ impl ModelsManager {
&cfg.endpoints, &cfg.endpoints,
fetch_auth, fetch_auth,
has_session, has_session,
&Default::default(),
&platform_keys, &platform_keys,
), ),
) )
@@ -1062,7 +1063,6 @@ impl ModelsManager {
/// Build a `SamplingConfig` from the current model + auth state. /// Build a `SamplingConfig` from the current model + auth state.
pub fn sampling_config(&self) -> SamplingConfig { pub fn sampling_config(&self) -> SamplingConfig {
let config = self.inner.cfg.read().clone(); 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 current_model_id = self.current_model_id();
let all_models = self.models(); let all_models = self.models();
let fallback; let fallback;
@@ -1079,9 +1079,12 @@ impl ModelsManager {
} }
}; };
let session_auth = auth_manager.current_or_expired(); // Resolve the session bearer PER the current model's platform: a
let credentials = // generic device-code OAuth platform (xai-grok) resolves from ITS OWN
resolve_credentials(current_model, session_auth.as_ref().map(|a| a.key.as_str())); // 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( sampling_config_for_model(
current_model, current_model,
@@ -1090,6 +1093,23 @@ impl ModelsManager {
) )
} }
/// The session bearer for a model catalog key. Generic device-code OAuth
/// platforms (xai-grok) resolve from their own persisted scope; every other
/// key resolves from the primary (Kimi) `AuthManager`. Reads the persisted
/// token (refresh happens in the async catalog/refresh paths); never logs
/// it.
fn session_key_for_catalog_key(&self, catalog_key: &str) -> Option<String> {
if let Some((platform, _)) = kigi_models::parse_managed_model_key(catalog_key)
&& let Some(oauth) = platform.oauth()
{
let kigi_home = crate::util::kigi_home::kigi_home();
return AuthManager::new_oauth_provider(&kigi_home, oauth)
.current_or_expired()
.map(|a| a.key);
}
self.inner.auth_manager.current_or_expired().map(|a| a.key)
}
/// Disk-cache origin key for this manager's current endpoints/auth shape /// Disk-cache origin key for this manager's current endpoints/auth shape
/// (see [`ModelsCache::origin`]). /// (see [`ModelsCache::origin`]).
fn cache_origin(&self) -> String { fn cache_origin(&self) -> String {
@@ -1100,10 +1120,15 @@ impl ModelsManager {
let fetch_auth = *self.inner.fetch_auth.read(); let fetch_auth = *self.inner.fetch_auth.read();
let has_oauth = self.inner.auth_manager.current_or_expired().is_some(); let has_oauth = self.inner.auth_manager.current_or_expired().is_some();
let platform_keys = PlatformApiKeys::resolve(&platforms); 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( crate::agent::models_fetch::models_fetch_origin(
&endpoints, &endpoints,
fetch_auth, fetch_auth,
has_oauth, has_oauth,
&Default::default(),
&platform_keys, &platform_keys,
) )
} }
@@ -1155,8 +1180,21 @@ impl ModelsManager {
let fetch_auth = *self.inner.fetch_auth.read(); let fetch_auth = *self.inner.fetch_auth.read();
let platform_keys = PlatformApiKeys::resolve(&cfg.platforms); let platform_keys = PlatformApiKeys::resolve(&cfg.platforms);
let auth = self.inner.auth_manager.auth().await.ok(); let auth = self.inner.auth_manager.auth().await.ok();
let outcome = // Resolve each generic device-code OAuth platform's OWN session token
fetch_models_async(endpoints.clone(), auth, fetch_auth, platform_keys.clone()).await; // (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() { if outcome.models.is_some() {
return outcome.models; return outcome.models;
} }
@@ -1173,7 +1211,12 @@ impl ModelsManager {
return None; return None;
} }
let auth = self.inner.auth_manager.auth().await.ok(); 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 { if retry.oauth_unauthorized {
tracing::warn!("model catalog: still unauthorized after token refresh"); tracing::warn!("model catalog: still unauthorized after token refresh");
} }
@@ -1614,12 +1657,14 @@ impl ModelsFetchOutcome {
pub(crate) fn prefetch_models_blocking( pub(crate) fn prefetch_models_blocking(
endpoints: &config::EndpointsConfig, endpoints: &config::EndpointsConfig,
auth: Option<&KimiAuth>, auth: Option<&KimiAuth>,
oauth_tokens: &crate::agent::models_fetch::OAuthSessionTokens,
fetch_auth: ModelFetchAuth, fetch_auth: ModelFetchAuth,
platform_keys: &PlatformApiKeys, platform_keys: &PlatformApiKeys,
) -> Option<IndexMap<String, ModelEntry>> { ) -> Option<IndexMap<String, ModelEntry>> {
prefetch_models_blocking_gated( prefetch_models_blocking_gated(
endpoints, endpoints,
auth, auth,
oauth_tokens,
fetch_auth, fetch_auth,
platform_keys, platform_keys,
crate::util::config::resolve_remote_fetch_enabled(), crate::util::config::resolve_remote_fetch_enabled(),
@@ -1632,6 +1677,7 @@ pub(crate) fn prefetch_models_blocking(
fn prefetch_models_blocking_gated( fn prefetch_models_blocking_gated(
endpoints: &config::EndpointsConfig, endpoints: &config::EndpointsConfig,
auth: Option<&KimiAuth>, auth: Option<&KimiAuth>,
oauth_tokens: &crate::agent::models_fetch::OAuthSessionTokens,
fetch_auth: ModelFetchAuth, fetch_auth: ModelFetchAuth,
platform_keys: &PlatformApiKeys, platform_keys: &PlatformApiKeys,
remote_fetch_enabled: bool, remote_fetch_enabled: bool,
@@ -1643,6 +1689,7 @@ fn prefetch_models_blocking_gated(
endpoints, endpoints,
fetch_auth, fetch_auth,
auth.is_some(), auth.is_some(),
oauth_tokens,
platform_keys, platform_keys,
); );
let cache = ModelsCacheManager::new(); let cache = ModelsCacheManager::new();
@@ -1666,7 +1713,7 @@ fn prefetch_models_blocking_gated(
} }
let _timer = crate::instrumentation_timer!("startup.fetch_models_blocking"); 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 { Ok(FetchModelsResult {
models, models,
etag, etag,
@@ -1826,6 +1873,10 @@ fn spawn_prefetch_thread(env: PrefetchEnv) -> EarlyPrefetchHandle {
let models = prefetch_models_blocking( let models = prefetch_models_blocking(
&env.endpoints, &env.endpoints,
env.auth.as_ref(), 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.model_fetch_auth,
&env.platform_keys, &env.platform_keys,
); );
@@ -2215,6 +2266,7 @@ pub(crate) fn validate_selectable(
pub(crate) async fn fetch_models_async( pub(crate) async fn fetch_models_async(
endpoints: config::EndpointsConfig, endpoints: config::EndpointsConfig,
auth: Option<KimiAuth>, auth: Option<KimiAuth>,
oauth_tokens: crate::agent::models_fetch::OAuthSessionTokens,
fetch_auth: ModelFetchAuth, fetch_auth: ModelFetchAuth,
platform_keys: PlatformApiKeys, platform_keys: PlatformApiKeys,
) -> ModelsFetchOutcome { ) -> ModelsFetchOutcome {
@@ -2222,6 +2274,7 @@ pub(crate) async fn fetch_models_async(
prefetch_models_blocking_gated( prefetch_models_blocking_gated(
&endpoints, &endpoints,
auth.as_ref(), auth.as_ref(),
&oauth_tokens,
fetch_auth, fetch_auth,
&platform_keys, &platform_keys,
crate::util::config::resolve_remote_fetch_enabled(), crate::util::config::resolve_remote_fetch_enabled(),
@@ -3950,6 +4003,7 @@ mod tests {
crate::agent::models_fetch::fetch_models_blocking( crate::agent::models_fetch::fetch_models_blocking(
&endpoints, &endpoints,
Some(&auth), Some(&auth),
&Default::default(),
ModelFetchAuth::Platforms, ModelFetchAuth::Platforms,
&PlatformApiKeys::default(), &PlatformApiKeys::default(),
) )
@@ -4019,6 +4073,7 @@ mod tests {
crate::agent::models_fetch::fetch_models_blocking( crate::agent::models_fetch::fetch_models_blocking(
&endpoints, &endpoints,
None, None,
&Default::default(),
ModelFetchAuth::Platforms, ModelFetchAuth::Platforms,
&keys, &keys,
) )
@@ -4170,6 +4225,7 @@ mod tests {
let outcome = prefetch_models_blocking_gated( let outcome = prefetch_models_blocking_gated(
&endpoints, &endpoints,
Some(&auth), Some(&auth),
&Default::default(),
ModelFetchAuth::Platforms, ModelFetchAuth::Platforms,
&keys, &keys,
true, true,
@@ -4185,6 +4241,7 @@ mod tests {
&endpoints, &endpoints,
ModelFetchAuth::Platforms, ModelFetchAuth::Platforms,
true, true,
&Default::default(),
&keys, &keys,
); );
let cache = ModelsCacheManager::new(); let cache = ModelsCacheManager::new();
@@ -4200,6 +4257,7 @@ mod tests {
let outcome = prefetch_models_blocking_gated( let outcome = prefetch_models_blocking_gated(
&endpoints, &endpoints,
Some(&auth), Some(&auth),
&Default::default(),
ModelFetchAuth::Platforms, ModelFetchAuth::Platforms,
&keys, &keys,
true, true,
@@ -4218,6 +4276,7 @@ mod tests {
let outcome = prefetch_models_blocking_gated( let outcome = prefetch_models_blocking_gated(
&endpoints, &endpoints,
Some(&auth), Some(&auth),
&Default::default(),
ModelFetchAuth::Platforms, ModelFetchAuth::Platforms,
&keys, &keys,
true, true,
@@ -4272,6 +4331,7 @@ mod tests {
&endpoints, &endpoints,
ModelFetchAuth::Platforms, ModelFetchAuth::Platforms,
true, true,
&Default::default(),
&PlatformApiKeys::test_keys(Some("sk"), None), &PlatformApiKeys::test_keys(Some("sk"), None),
); );
let cache = ModelsCacheManager::new(); let cache = ModelsCacheManager::new();
@@ -4291,6 +4351,7 @@ mod tests {
let outcome = prefetch_models_blocking_gated( let outcome = prefetch_models_blocking_gated(
&endpoints, &endpoints,
Some(&auth), Some(&auth),
&Default::default(),
ModelFetchAuth::Platforms, ModelFetchAuth::Platforms,
&PlatformApiKeys::default(), &PlatformApiKeys::default(),
true, true,
@@ -28,6 +28,44 @@ pub(crate) const DEFAULT_CONTEXT_WINDOW: u64 = 256_000;
struct ModelsResponse { struct ModelsResponse {
data: Vec<serde_json::Value>, data: Vec<serde_json::Value>,
} }
/// Session bearer tokens for the GENERIC device-code OAuth platforms
/// (`platform.oauth().is_some()` — xai-grok today), resolved per provider from
/// its own scope-keyed `AuthManager` (refreshed on expiry) before the blocking
/// fetch. kimi-code is NOT here — it still rides the single `auth: &KimiAuth`.
///
/// SECURITY: values are access tokens — never logged, never persisted here.
pub(crate) type OAuthSessionTokens = std::collections::BTreeMap<kigi_models::PlatformId, String>;
/// Resolve session bearers for every generic device-code OAuth platform whose
/// own `AuthManager` holds a usable (refreshed-on-expiry) session. Each
/// platform resolves from ITS OWN scope (`oauth/xai`, …) — independent of the
/// Kimi session. Providers without a stored session are simply absent. Only a
/// non-secret count is logged.
pub(crate) async fn resolve_generic_oauth_tokens(
kigi_home: &std::path::Path,
) -> OAuthSessionTokens {
let mut tokens = OAuthSessionTokens::new();
for platform in kigi_models::PlatformId::ALL {
let Some(oauth) = platform.oauth() else {
continue;
};
let manager = std::sync::Arc::new(crate::auth::AuthManager::new_oauth_provider(
kigi_home, oauth,
));
manager.configure_refresher();
if let Ok(auth) = manager.auth().await {
tokens.insert(platform, auth.key);
}
}
if !tokens.is_empty() {
tracing::info!(
count = tokens.len(),
"resolved generic oauth platform session tokens"
);
}
tokens
}
/// The models-fetch origin key for this endpoints/auth shape. Used as the /// 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 /// 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 /// 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, endpoints: &crate::agent::config::EndpointsConfig,
fetch_auth: crate::agent::models::ModelFetchAuth, fetch_auth: crate::agent::models::ModelFetchAuth,
has_oauth: bool, has_oauth: bool,
oauth_tokens: &OAuthSessionTokens,
platform_keys: &crate::agent::models::PlatformApiKeys, platform_keys: &crate::agent::models::PlatformApiKeys,
) -> String { ) -> String {
match fetch_auth { match fetch_auth {
crate::agent::models::ModelFetchAuth::CustomEndpoint => endpoints.resolve_models_list_url(), crate::agent::models::ModelFetchAuth::CustomEndpoint => endpoints.resolve_models_list_url(),
crate::agent::models::ModelFetchAuth::Platforms => { crate::agent::models::ModelFetchAuth::Platforms => {
let parts: Vec<String> = enabled_platforms(has_oauth, platform_keys) let parts: Vec<String> = enabled_platforms(has_oauth, oauth_tokens, platform_keys)
.into_iter() .into_iter()
.map(|p| format!("{}={}", p.as_str(), platform_models_url(p, endpoints))) .map(|p| format!("{}={}", p.as_str(), platform_models_url(p, endpoints)))
.collect(); .collect();
@@ -53,14 +92,23 @@ pub(crate) fn models_fetch_origin(
} }
/// The platforms with usable credentials, in registry order (kimi-code first /// The platforms with usable credentials, in registry order (kimi-code first
/// so "default model = first list item" favors the subscription). /// 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( fn enabled_platforms(
has_oauth: bool, has_oauth: bool,
oauth_tokens: &OAuthSessionTokens,
platform_keys: &crate::agent::models::PlatformApiKeys, platform_keys: &crate::agent::models::PlatformApiKeys,
) -> Vec<kigi_models::PlatformId> { ) -> Vec<kigi_models::PlatformId> {
kigi_models::PlatformId::ALL kigi_models::PlatformId::ALL
.into_iter() .into_iter()
.filter(|p| { .filter(|p| {
if p.uses_oauth() { if p.oauth().is_some() {
oauth_tokens.contains_key(p)
} else if p.uses_oauth() {
has_oauth has_oauth
} else { } else {
platform_keys.key_for(*p).is_some() platform_keys.key_for(*p).is_some()
@@ -68,20 +116,30 @@ fn enabled_platforms(
}) })
.collect() .collect()
} }
/// `{base}/models` for one platform. The subscription platform resolves its /// The inference/listing base for one platform's `/models` fetch:
/// base through the endpoints config (`coding_api_base_url` override, /// - kimi-code (`uses_oauth`, no `OAuthConfig`) → the Kimi subscription base
/// else `KIGI_CODE_BASE_URL` / production default via kigi-env); the open /// via the endpoints config (`proxy_url`);
/// platforms use their fixed bases. /// - 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( fn platform_models_url(
platform: kigi_models::PlatformId, platform: kigi_models::PlatformId,
endpoints: &crate::agent::config::EndpointsConfig, endpoints: &crate::agent::config::EndpointsConfig,
) -> String { ) -> String {
let base = if platform.uses_oauth() { format!(
endpoints.proxy_url() "{}/models",
} else { platform_fetch_base(platform, endpoints).trim_end_matches('/')
platform.base_url() )
};
format!("{}/models", base.trim_end_matches('/'))
} }
/// Fetch result: model entries + optional etag from the subscription platform. /// Fetch result: model entries + optional etag from the subscription platform.
pub struct FetchModelsResult { pub struct FetchModelsResult {
@@ -106,6 +164,7 @@ pub struct FetchModelsResult {
pub(crate) fn fetch_models_blocking( pub(crate) fn fetch_models_blocking(
endpoints: &crate::agent::config::EndpointsConfig, endpoints: &crate::agent::config::EndpointsConfig,
auth: Option<&KimiAuth>, auth: Option<&KimiAuth>,
oauth_tokens: &OAuthSessionTokens,
fetch_auth: crate::agent::models::ModelFetchAuth, fetch_auth: crate::agent::models::ModelFetchAuth,
platform_keys: &crate::agent::models::PlatformApiKeys, platform_keys: &crate::agent::models::PlatformApiKeys,
) -> Result<FetchModelsResult, BackendError> { ) -> Result<FetchModelsResult, BackendError> {
@@ -114,7 +173,7 @@ pub(crate) fn fetch_models_blocking(
fetch_custom_endpoint_models_blocking(endpoints, auth) fetch_custom_endpoint_models_blocking(endpoints, auth)
} }
crate::agent::models::ModelFetchAuth::Platforms => { 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( fn fetch_platform_models_blocking(
endpoints: &crate::agent::config::EndpointsConfig, endpoints: &crate::agent::config::EndpointsConfig,
auth: Option<&KimiAuth>, auth: Option<&KimiAuth>,
oauth_tokens: &OAuthSessionTokens,
platform_keys: &crate::agent::models::PlatformApiKeys, platform_keys: &crate::agent::models::PlatformApiKeys,
) -> Result<FetchModelsResult, BackendError> { ) -> Result<FetchModelsResult, BackendError> {
let enabled = enabled_platforms(auth.is_some(), platform_keys); let enabled = enabled_platforms(auth.is_some(), oauth_tokens, platform_keys);
if enabled.is_empty() { if enabled.is_empty() {
return Err(BackendError::Auth( return Err(BackendError::Auth(
"No platform credentials: log in with `kigi login`, paste a platform API key in \ "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). // serves its own metadata (kimi/moonshot today).
let enrichment = crate::agent::enrichment_fetch::load_enrichment_catalog(&enabled); let enrichment = crate::agent::enrichment_fetch::load_enrichment_catalog(&enabled);
for platform in &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()) auth.map(|a| a.key.clone())
.expect("enabled_platforms gated on auth presence") .expect("enabled_platforms gated on auth presence")
} else { } else {
@@ -211,7 +278,9 @@ fn fetch_platform_models_blocking(
"platform models fetch succeeded" "platform models fetch succeeded"
); );
successes += 1; 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; etag = platform_etag;
} }
models.extend(platform_models); models.extend(platform_models);
@@ -383,11 +452,7 @@ fn fetch_one_platform_models(
); );
} }
} }
let base_url = if platform.uses_oauth() { let base_url = platform_fetch_base(platform, endpoints);
endpoints.proxy_url()
} else {
platform.base_url()
};
let models = filtered let models = filtered
.into_iter() .into_iter()
.map(|mut wire| { .map(|mut wire| {
@@ -835,7 +900,7 @@ mod tests {
"sk-oai", "sk-oai",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -983,7 +1048,7 @@ mod tests {
"sk-ant", "sk-ant",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -1095,7 +1160,7 @@ mod tests {
"sk-ds", "sk-ds",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -1193,7 +1258,7 @@ mod tests {
"gsk-1", "gsk-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -1277,7 +1342,7 @@ mod tests {
"msk-1", "msk-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -1363,7 +1428,7 @@ mod tests {
"fw-1", "fw-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -1471,7 +1536,7 @@ mod tests {
"gk-1", "gk-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -1555,7 +1620,7 @@ mod tests {
"or-1", "or-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -1660,7 +1725,7 @@ mod tests {
"tg-1", "tg-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -1753,7 +1818,7 @@ mod tests {
"cb-1", "cb-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -1861,7 +1926,7 @@ mod tests {
"nvapi-1", "nvapi-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -1955,7 +2020,7 @@ mod tests {
"vg-1", "vg-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -2048,7 +2113,7 @@ mod tests {
"xai-1", "xai-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -2139,7 +2204,7 @@ mod tests {
"qtp-1", "qtp-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -2221,7 +2286,7 @@ mod tests {
"kc-1", "kc-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -2312,7 +2377,7 @@ mod tests {
"zai-1", "zai-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -2402,7 +2467,7 @@ mod tests {
"xm-1", "xm-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -2493,7 +2558,7 @@ mod tests {
"mm-1", "mm-1",
); );
let result = tokio::task::spawn_blocking(move || { let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &keys) fetch_platform_models_blocking(&endpoints, None, &Default::default(), &keys)
}) })
.await .await
.unwrap() .unwrap()
@@ -2991,6 +3056,125 @@ mod tests {
"https://registry.acme.com/api/list-models" "https://registry.acme.com/api/list-models"
); );
} }
/// xai-grok OAuth-cycle e2e (mock wire): the generic device-code OAuth
/// bearer (`grok-oauth-tok`, resolved from a stored `oauth/xai` session —
/// mocked here as the token map) fetches `GET {base}/models` against the
/// platform's OWN base (api.x.ai/v1 via its env), enriches from models.dev
/// "xai", restricts to tool-calling chat models, and keys each entry under
/// `xai-grok/<id>` with the Passthrough dialect — NOT the API-key `xai/`.
#[tokio::test(flavor = "multi_thread")]
#[serial_test::serial]
async fn xai_grok_oauth_listing_is_enriched_restricted_and_keyed() {
let platform_server = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("GET"))
.and(wiremock::matchers::path("/models"))
.and(wiremock::matchers::header(
"Authorization",
"Bearer grok-oauth-tok",
))
.respond_with(wiremock::ResponseTemplate::new(200).set_body_json(
serde_json::json!({ "data": [
{ "id": "grok-4.5", "object": "model" },
// enrichment-known but NOT tool-calling → dropped by restrict.
{ "id": "grok-2-image", "object": "model" }
]}),
))
.expect(1)
.mount(&platform_server)
.await;
let modelsdev_server = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("GET"))
.and(wiremock::matchers::path("/api.json"))
.respond_with(wiremock::ResponseTemplate::new(200).set_body_json(
serde_json::json!({ "xai": { "models": {
"grok-4.5": {
"name": "Grok 4.5",
"reasoning": true,
"limit": {"context": 256000, "output": 64000},
"tool_call": true
},
"grok-2-image": { "limit": {"context": 8192} }
}}}),
))
.expect(1)
.mount(&modelsdev_server)
.await;
let cache_dir = tempfile::tempdir().unwrap();
let _base = kigi_test_support::EnvGuard::set(
kigi_models::XAI_GROK_BASE_URL_ENV,
platform_server.uri(),
);
let _mdev = kigi_test_support::EnvGuard::set(
crate::agent::enrichment_fetch::MODELS_DEV_URL_ENV,
format!("{}/api.json", modelsdev_server.uri()),
);
let _mdev_cache = kigi_test_support::EnvGuard::set(
crate::agent::enrichment_fetch::MODELS_DEV_CACHE_DIR_ENV,
cache_dir.path(),
);
let endpoints = crate::agent::config::EndpointsConfig::default();
let keys = crate::agent::models::PlatformApiKeys::default();
let mut tokens = OAuthSessionTokens::new();
tokens.insert(kigi_models::PlatformId::XaiGrok, "grok-oauth-tok".into());
let result = tokio::task::spawn_blocking(move || {
fetch_platform_models_blocking(&endpoints, None, &tokens, &keys)
})
.await
.unwrap()
.expect("fetch must succeed");
assert_eq!(
result
.models
.iter()
.map(|m| m.id.as_deref().unwrap_or_default())
.collect::<Vec<_>>(),
vec!["xai-grok/grok-4.5"],
"grok-2-image (enrichment-known, not tool-calling) must be dropped; \
the surviving model is keyed under xai-grok/ (not xai/)"
);
let entry = &result.models[0];
assert_eq!(entry.model, "grok-4.5");
assert_eq!(
entry.base_url,
platform_server.uri(),
"xai-grok fetches against its OWN base (its base env), not proxy_url"
);
assert_eq!(
entry.context_window.get(),
256_000,
"context window must come from models.dev \"xai\" enrichment"
);
assert_eq!(entry.max_completion_tokens, Some(64_000));
assert_eq!(
entry.api_backend,
crate::sampling::ApiBackend::ChatCompletions
);
assert_eq!(entry.name.as_deref(), Some("Grok 4.5"));
assert!(
entry.env_key.is_none(),
"an OAuth channel carries no api-key env"
);
assert!(
!entry.supported_in_api,
"subscription models require the OAuth session (not the public API)"
);
// Passthrough dialect (identical to the API-key xai wire).
let model_entry = crate::agent::config::ModelEntry::from_config_entry(entry);
let creds = crate::agent::config::ResolvedCredentials {
api_key: Some("grok-oauth-tok".into()),
base_url: entry.base_url.clone(),
auth_type: kigi_chat_state::AuthType::SessionToken,
auth_scheme: Default::default(),
};
let cfg = crate::agent::config::sampling_config_for_model(&model_entry, creds, None);
assert_eq!(
cfg.chat_compat,
kigi_sampling_types::ChatCompat::Passthrough
);
}
/// INVARIANT: each platform's `/models` URL matches its registry base — /// INVARIANT: each platform's `/models` URL matches its registry base —
/// kimi-code → the subscription proxy (config override respected, else the /// kimi-code → the subscription proxy (config override respected, else the
/// kigi-env default), moonshot platforms → their fixed bases — and 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), platform_models_url(kigi_models::PlatformId::OpenAi, &cfg),
"https://api.openai.com/v1/models" "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. // Proxy override re-points the subscription platform only.
let proxied = EndpointsConfig::from_config_value( let proxied = EndpointsConfig::from_config_value(
&toml::from_str( &toml::from_str(
@@ -3047,6 +3237,7 @@ mod tests {
&cfg, &cfg,
ModelFetchAuth::Platforms, ModelFetchAuth::Platforms,
true, true,
&Default::default(),
&PlatformApiKeys::default(), &PlatformApiKeys::default(),
); );
assert_eq!( assert_eq!(
@@ -3057,6 +3248,7 @@ mod tests {
&cfg, &cfg,
ModelFetchAuth::Platforms, ModelFetchAuth::Platforms,
true, true,
&Default::default(),
&crate::agent::models::PlatformApiKeys::test_keys(Some("sk-secret-cn"), None), &crate::agent::models::PlatformApiKeys::test_keys(Some("sk-secret-cn"), None),
); );
assert_ne!( assert_ne!(
@@ -3082,6 +3274,7 @@ mod tests {
&custom, &custom,
ModelFetchAuth::CustomEndpoint, ModelFetchAuth::CustomEndpoint,
false, false,
&Default::default(),
&PlatformApiKeys::default(), &PlatformApiKeys::default(),
), ),
"https://models.acme.com/v1/models" "https://models.acme.com/v1/models"
@@ -509,7 +509,13 @@ impl acp::Agent for MvpAgent {
Ok(self.auth_response_with_meta()) 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) auth_method::platform_for_method_id(&arguments.method_id)
{ {
self.authenticate_api_key_platform( self.authenticate_api_key_platform(
@@ -25,7 +25,16 @@ impl MvpAgent {
primary: &SamplingConfig, primary: &SamplingConfig,
) -> Result<(OaiCompatClient, String), acp::Error> { ) -> Result<(OaiCompatClient, String), acp::Error> {
let slug = self.resolve_session_summary_model(); 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 models = self.models_manager.models();
let endpoints = self.models_manager.endpoints(); let endpoints = self.models_manager.endpoints();
let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone(); 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()); .and_then(|v| v.as_object().cloned());
Ok(AuthenticateResponse::new().meta(meta)) Ok(AuthenticateResponse::new().meta(meta))
} }
/// `authenticate(<generic-oauth platform id>)`: interactive device-code
/// login for a `uses_oauth` platform carrying an `OAuthConfig` (xai-grok).
///
/// Uses a per-provider [`AuthManager`] scoped to the platform's `scope_key`
/// (NOT the primary Kimi manager) so the minted session is persisted under
/// its own `auth.json` scope, then triggers a catalog re-sync so the
/// platform's models appear (their bearer is resolved per-provider at fetch
/// / sampling time). The tokens are never logged.
pub(super) async fn authenticate_oauth_platform(
&self,
platform: kigi_models::PlatformId,
arguments: acp::AuthenticateRequest,
) -> Result<AuthenticateResponse, acp::Error> {
let method_id = arguments.method_id.clone();
let oauth = platform
.oauth()
.expect("oauth_platform() guarantees a device-code OAuthConfig");
let auth_meta = AuthRequestMeta::from_json(arguments.meta.as_ref());
tracing::info!(
method = method_id.0.as_ref(),
headless = auth_meta.headless,
reauth = auth_meta.reauth,
"auth: generic oauth device login",
);
let kigi_home = crate::util::kigi_home::kigi_home();
let auth_manager =
std::sync::Arc::new(crate::auth::AuthManager::new_oauth_provider(&kigi_home, oauth));
auth_manager.configure_refresher();
let flow_result = if !auth_meta.headless {
let (url_tx, url_rx) = tokio::sync::oneshot::channel();
let (code_tx, code_rx) = tokio::sync::mpsc::channel(1);
*self.auth_code_tx.borrow_mut() = Some(code_tx);
*self.auth_url_rx.borrow_mut() = Some(url_rx);
let result = crate::auth::run_oauth_provider_flow(
&auth_manager,
oauth,
auth_meta.reauth,
Some(crate::auth::AuthChannels {
url_tx: Some(url_tx),
code_rx,
}),
)
.await;
*self.auth_code_tx.borrow_mut() = None;
*self.auth_url_rx.borrow_mut() = None;
result
} else {
crate::auth::run_oauth_provider_flow(&auth_manager, oauth, auth_meta.reauth, None).await
};
let (_auth, _did_auth) = flow_result.map_err(|e| {
emit_login_span(false, method_id.0.as_ref(), None, Some("login_flow_failed"));
let mut err = acp::Error::auth_required();
err.message = e.to_string();
err
})?;
// Do NOT stamp this token onto the shared sampling_config: it authorizes
// ONLY this platform's models (api.x.ai/v1), not the primary session.
// The catalog re-sync below resolves it per-provider.
self.set_auth_method(method_id.clone());
self.models_manager.on_auth_changed().await;
emit_login_span(true, method_id.0.as_ref(), None, None);
Ok(self.auth_response_with_meta())
}
pub(crate) fn deployment_key(&self) -> Option<String> { pub(crate) fn deployment_key(&self) -> Option<String> {
self.cfg.borrow().endpoints.deployment_key.clone() self.cfg.borrow().endpoints.deployment_key.clone()
} }
@@ -593,16 +669,50 @@ impl MvpAgent {
); );
Ok(entry.clone()) Ok(entry.clone())
} }
/// Resolve the SESSION token for `model` by the model's OWN platform — the
/// single guard against the api_key-channel token leak.
///
/// An oauth-platform model (xai-grok) draws its session token from ITS OWN
/// process-global pool manager (build-on-demand from the on-disk grok token,
/// proactively refreshed), INDEPENDENT of the primary `auth_method`; when
/// that provider has no stored session the token is `None` — NEVER the
/// primary Kimi key. Every other model (first-party / Kimi) uses the primary
/// session manager, and only under a session-based auth method —
/// byte-identical to the pre-fix path. SECURITY: the resolved token is never
/// logged.
fn session_token_for_model(&self, model: &ModelEntry) -> Option<crate::auth::KimiAuth> {
if let Some(oauth) = model
.info()
.id
.as_deref()
.and_then(kigi_models::parse_managed_model_key)
.and_then(|(platform, _)| platform.oauth())
{
return crate::auth::oauth_registry::global_manager_for(
&crate::util::kigi_home::kigi_home(),
oauth,
)
.current_or_expired();
}
if self.is_session_based_auth() {
self.auth_manager.current_or_expired()
} else {
None
}
}
pub(crate) fn prepare_sampling_config_for_model( pub(crate) fn prepare_sampling_config_for_model(
&self, &self,
model: &ModelEntry, model: &ModelEntry,
origin_client: Option<crate::http::OriginClientInfo>, origin_client: Option<crate::http::OriginClientInfo>,
) -> SamplingConfig { ) -> SamplingConfig {
let session = if self.is_session_based_auth() { // Resolve the session token by the MODEL's platform, not the primary
self.auth_manager.current_or_expired() // auth method: an oauth-platform model (xai-grok) uses its OWN
} else { // pool-backed token (`None` — never the Kimi key — when the user has not
None // 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 has_session_key = session.is_some();
let mut credentials = resolve_credentials( let mut credentials = resolve_credentials(
model, model,
@@ -1158,6 +1158,72 @@ fn build_agent_with_auth(auth: crate::auth::KimiAuth) -> MvpAgent {
let cfg = AgentConfig::default(); let cfg = AgentConfig::default();
MvpAgent::new(gateway, &cfg, auth_manager, None).expect("valid test config") 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 /// Regression: boot-time plugin discovery is deferred past ACP
/// `initialize`, so the shared plugin registry starts empty. /// `initialize`, so the shared plugin registry starts empty.
/// `resolve_mcp_servers` reads that snapshot to merge plugin-contributed /// `resolve_mcp_servers` reads that snapshot to merge plugin-contributed
@@ -176,6 +176,7 @@ async fn handle_connection(ws: WebSocket, state: Arc<ServerState>, peer_addr: So
prefetch_models_blocking( prefetch_models_blocking(
&agent_config.endpoints, &agent_config.endpoints,
auth.as_ref(), auth.as_ref(),
&Default::default(),
fetch_auth, fetch_auth,
&platform_keys, &platform_keys,
) )
@@ -982,9 +982,20 @@ fn resolve_model_override_to_config(
} else { } else {
acp::ModelId::new(entry.info().model.clone()) 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 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); credentials.auth_type = subagent_auth_type(Some(&entry), &ctx.auth_method_id);
let resolved_auth_type = credentials.auth_type; let resolved_auth_type = credentials.auth_type;
let config = sampling_config_for_model(&entry, credentials, ctx.alpha_test_key.clone()); let config = sampling_config_for_model(&entry, credentials, ctx.alpha_test_key.clone());
@@ -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!(config.model, "kigi-4.5");
assert_eq!(model_id.0.as_ref(), "kigi-4.5"); assert_eq!(model_id.0.as_ref(), "kigi-4.5");
} }
/// Build an `Arc<AuthManager>` (primary Kimi) holding `key` as its live bearer.
/// The `TempDir` is returned so the caller keeps it alive.
fn kimi_primary_with_token(key: &str) -> (tempfile::TempDir, std::sync::Arc<crate::auth::AuthManager>) {
let dir = tempfile::tempdir().unwrap();
let manager = std::sync::Arc::new(crate::auth::AuthManager::new(
dir.path(),
crate::auth::KimiCodeConfig::default(),
));
manager.hot_swap(crate::auth::KimiAuth {
key: key.to_string(),
auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
(dir, manager)
}
/// LEAK 2 (subagent model-override): a grok (oauth-platform) override with a
/// Kimi primary must NEVER receive the primary Kimi session token as its
/// `api_key` — it draws grok's own pooled token (or `None`). Revert-to-red: the
/// pre-fix code passed `ctx.auth` (Kimi) straight to `resolve_credentials`, so
/// `config.api_key == "kimi-secret"` and this assertion fails.
#[tokio::test]
async fn subagent_override_grok_model_never_leaks_kimi_session_token() {
let (_kd, manager) = kimi_primary_with_token("kimi-secret");
let mut grok = test_model_entry("grok-4-latest");
grok.info.id = Some("xai-grok/grok-4-latest".to_string());
grok.info.base_url = "https://api.x.ai/v1".to_string();
let mut models = indexmap::IndexMap::new();
models.insert("grok".to_string(), grok);
let mut ctx = ctx_with_toggle(HashMap::new());
ctx.available_models = models;
ctx.auth = Some(crate::auth::KimiAuth {
key: "kimi-secret".to_string(),
auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
ctx.auth_manager = manager;
let (config, _model_id) =
resolve_model_override_to_config("grok", &ctx).expect("grok override resolves to a config");
assert_ne!(
config.api_key.as_deref(),
Some("kimi-secret"),
"a grok override must never receive the primary Kimi session token",
);
}
/// Byte-identical guard: a non-oauth override with a Kimi primary still resolves
/// to the primary session token — passes both before and after the fix (the
/// non-oauth path is unchanged).
#[tokio::test]
async fn subagent_override_non_oauth_model_still_gets_primary_token() {
let (_kd, manager) = kimi_primary_with_token("kimi-secret");
let mut entry = test_model_entry("kimi-k2-0905-preview");
entry.info.id = Some("moonshot-cn/kimi-k2".to_string());
let mut models = indexmap::IndexMap::new();
models.insert("k2".to_string(), entry);
let mut ctx = ctx_with_toggle(HashMap::new());
ctx.available_models = models;
ctx.auth = Some(crate::auth::KimiAuth {
key: "kimi-secret".to_string(),
auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
ctx.auth_manager = manager;
let (config, _model_id) =
resolve_model_override_to_config("k2", &ctx).expect("non-oauth override resolves to a config");
assert_eq!(
config.api_key.as_deref(),
Some("kimi-secret"),
"a non-oauth override must still receive the primary session token",
);
}
/// An unresolvable `AgentDefinition.model` pin (model absent from /// An unresolvable `AgentDefinition.model` pin (model absent from
/// `available_models`) falls through to inherit the parent model. /// `available_models`) falls through to inherit the parent model.
#[tokio::test] #[tokio::test]
@@ -13,15 +13,47 @@
use std::sync::Arc; use std::sync::Arc;
use crate::auth::kimi_oauth::{ use kigi_models::OAuthConfig;
DeviceAuthorization, DevicePollResult, poll_device_token, request_device_authorization,
}; use crate::auth::kimi_oauth::{DeviceAuthorization, DevicePollResult};
use crate::auth::{AuthChannels, AuthManager, AuthUrlInfo, AuthUrlMode, KimiAuth}; use crate::auth::{AuthChannels, AuthManager, AuthUrlInfo, AuthUrlMode, KimiAuth};
/// Extra wait added to the poll interval when the server answers `slow_down` /// Extra wait added to the poll interval when the server answers `slow_down`
/// (OAuth-standard device-flow backpressure). /// (OAuth-standard device-flow backpressure).
const SLOW_DOWN_INCREMENT_SECS: u64 = 5; 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<DeviceAuthorization> {
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<DevicePollResult> {
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. /// Outcome of one full poll loop over a single device authorization.
enum PollLoopOutcome { enum PollLoopOutcome {
/// Access token issued. /// Access token issued.
@@ -40,11 +72,29 @@ pub async fn run_device_code_login_channels(
host: &str, host: &str,
auth_manager: &Arc<AuthManager>, auth_manager: &Arc<AuthManager>,
channels: &mut Option<AuthChannels>, channels: &mut Option<AuthChannels>,
) -> 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<AuthManager>,
channels: &mut Option<AuthChannels>,
) -> 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<AuthManager>,
channels: &mut Option<AuthChannels>,
) -> anyhow::Result<(KimiAuth, bool)> { ) -> anyhow::Result<(KimiAuth, bool)> {
let interactive_tui = channels.is_some(); let interactive_tui = channels.is_some();
let mut channels = channels.take(); let mut channels = channels.take();
loop { loop {
let device_auth = request_device_authorization(host).await?; let device_auth = backend.request().await?;
let display_uri = device_auth.verification_uri_complete.clone(); let display_uri = device_auth.verification_uri_complete.clone();
if interactive_tui { if interactive_tui {
@@ -62,7 +112,7 @@ pub async fn run_device_code_login_channels(
prompt_on_stderr(&device_auth).await; 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) => { PollLoopOutcome::Done(auth) => {
let auth = auth_manager let auth = auth_manager
.update(*auth) .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 /// Poll the token endpoint until the user approves, the device code expires
/// (→ [`PollLoopOutcome::Restart`]), or the wire fails. /// (→ [`PollLoopOutcome::Restart`]), or the wire fails.
async fn complete_device_code_login( async fn complete_device_code_login(
host: &str, backend: &DeviceFlowBackend<'_>,
device_auth: &DeviceAuthorization, device_auth: &DeviceAuthorization,
) -> anyhow::Result<PollLoopOutcome> { ) -> anyhow::Result<PollLoopOutcome> {
let mut poll_interval = std::time::Duration::from_secs(device_auth.interval.max(1) as u64); 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 // Sleep first: an immediate poll on a fresh code only returns
// authorization_pending (and risks slow_down). // authorization_pending (and risks slow_down).
tokio::time::sleep(poll_interval).await; 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) => { DevicePollResult::Success(auth) => {
tracing::info!("auth: device login authorized"); tracing::info!("auth: device login authorized");
return Ok(PollLoopOutcome::Done(auth)); return Ok(PollLoopOutcome::Done(auth));
@@ -67,6 +67,37 @@ pub async fn run_auth_flow(
run_auth_flow_inner(auth_manager, kimi_code_config, reauth, false, channels).await 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<AuthManager>,
oauth: &'static kigi_models::OAuthConfig,
reauth: bool,
channels: Option<AuthChannels>,
) -> 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( async fn run_auth_flow_inner(
auth_manager: &Arc<AuthManager>, auth_manager: &Arc<AuthManager>,
_kimi_code_config: &KimiCodeConfig, _kimi_code_config: &KimiCodeConfig,
@@ -130,8 +130,9 @@ fn with_device_headers(
} }
/// Defend against control characters / non-https redirects from a /// Defend against control characters / non-https redirects from a
/// compromised or mis-configured OAuth host. /// compromised or mis-configured OAuth host. Shared with the generic
fn validate_verification_uri(uri: &str) -> anyhow::Result<()> { /// 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()) { if uri.chars().any(|c| c.is_ascii_control()) {
anyhow::bail!("Server returned invalid verification URI"); anyhow::bail!("Server returned invalid verification URI");
} }
@@ -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 /// Single field-assembly point for [`Self::new`]'s two construction paths
/// (inline `KIGI_AUTH` vs. on-disk `auth.json`), which differ only in the /// (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 /// threaded fields. One literal means a newly added field can't be silently
@@ -789,6 +842,13 @@ impl AuthManager {
&self.kimi_code_config &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. /// Handle notified after every successful token refresh.
/// ///
/// Used by [`ModelsManager`] to trigger model catalog recovery /// Used by [`ModelsManager`] to trigger model catalog recovery
+4 -1
View File
@@ -8,6 +8,8 @@ mod flow;
pub(crate) mod kimi_oauth; pub(crate) mod kimi_oauth;
pub(crate) mod manager; pub(crate) mod manager;
mod model; mod model;
pub(crate) mod oauth_device;
pub(crate) mod oauth_registry;
pub(crate) mod recovery; pub(crate) mod recovery;
pub(crate) mod refresh; pub(crate) mod refresh;
mod storage; mod storage;
@@ -17,7 +19,8 @@ pub(crate) use flow::try_ensure_session_noninteractive;
pub use flow::{ pub use flow::{
AuthChannels, AuthUrlInfo, AuthUrlMode, LogoutResult, ensure_authenticated, AuthChannels, AuthUrlInfo, AuthUrlMode, LogoutResult, ensure_authenticated,
ensure_authenticated_or_noninteractive, perform_logout, run_auth_flow, 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; mod meta;
pub use device::device_headers; pub use device::device_headers;
@@ -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<String>,
/// 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<String>,
#[serde(default)]
expires_in: Option<i64>,
#[serde(default)]
interval: Option<i64>,
}
#[derive(Deserialize, Default)]
struct OAuthErrorBody {
#[serde(default)]
error: Option<String>,
#[serde(default)]
error_description: Option<String>,
}
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<DeviceAuthorization> {
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<DevicePollResult> {
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::<TokenResponse>(&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<super::model::KimiAuth, RefreshError> {
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::<TokenResponse>(&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:?}"),
}
}
}
@@ -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<HashMap<&'static str, Arc<AuthManager>>> {
static POOL: OnceLock<Mutex<HashMap<&'static str, Arc<AuthManager>>>> = 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<AuthManager> {
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<AuthManager>>,
) -> Option<Arc<AuthManager>> {
if let Some((platform, _)) = kigi_models::parse_managed_model_key(managed_key)
&& let Some(oauth) = platform.oauth()
{
return Some(global_manager_for(kigi_home, oauth));
}
primary.cloned()
}
/// The SESSION token (the raw bearer/key string) that 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<AuthManager>>,
) -> Option<String> {
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<AuthManager>) {
let dir = tempfile::tempdir().unwrap();
let manager = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
manager.hot_swap(KimiAuth {
key: key.to_string(),
auth_mode: AuthMode::OAuth,
..KimiAuth::test_default()
});
(dir, manager)
}
fn xai_oauth() -> &'static kigi_models::OAuthConfig {
kigi_models::PlatformId::XaiGrok
.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/<model>` turn resolves to the process-global pooled xai
/// manager, NEVER the primary Kimi manager — the pool is the single source.
#[tokio::test]
async fn grok_model_resolves_to_pooled_xai_manager_not_kimi() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi))
.expect("grok model resolves to the pooled xai manager");
assert!(
!Arc::ptr_eq(&resolved, &kimi),
"grok model must NOT resolve to the Kimi manager"
);
assert!(
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), xai_oauth())),
"grok model must resolve to the process-global pooled xai manager"
);
}
/// Facet B guard: the resolver routes purely by the model's platform, with
/// no auth-method input — so even when the session's primary is a Kimi
/// (session) manager holding "kimi-tok", a grok model never resolves that
/// Kimi token.
#[tokio::test]
async fn grok_model_under_kimi_primary_never_yields_kimi_token() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-fast", Some(&kimi))
.expect("grok model resolves to its own pooled manager regardless of primary");
assert!(!Arc::ptr_eq(&resolved, &kimi));
assert_ne!(
resolved.current_or_expired().map(|a| a.key),
Some("kimi-tok".to_string()),
"the Kimi bearer must never be what a grok turn resolves"
);
}
/// Fail-fast: a grok key resolves to the pooled xai manager (never the Kimi
/// primary) even with no stored grok session in the pool — the pooled
/// manager then simply holds no token, so nothing (least of all the Kimi
/// bearer) is sent to api.x.ai.
#[tokio::test]
async fn grok_never_falls_back_to_kimi_primary() {
let (_kd, kimi) = primary_with_token("kimi-tok");
let home = tempfile::tempdir().unwrap();
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi))
.expect("grok routes to the pooled xai manager, not None");
assert!(
!Arc::ptr_eq(&resolved, &kimi),
"an OAuth platform must never fall back to the primary Kimi manager"
);
}
/// `session_key_for_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()),
);
}
}
@@ -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<dyn AuthSnapshot>,
cfg: &'static OAuthConfig,
}
impl GenericDeviceRefresher {
pub(crate) fn new(auth: Arc<dyn AuthSnapshot>, 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<RefreshOutcome> {
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<Option<KimiAuth>>,
disk: Mutex<Option<KimiAuth>>,
}
impl FakeSnapshot {
fn new(current: Option<KimiAuth>, disk: Option<KimiAuth>) -> Arc<Self> {
Arc::new(Self {
current: Mutex::new(current),
disk: Mutex::new(disk),
})
}
}
impl AuthSnapshot for FakeSnapshot {
fn current(&self) -> Option<KimiAuth> {
self.current
.lock()
.clone()
.filter(|a| !crate::auth::is_expired(a))
}
fn expired_auth(&self) -> Option<KimiAuth> {
self.current.lock().clone().filter(crate::auth::is_expired)
}
fn read_disk_auth(&self) -> Option<KimiAuth> {
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"));
}
}
@@ -1,3 +1,4 @@
mod generic_refresher;
mod kimi_refresher; mod kimi_refresher;
use std::sync::Arc; use std::sync::Arc;
@@ -6,6 +7,7 @@ use crate::auth::manager::AuthManager;
pub(crate) use crate::auth::manager::RefreshReason; pub(crate) use crate::auth::manager::RefreshReason;
use crate::auth::model::KimiAuth; use crate::auth::model::KimiAuth;
pub(crate) use generic_refresher::GenericDeviceRefresher;
pub(crate) use kimi_refresher::KimiRefresher; pub(crate) use kimi_refresher::KimiRefresher;
/// Read-only view of `AuthManager` for refreshers. Enforces the /// 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; 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<AuthManager>) -> Arc<dyn TokenRefresher> { pub(crate) fn build_refresher(auth_manager: Arc<AuthManager>) -> Arc<dyn TokenRefresher> {
match kigi_models::oauth_config_for_scope_key(auth_manager.scope()) {
Some(cfg) => {
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
Arc::new(GenericDeviceRefresher::new(snapshot, cfg))
}
None => {
let snapshot: Arc<dyn AuthSnapshot> = auth_manager; let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host())) Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host()))
}
}
} }
@@ -60,9 +60,12 @@ impl SessionActor {
stream_tool_calls: Some(sampling_config.stream_tool_calls), stream_tool_calls: Some(sampling_config.stream_tool_calls),
}); });
let existing = self.chat_state_handle.get_credentials().await; 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 let session_key = self
.auth_manager .auth_manager_for_model(&sampling_config.model)
.as_ref()
.and_then(|am| am.current_or_expired().map(|a| a.key)); .and_then(|am| am.current_or_expired().map(|a| a.key));
self.chat_state_handle self.chat_state_handle
.update_credentials(kigi_chat_state::Credentials { .update_credentials(kigi_chat_state::Credentials {
@@ -624,12 +624,25 @@ impl SessionActor {
let resolved_describe = self let resolved_describe = self
.resolve_aux_sampler_config(&self.image_description_model) .resolve_aux_sampler_config(&self.image_description_model)
.await; .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( crate::agent::config::finalize_image_describe_sampler_config(
resolved_describe, resolved_describe,
&active_session_config, &active_session_config,
Some(self.max_retries), 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| { let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| {
acp::Error::internal_error().data(format!( acp::Error::internal_error().data(format!(
"failed to build image-describe sampling client: {e}" "failed to build image-describe sampling client: {e}"
@@ -103,6 +103,30 @@ where
result 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<crate::auth::AuthManager>);
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<String> {
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<crate::auth::AuthManager>,
) -> kigi_sampler::SharedBearerResolver {
std::sync::Arc::new(AuthManagerBearerResolver(am))
}
impl SessionActor { impl SessionActor {
pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) { pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
let mcp_wait_start = std::time::Instant::now(); let mcp_wait_start = std::time::Instant::now();
@@ -184,6 +208,76 @@ impl SessionActor {
let auth_method = self.auth_method_id.load(); let auth_method = self.auth_method_id.load();
SessionTokenAuthGate::new(auth_method.as_deref(), byok, base_url) SessionTokenAuthGate::new(auth_method.as_deref(), byok, base_url)
} }
/// 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<std::sync::Arc<crate::auth::AuthManager>> {
// `model` is the bare routing slug; recover the managed catalog key
// (`{platform}/{model}`) so the platform — and thus its OAuth scope — is
// unambiguous. A bare / config / unlisted model yields no managed key
// and resolves to the primary.
let managed_key = self.managed_key_for_model(model);
crate::auth::oauth_registry::manager_for_model(
&crate::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<String> {
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 /// Emit a unified-log breadcrumb whenever the session-token refresh gate is
/// evaluated with an **`Unknown`** per-model BYOK status on a session-based /// evaluated with an **`Unknown`** per-model BYOK status on a session-based
/// method — the condition that (pre-fix) silently demoted live sessions to /// 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<crate::auth::AuthManager>);
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<String> {
self.0.current_or_expired().map(|a| a.key)
}
}
let cfg = self let cfg = self
.chat_state_handle .chat_state_handle
.get_sampling_config() .get_sampling_config()
@@ -273,6 +355,16 @@ impl SessionActor {
SessionTokenAuthGate::new(auth_method.as_deref(), model_facts.byok, &cfg.base_url); SessionTokenAuthGate::new(auth_method.as_deref(), model_facts.byok, &cfg.base_url);
let use_bearer_resolver = gate.active(); let use_bearer_resolver = gate.active();
self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url); 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 auth_scheme = model_facts.auth_scheme;
let mut extra_headers = cfg.extra_headers; let mut extra_headers = cfg.extra_headers;
crate::agent::config::inject_url_derived_headers( crate::agent::config::inject_url_derived_headers(
@@ -323,15 +415,7 @@ impl SessionActor {
idle_timeout_secs: None, idle_timeout_secs: None,
origin_client: self.origin_client.clone(), origin_client: self.origin_client.clone(),
attribution_callback: self.attribution_callback.clone(), attribution_callback: self.attribution_callback.clone(),
bearer_resolver: if use_bearer_resolver { bearer_resolver: inference_auth_manager.map(auth_manager_bearer_resolver),
self.auth_manager
.as_ref()
.map(|am| -> kigi_sampler::SharedBearerResolver {
std::sync::Arc::new(AuthManagerBearerResolver(am.clone()))
})
} else {
None
},
supports_backend_search: self.supports_backend_search.get(), supports_backend_search: self.supports_backend_search.get(),
compactions_remaining: self.compactions_remaining.get(), compactions_remaining: self.compactions_remaining.get(),
compaction_at_tokens: self.compaction_at_tokens.get(), compaction_at_tokens: self.compaction_at_tokens.get(),
@@ -461,10 +545,16 @@ impl SessionActor {
slug: &str, slug: &str,
) -> Option<kigi_sampler::SamplerConfig> { ) -> Option<kigi_sampler::SamplerConfig> {
let creds = self.chat_state_handle.get_credentials().await; let creds = self.chat_state_handle.get_credentials().await;
let session_key = self // Resolve the aux token by the aux model's OWN platform: a grok
.auth_manager // (oauth-platform) aux model draws its pooled grok token or `None` —
.as_ref() // NEVER the primary Kimi session token (which `resolve_credentials`
.and_then(|am| am.current_or_expired().map(|a| a.key.clone())); // 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 models = self.models_manager.models();
let endpoints = self.models_manager.endpoints(); let endpoints = self.models_manager.endpoints();
crate::agent::config::resolve_aux_model_sampling_config( crate::agent::config::resolve_aux_model_sampling_config(
@@ -491,6 +581,10 @@ impl SessionActor {
&active_session_config, &active_session_config,
Some(self.max_retries), 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 model = cfg.model.clone();
let client = kigi_sampler::SamplingClient::new(cfg) let client = kigi_sampler::SamplingClient::new(cfg)
.map_err(|e| { .map_err(|e| {
@@ -661,7 +755,17 @@ impl SessionActor {
)), )),
); );
} }
if auth_recovery_eligible && let Some(ref am) = self.auth_manager { // 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 { if am.try_recover_unauthorized().await {
tracing::info!( tracing::info!(
session_id = % self.session_info.id.0, session_id = % self.session_info.id.0,
@@ -685,6 +789,7 @@ impl SessionActor {
None, None,
); );
} }
}
if matches!(error.kind, SamplingErrorKind::IdleTimeout) { if matches!(error.kind, SamplingErrorKind::IdleTimeout) {
self.signals_handle().record_idle_timeout(); self.signals_handle().record_idle_timeout();
} }
@@ -846,14 +951,17 @@ impl SessionActor {
} }
/// Proactively refresh the auth token if near expiry. /// Proactively refresh the auth token if near expiry.
pub(super) async fn refresh_token_if_expired(&self) { pub(super) async fn refresh_token_if_expired(&self) {
if let Some(ref am) = self.auth_manager {
let creds = self.chat_state_handle.get_credentials().await;
let (model_id, base_url) = self let (model_id, base_url) = self
.chat_state_handle .chat_state_handle
.get_sampling_config() .get_sampling_config()
.await .await
.map(|c| (c.model, c.base_url)) .map(|c| (c.model, c.base_url))
.unwrap_or_default(); .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;
if self.auth_gate(&model_id, &base_url).active() if self.auth_gate(&model_id, &base_url).active()
&& let Ok(key) = am.get_valid_token().await && let Ok(key) = am.get_valid_token().await
{ {
@@ -975,3 +1083,47 @@ impl SessionActor {
.push_assistant_response(assistant_item); .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",
);
}
}
+37 -30
View File
@@ -6894,28 +6894,33 @@ pub(crate) mod tests {
#[test] #[test]
fn pending_menu_items_lists_interactive_methods_plus_quit() { fn pending_menu_items_lists_interactive_methods_plus_quit() {
let items = pending_menu_items(&fresh_user_auth_methods(), None); 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!( assert!(
matches!(&items[0], PendingMenuItem::Login { label } if label == "Kimi Code (OAuth)"), matches!(&items[0], PendingMenuItem::Login { label } if label == "Kimi Code (OAuth)"),
"row 0 must be the OAuth login, got {:?}", "row 0 must be the OAuth login, got {:?}",
items[0] 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!( assert_eq!(
items[1], items[2],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::MoonshotCn), target: PlatformLogin(kigi_shell::models::PlatformId::MoonshotCn),
label: "Moonshot Open Platform (API key \u{b7} moonshot.cn)".into(), label: "Moonshot Open Platform (API key \u{b7} moonshot.cn)".into(),
} }
); );
assert_eq!( assert_eq!(
items[2], items[3],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::MoonshotAi), target: PlatformLogin(kigi_shell::models::PlatformId::MoonshotAi),
label: "Moonshot Open Platform (API key \u{b7} moonshot.ai)".into(), label: "Moonshot Open Platform (API key \u{b7} moonshot.ai)".into(),
} }
); );
assert_eq!( assert_eq!(
items[3], items[4],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::OpenAi), target: PlatformLogin(kigi_shell::models::PlatformId::OpenAi),
label: "OpenAI (API key)".into(), 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" "new registry rows must appear in the picker with zero TUI changes"
); );
assert_eq!( assert_eq!(
items[4], items[5],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Anthropic), target: PlatformLogin(kigi_shell::models::PlatformId::Anthropic),
label: "Anthropic (API key)".into(), label: "Anthropic (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[5], items[6],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::DeepSeek), target: PlatformLogin(kigi_shell::models::PlatformId::DeepSeek),
label: "DeepSeek (API key)".into(), label: "DeepSeek (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[6], items[7],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Groq), target: PlatformLogin(kigi_shell::models::PlatformId::Groq),
label: "Groq (API key)".into(), label: "Groq (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[7], items[8],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Mistral), target: PlatformLogin(kigi_shell::models::PlatformId::Mistral),
label: "Mistral (API key)".into(), label: "Mistral (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[8], items[9],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Fireworks), target: PlatformLogin(kigi_shell::models::PlatformId::Fireworks),
label: "Fireworks AI (API key)".into(), label: "Fireworks AI (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[9], items[10],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Google), target: PlatformLogin(kigi_shell::models::PlatformId::Google),
label: "Google Gemini (API key)".into(), label: "Google Gemini (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[10], items[11],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::OpenRouter), target: PlatformLogin(kigi_shell::models::PlatformId::OpenRouter),
label: "OpenRouter (API key)".into(), label: "OpenRouter (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[11], items[12],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Together), target: PlatformLogin(kigi_shell::models::PlatformId::Together),
label: "Together AI (API key)".into(), label: "Together AI (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[12], items[13],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Cerebras), target: PlatformLogin(kigi_shell::models::PlatformId::Cerebras),
label: "Cerebras (API key)".into(), label: "Cerebras (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[13], items[14],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Nvidia), target: PlatformLogin(kigi_shell::models::PlatformId::Nvidia),
label: "NVIDIA NIM (API key)".into(), label: "NVIDIA NIM (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[14], items[15],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Vercel), target: PlatformLogin(kigi_shell::models::PlatformId::Vercel),
label: "Vercel AI Gateway (API key)".into(), label: "Vercel AI Gateway (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[15], items[16],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Xai), target: PlatformLogin(kigi_shell::models::PlatformId::Xai),
label: "xAI (Grok) (API key)".into(), label: "xAI (Grok) (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[16], items[17],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::QwenTokenPlan), target: PlatformLogin(kigi_shell::models::PlatformId::QwenTokenPlan),
label: "Qwen Token Plan (API key)".into(), label: "Qwen Token Plan (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[17], items[18],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::QwenTokenPlanCn), target: PlatformLogin(kigi_shell::models::PlatformId::QwenTokenPlanCn),
label: "Qwen Token Plan China (API key)".into(), label: "Qwen Token Plan China (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[18], items[19],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::KimiCoding), target: PlatformLogin(kigi_shell::models::PlatformId::KimiCoding),
label: "Kimi For Coding (API key)".into(), label: "Kimi For Coding (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[19], items[20],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Zai), target: PlatformLogin(kigi_shell::models::PlatformId::Zai),
label: "Z.AI (API key)".into(), label: "Z.AI (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[20], items[21],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::ZaiCodingCn), target: PlatformLogin(kigi_shell::models::PlatformId::ZaiCodingCn),
label: "Z.AI Coding China (API key)".into(), label: "Z.AI Coding China (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[21], items[22],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Xiaomi), target: PlatformLogin(kigi_shell::models::PlatformId::Xiaomi),
label: "Xiaomi MiMo (API key)".into(), label: "Xiaomi MiMo (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[22], items[23],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::XiaomiTokenPlanCn), target: PlatformLogin(kigi_shell::models::PlatformId::XiaomiTokenPlanCn),
label: "Xiaomi Token Plan China (API key)".into(), label: "Xiaomi Token Plan China (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[23], items[24],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::Minimax), target: PlatformLogin(kigi_shell::models::PlatformId::Minimax),
label: "MiniMax (API key)".into(), label: "MiniMax (API key)".into(),
} }
); );
assert_eq!( assert_eq!(
items[24], items[25],
PendingMenuItem::ApiKey { PendingMenuItem::ApiKey {
target: PlatformLogin(kigi_shell::models::PlatformId::MinimaxCn), target: PlatformLogin(kigi_shell::models::PlatformId::MinimaxCn),
label: "MiniMax China (API key)".into(), 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. // The non-interactive methods must never appear as rows.
let byok = kigi_shell::agent::auth_method::build_auth_methods( let byok = kigi_shell::agent::auth_method::build_auth_methods(
kigi_shell::agent::auth_method::AuthMethodsBuildInputs { kigi_shell::agent::auth_method::AuthMethodsBuildInputs {
@@ -7080,8 +7085,8 @@ pub(crate) mod tests {
); );
assert_eq!( assert_eq!(
pending_menu_items(&byok.methods, None).len(), pending_menu_items(&byok.methods, None).len(),
26, 27,
"xai.api_key / cached_token must not add rows" "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 /// 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(&kimi_only));
assert!(!login_picker_has_choice(&[])); 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. /// and 'l' keeps selecting the first (OAuth) row.
#[test] #[test]
fn welcome_pending_arrows_select_moonshot_row() { fn welcome_pending_arrows_select_moonshot_row() {
@@ -7102,7 +7107,9 @@ pub(crate) mod tests {
app.auth_methods = fresh_user_auth_methods(); app.auth_methods = fresh_user_auth_methods();
app.auth_state = AuthState::Pending { error: None }; app.auth_state = AuthState::Pending { error: None };
app.welcome_prompt_focused = false; 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));
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)); let outcome = app.handle_input(&key_event(KeyCode::Enter, KeyModifiers::NONE));