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

M1/F1: Kimi Code OAuth device-code flow

Browse Source 
Replace the xAI OAuth stack with the Kimi device authorization grant:
- kimi_oauth.rs wire layer (device_authorization + token poll + refresh
  against kigi_env::oauth_host(); client_id per PRD; retryable statuses
  429/5xx with backoff; expired_token restarts authorization)
- X-Msh-Device-{Name,Model,Id} headers; device_id minted uuid4-hex at
  ~/.kigi/device_id (0600)
- Storage: system keyring service `kigi`, entry `oauth/kimi-code`
  (macOS/Windows native backends), atomic-file fallback under ~/.kigi;
  official client's keyring/~/.kimi never touched
- Refresh manager: 60s tick, threshold max(300, expires_in*0.5),
  401-tombstone keyed by rejected refresh token with 300s cooldown and
  rotation auto-clear, cross-process lock with sibling-adoption
  triple-check, sleep/wake forced refresh
- Deleted xAI machinery: enterprise OIDC (PKCE/JWKS/teams), devbox login,
  external auth provider, JWT tier gating + subscription paywall stack,
  X-XAI-Token-Auth marker headers, ZDR gates, /user enrichment
- kigi login / TUI /login both drive the device flow; login-host display
  now derives from kigi_env::oauth_host()
- 264 auth unit/wiremock tests; live contract probe of
  auth.kimi.com/api/oauth/device_authorization matches the wire shapes

Gates: check/clippy --all-targets clean, fmt, deny ok, kigi-shell lib
5131 tests green.
This commit is contained in:
ZacharyZhang-NY
2026-07-17 07:37:29 -04:00
parent d6c20fc13f
commit 021b82443d
117 changed files with 4052 additions and 19900 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-shell/src/agent/app.rs
+5 -78
View File
@@ -20,7 +20,7 @@ use crate::agent::config::{Config as AgentConfig, ModelEntry};
use crate::agent::init::{bootstrap, exit_on_config_error};
use crate::agent::models::{ModelFetchAuth, prefetch_models_blocking};
use crate::agent::mvp_agent::MvpAgent;
use crate::auth::{AuthManager, AuthMode, GrokAuth, run_auth_flow};
use crate::auth::{AuthManager, AuthMode, KimiAuth, run_auth_flow};
use crate::util::kigi_home;
use dirs;
@@ -399,76 +399,6 @@ pub async fn run_stdio_agent(
result
}
async fn migrate_devbox_auth_if_legacy(
auth: Option<GrokAuth>,
agent_config: &AgentConfig,
) -> Option<GrokAuth> {
let auth = auth?;
if !crate::auth::devbox_login::is_devbox_environment() || auth.auth_mode != AuthMode::WebLogin {
return Some(auth);
}
info!("Devbox legacy auth detected, attempting migration to OIDC");
kigi_log::unified_log::info(
"devbox legacy auth migration: starting",
None,
Some(serde_json::json!({
"user_id": auth.user_id,
"auth_mode": format!("{:?}", auth.auth_mode),
})),
);
// save + remove_scope are two non-atomic writes to auth.json (no lock). Safe
// at startup: no concurrent writer yet, and `lookup_auth` prefers the primary
// scope if a reader sees the intermediate state.
let migration_auth_manager = agent_config.create_auth_manager();
let new_auth = match crate::auth::devbox_login::mint_devbox_auth(&migration_auth_manager).await
{
Ok(new_auth) => new_auth,
Err(e) => {
tracing::warn!(error = ?e, "devbox legacy auth migration: devbox login helper call failed, continuing with legacy auth");
kigi_log::unified_log::error(
"devbox legacy auth migration: mint failed",
None,
Some(serde_json::json!({ "error": e.to_string() })),
);
return Some(auth);
}
};
match migration_auth_manager
.save_without_enrichment(new_auth)
.await
{
Ok(saved_auth) => {
if let Err(e) = migration_auth_manager.remove_scope(crate::auth::LEGACY_AUTH_SCOPE) {
tracing::warn!(error = ?e, "Failed to remove legacy auth scope entry (non-fatal)");
}
kigi_log::unified_log::info(
"devbox legacy auth migration: succeeded",
None,
Some(serde_json::json!({
"user_id": saved_auth.user_id,
"has_refresh_token": saved_auth.refresh_token.is_some(),
"expires_at": saved_auth.expires_at.map(|e| e.to_rfc3339()),
"auth_mode": format!("{:?}", saved_auth.auth_mode),
})),
);
info!(user_id = %saved_auth.user_id, "Devbox legacy auth migrated to OIDC successfully");
Some(saved_auth)
}
Err(e) => {
tracing::warn!(error = ?e, "devbox legacy auth migration: failed to save new auth, continuing with legacy");
kigi_log::unified_log::error(
"devbox legacy auth migration: save failed",
None,
Some(serde_json::json!({ "error": e.to_string() })),
);
Some(auth)
}
}
}
/// Run the agent in leader mode, accepting IPC connections from multiple clients.
///
/// Startup sequence:
@@ -681,14 +611,11 @@ pub async fn run_leader(
// The IPC server is already accepting connections. Clients that send ACP
// messages during this window receive a `leader_starting` error and can retry.
let ctx = &agent_config.grok_com_config;
let ctx = &agent_config.kimi_code_config;
// Never interactive: a detached leader has no TTY (forcing OAuth here hung BYOK).
let auth: Option<GrokAuth> = crate::auth::try_ensure_session_noninteractive(ctx).await;
let auth: Option<KimiAuth> = crate::auth::try_ensure_session_noninteractive(ctx).await;
// ── Phase 6b: Legacy devbox auth migration ─────────────────────────────
let auth: Option<GrokAuth> = migrate_devbox_auth_if_legacy(auth, &agent_config).await;
let auth_for_prefetch: Option<GrokAuth> = auth.clone();
let auth_for_prefetch: Option<KimiAuth> = auth.clone();
let endpoints_for_prefetch = agent_config.endpoints.clone();
let fetch_auth_for_prefetch = ModelFetchAuth::resolve(&endpoints_for_prefetch, auth.is_some());
// The shared pair helper owns the remote_fetch gate for both halves, so a
@@ -894,7 +821,7 @@ pub async fn run_leader(
if let Some(home) = dirs::home_dir() {
watch_paths.push(home.join(".claude.json"));
}
let auth_scope = agent_config.grok_com_config.auth_scope();
let auth_scope = agent_config.kimi_code_config.auth_scope();
// Gated on user_kigi_home() so a cwd-relative .kigi/auth.json is never
// read as the user auth store when no home resolves.
let initial_auth_key_hash = kigi_config::user_kigi_home()
crates/codegen/kigi-shell/src/agent/auth_method.rs
+176 -759
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/agent/config.rs
+40 -361
View File
@@ -1,5 +1,5 @@
use crate::agent::auth_method::ModelByok;
use crate::auth::{AuthManager, GrokComConfig, OidcAuthConfig};
use crate::auth::{AuthManager, KimiCodeConfig};
use crate::remote::DEFAULT_CONTEXT_WINDOW;
use crate::{config::StorageMode, sampling::ApiBackend, tools::config::ShellToolsetConfig};
use agent_client_protocol as acp;
@@ -1108,7 +1108,7 @@ pub struct Config {
/// Warnings from `[model.*]` parsing; surfaced by `grok inspect`.
#[serde(skip)]
pub model_override_warnings: Vec<super::config_model_override_parse::ModelOverrideWarning>,
pub grok_com_config: GrokComConfig,
pub kimi_code_config: KimiCodeConfig,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub shortcuts: Option<toml::Value>,
/// Written by the client via `config_toml_edit`; absorbed so it isn't
@@ -1178,9 +1178,9 @@ pub struct Config {
#[serde(default, skip_serializing)]
pub managed_mcps: crate::config::ManagedMcpsConfig,
/// `[auth]` alias — consumed by `expand_auth_alias` before serde.
/// Typed as `GrokComConfig` (same schema) so sub-field typos are caught.
/// Typed as `KimiCodeConfig` (same schema) so sub-field typos are caught.
#[serde(default, skip_serializing)]
pub auth: Option<GrokComConfig>,
pub auth: Option<KimiCodeConfig>,
/// `[desktop]` section — owned by grok-desktop (Electron app), opaque to the CLI agent.
#[serde(default, skip_serializing)]
pub desktop: Option<toml::Value>,
@@ -1525,7 +1525,7 @@ impl Default for Config {
auto_mode: AutoModeConfig::default(),
config_models: IndexMap::new(),
model_override_warnings: Vec::new(),
grok_com_config: GrokComConfig::default(),
kimi_code_config: KimiCodeConfig::default(),
shortcuts: None,
hints: None,
ui: UiConfig::default(),
@@ -1620,13 +1620,12 @@ impl Config {
}
Ok(())
}
/// Build an `AuthManager` with the configured proxy URL applied.
/// Build an `AuthManager` for this configuration.
pub fn create_auth_manager(&self) -> AuthManager {
AuthManager::new(
&crate::util::kigi_home::kigi_home(),
self.grok_com_config.clone(),
self.kimi_code_config.clone(),
)
.with_proxy_base_url(&self.endpoints.proxy_url())
}
/// Deserialize the merged `base` document, also returning the ignored key
/// paths whose top-level key appears in `user_config`. Paths outside it
@@ -1679,12 +1678,6 @@ impl Config {
}
config.config_models = config_models;
config.model_override_warnings = model_override_warnings;
if config.grok_com_config.oidc.is_none() {
config.grok_com_config.oidc = OidcAuthConfig::from_env();
}
if config.grok_com_config.oidc.is_none() && config.grok_com_config.oauth2.is_none() {
config.grok_com_config.oauth2 = crate::auth::OAuth2ProviderConfig::from_env();
}
if config.client_version.is_none() {
config.client_version = Self::default().client_version;
}
@@ -1818,16 +1811,16 @@ impl Config {
self.resolve_runtime_fields(&ctx);
crate::util::config::set_remote_campaigns_from_settings(self.remote_settings.as_ref());
}
/// If the TOML contains `[auth]`, copy its contents under `[grok_com_config]`.
/// `[grok_com_config]` takes precedence if both are present (explicit wins).
/// If the TOML contains `[auth]`, copy its contents under `[kimi_code_config]`.
/// `[kimi_code_config]` takes precedence if both are present (explicit wins).
///
/// This lets customers write the shorter `[auth.oidc]` instead of `[grok_com_config.oidc]`.
/// This lets customers write the shorter `[auth.oidc]` instead of `[kimi_code_config.oidc]`.
fn expand_auth_alias(raw_config: &toml::Value) -> toml::Value {
let mut config = raw_config.clone();
if let toml::Value::Table(ref mut table) = config
&& let Some(auth) = table.remove("auth")
{
if let Some(gcc) = table.get_mut("grok_com_config") {
if let Some(gcc) = table.get_mut("kimi_code_config") {
if let (toml::Value::Table(gcc_table), toml::Value::Table(auth_table)) =
(gcc, &auth)
{
@@ -1836,7 +1829,7 @@ impl Config {
}
}
} else {
table.insert("grok_com_config".to_owned(), auth);
table.insert("kimi_code_config".to_owned(), auth);
}
}
config
@@ -2320,18 +2313,6 @@ impl Config {
.default(true)
.resolve()
}
/// Resolve whether to use grok's default OAuth2 (xAI auth.x.ai).
///
/// Enterprise OIDC (`oidc` in config.toml) always wins — this only gates
/// the default xAI OAuth2 fallback when no enterprise OIDC is configured.
///
/// Priority: `--oauth` > KIGI_OAUTH_ENABLED env > default (true = OAuth).
pub fn resolve_grok_oauth(&self, cli_oidc: Option<bool>) -> Resolved<bool> {
BoolFlag::env("KIGI_OAUTH_ENABLED")
.cli(cli_oidc)
.default(true)
.resolve()
}
/// Resolve whether to spawn the per-`Ready`-client transport
/// liveness pollers and the session-actor `StatusDispatcher`.
///
@@ -3918,42 +3899,6 @@ pub fn resolve_credentials(model: &ModelEntry, session_key: Option<&str>) -> Res
auth_scheme,
}
}
/// `disable_api_key_auth` at the credential seam: swap a first-party xAI API
/// key for the IdP session (absent => request fails => forces login). BYOK
/// (non-xAI `base_url`) is untouched; no-op when the switch is off.
pub fn enforce_disable_api_key_auth(
creds: &mut ResolvedCredentials,
disable_api_key_auth: bool,
session_key: Option<&str>,
) {
if disable_api_key_auth
&& creds.auth_type == kigi_chat_state::AuthType::ApiKey
&& crate::util::is_first_party_xai_url(&creds.base_url)
{
creds.auth_type = kigi_chat_state::AuthType::SessionToken;
creds.api_key = session_key.map(str::to_owned);
kigi_log::unified_log::debug(
"auth: kill switch blocked a first-party API key at the credential seam",
None,
Some(serde_json::json!(
{ "replaced_with_session" : session_key.is_some(), "base_url" : creds
.base_url, }
)),
);
}
}
/// Resolve credentials for an auxiliary sampling path (web search, image
/// description) with the first-party API-key kill switch applied, so these
/// paths honor `disable_api_key_auth` exactly like the main chat path.
fn resolve_credentials_enforced(
entry: &ModelEntry,
session_key: Option<&str>,
disable_api_key_auth: bool,
) -> ResolvedCredentials {
let mut credentials = resolve_credentials(entry, session_key);
enforce_disable_api_key_auth(&mut credentials, disable_api_key_auth, session_key);
credentials
}
/// Try to resolve credentials for a model by loading the effective config.
/// Returns `None` (with a warning) if config loading, parsing, or model
@@ -3971,12 +3916,7 @@ pub fn try_resolve_model_credentials(
.ok()?;
let models = resolve_model_list(&cfg, None);
let entry = find_model_by_id(&models, model_id)?;
let mut credentials = resolve_credentials(entry, session_key);
enforce_disable_api_key_auth(
&mut credentials,
cfg.grok_com_config.api_key_auth_disabled(),
session_key,
);
let credentials = resolve_credentials(entry, session_key);
Some(credentials)
}
/// Per-model auth facts (BYOK status + auth scheme) from one effective-config
@@ -4045,13 +3985,12 @@ pub fn resolve_aux_model_sampling_config(
models: &IndexMap<String, ModelEntry>,
endpoints: &EndpointsConfig,
session_key: Option<&str>,
disable_api_key_auth: bool,
alpha_test_key: Option<String>,
client_version: Option<String>,
) -> Option<SamplerConfig> {
let catalog_entry = find_model_by_id(models, model_id).cloned();
if let Some(entry) = &catalog_entry {
let credentials = resolve_credentials_enforced(entry, session_key, disable_api_key_auth);
let credentials = resolve_credentials(entry, session_key);
let sampler = sampling_config_for_model(
entry,
credentials,
@@ -4108,7 +4047,7 @@ pub fn resolve_aux_model_sampling_config(
env_key: None,
api_base_url: None,
};
let credentials = resolve_credentials_enforced(&entry, session_key, disable_api_key_auth);
let credentials = resolve_credentials(&entry, session_key);
let sampler = sampling_config_for_model(
&entry,
credentials,
@@ -4240,11 +4179,7 @@ pub fn sampling_config_for_model(
/// URL-derived header logic at the shell boundary so callers downstream see a
/// single homogenous header bag.
///
/// * cli-chat-proxy bases get `X-XAI-Token-Auth` and
/// `x-authenticateresponse` headers (mirrors the inline match in the legacy
/// `sampling::Client::new` on `is_cli_chat_proxy_url`).
/// * With the optional non-production feature, matching first-party hosts may
/// get an extra access header from the corresponding key argument.
/// * First-party bases get the client-mode header.
///
/// Existing entries are never overwritten so callers can pre-set a value.
pub fn inject_url_derived_headers(
@@ -4253,12 +4188,6 @@ pub fn inject_url_derived_headers(
base_url: &str,
) {
if crate::util::is_cli_chat_proxy_url(base_url) {
headers
.entry("X-XAI-Token-Auth".to_string())
.or_insert_with(|| "xai-grok-cli".to_string());
headers
.entry("x-authenticateresponse".to_string())
.or_insert_with(|| "authenticate-response".to_string());
headers
.entry(crate::http::CLIENT_MODE_HEADER.to_string())
.or_insert_with(|| crate::http::process_client_mode().to_string());
@@ -4289,7 +4218,6 @@ pub fn resolve_model_to_sampling_config(
fn resolve_hidden_default_web_search_sampling_config(
model_id: &str,
session_key: Option<&str>,
disable_api_key_auth: bool,
alpha_test_key: Option<String>,
client_version: Option<String>,
endpoints: &EndpointsConfig,
@@ -4331,7 +4259,7 @@ fn resolve_hidden_default_web_search_sampling_config(
env_key: None,
api_base_url: None,
};
let credentials = resolve_credentials_enforced(&entry, session_key, disable_api_key_auth);
let credentials = resolve_credentials(&entry, session_key);
sampling_config_for_model(
&entry,
credentials,
@@ -4345,13 +4273,12 @@ pub fn resolve_web_search_sampling_config(
model_id: &str,
models: &IndexMap<String, ModelEntry>,
session_key: Option<&str>,
disable_api_key_auth: bool,
alpha_test_key: Option<String>,
client_version: Option<String>,
endpoints: &EndpointsConfig,
) -> Option<SamplerConfig> {
let resolved = if let Some(entry) = find_model_by_id(models, model_id).cloned() {
let credentials = resolve_credentials_enforced(&entry, session_key, disable_api_key_auth);
let credentials = resolve_credentials(&entry, session_key);
Some(sampling_config_for_model(
&entry,
credentials,
@@ -4364,7 +4291,6 @@ pub fn resolve_web_search_sampling_config(
Some(resolve_hidden_default_web_search_sampling_config(
model_id,
session_key,
disable_api_key_auth,
alpha_test_key,
client_version,
endpoints,
@@ -4643,28 +4569,21 @@ reasoning_effort = "low"
}
}
#[test]
fn inject_url_derived_headers_adds_proxy_headers_for_cli_chat_proxy_url() {
fn inject_url_derived_headers_adds_client_mode_for_first_party_url() {
let mut headers = IndexMap::new();
inject_url_derived_headers(
&mut headers,
None,
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
);
assert_eq!(
headers.get("X-XAI-Token-Auth").map(String::as_str),
Some("xai-grok-cli")
);
assert_eq!(
headers.get("x-authenticateresponse").map(String::as_str),
Some("authenticate-response")
);
assert!(headers.get(crate::http::CLIENT_MODE_HEADER).is_some());
assert!(headers.get("X-XAI-Token-Auth").is_none());
}
#[test]
fn inject_url_derived_headers_skips_proxy_headers_for_external_url() {
fn inject_url_derived_headers_skips_headers_for_external_url() {
let mut headers = IndexMap::new();
inject_url_derived_headers(&mut headers, None, "https://api.x.ai/v1");
assert!(headers.get("X-XAI-Token-Auth").is_none());
assert!(headers.get("x-authenticateresponse").is_none());
inject_url_derived_headers(&mut headers, None, "https://api.example.com/v1");
assert!(headers.get(crate::http::CLIENT_MODE_HEADER).is_none());
}
#[test]
fn inject_url_derived_headers_preserves_caller_extra_headers() {
@@ -4679,24 +4598,6 @@ reasoning_effort = "low"
headers.get("x-custom-byok").map(String::as_str),
Some("value")
);
assert_eq!(
headers.get("X-XAI-Token-Auth").map(String::as_str),
Some("xai-grok-cli")
);
}
#[test]
fn inject_url_derived_headers_does_not_overwrite_existing_entries() {
let mut headers = IndexMap::new();
headers.insert("X-XAI-Token-Auth".to_string(), "caller-set".to_string());
inject_url_derived_headers(
&mut headers,
None,
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
);
assert_eq!(
headers.get("X-XAI-Token-Auth").map(String::as_str),
Some("caller-set"),
);
}
#[test]
fn parses_toolset_overrides() {
@@ -4834,7 +4735,6 @@ reasoning_effort = "low"
crate::models::default_web_search_model(),
&IndexMap::new(),
Some("session-token"),
false,
None,
None,
&endpoints,
@@ -4891,51 +4791,14 @@ reasoning_effort = "low"
None,
),
);
let resolved = resolve_aux_model_sampling_config(
"grok-build",
&catalog,
&endpoints,
None,
false,
None,
None,
)
.expect("override entry has an API key, so resolution succeeds");
let resolved =
resolve_aux_model_sampling_config("grok-build", &catalog, &endpoints, None, None, None)
.expect("override entry has an API key, so resolution succeeds");
assert_eq!(resolved.model, "v9m-rl-learnability-tp8");
assert_eq!(resolved.base_url, "https://vendor.example/v1");
assert_eq!(resolved.api_key.as_deref(), Some("vendor-key"));
}
#[test]
fn web_search_disable_api_key_auth_swaps_first_party_key_for_session() {
let endpoints = EndpointsConfig::default();
let mut models = IndexMap::new();
models.insert(
"ws-model".to_string(),
test_model_entry(
"ws-model",
"https://api.x.ai/v1",
Some("first-party-key"),
None,
None,
),
);
let resolved = resolve_web_search_sampling_config(
"ws-model",
&models,
Some("session-token"),
true,
None,
None,
&endpoints,
)
.expect("web search model should resolve");
assert_eq!(
resolved.api_key.as_deref(),
Some("session-token"),
"first-party API key must be swapped for the session token when disabled"
);
}
#[test]
fn parses_model_api_key() {
let raw_config: toml::Value = toml::from_str(
r#"
@@ -5351,12 +5214,9 @@ reasoning_effort = "low"
config.base_url,
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url
);
assert_eq!(
config
.extra_headers
.get("X-XAI-Token-Auth")
.map(String::as_str),
Some("xai-grok-cli")
assert!(
config.extra_headers.get("X-XAI-Token-Auth").is_none(),
"the xAI token-auth marker header must be gone"
);
}
/// Regression: without a session key, `resolve_credentials` falls through
@@ -5376,75 +5236,6 @@ reasoning_effort = "low"
auth_scheme: Default::default(),
}
}
/// `disable_api_key_auth` kill switch (Claude `forceLoginMethod` parity).
#[test]
fn enforce_disable_api_key_auth_blocks_first_party_only() {
use kigi_chat_state::AuthType;
let mut creds = api_key_creds("https://api.x.ai/v1");
enforce_disable_api_key_auth(&mut creds, false, Some("session-jwt"));
assert_eq!(creds.auth_type, AuthType::ApiKey);
assert_eq!(creds.api_key.as_deref(), Some("xai-secret"));
let mut creds = api_key_creds("https://api.x.ai/v1");
enforce_disable_api_key_auth(&mut creds, true, Some("session-jwt"));
assert_eq!(creds.auth_type, AuthType::SessionToken);
assert_eq!(creds.api_key.as_deref(), Some("session-jwt"));
let mut creds = api_key_creds("https://api.x.ai/v1");
enforce_disable_api_key_auth(&mut creds, true, None);
assert_eq!(creds.auth_type, AuthType::SessionToken);
assert_eq!(creds.api_key, None);
let mut creds = api_key_creds("https://api.example.com/v1");
enforce_disable_api_key_auth(&mut creds, true, Some("session-jwt"));
assert_eq!(creds.auth_type, AuthType::ApiKey);
assert_eq!(creds.api_key.as_deref(), Some("xai-secret"));
let mut creds = ResolvedCredentials {
auth_type: AuthType::SessionToken,
..api_key_creds("https://api.x.ai/v1")
};
enforce_disable_api_key_auth(&mut creds, true, Some("session-jwt"));
assert_eq!(creds.auth_type, AuthType::SessionToken);
}
/// Regression for the OVERRIDE_MODEL kill-switch bypass: a first-party model
/// with its own api_key resolves to `ApiKey` (priority 1, beating the
/// session), and the kill switch — now applied inside
/// `try_resolve_model_credentials` — swaps it for the session token. BYOK
/// (non-x.ai) own keys are preserved. (`try_resolve_model_credentials`
/// loads global config, so this exercises its resolve + enforce core.)
#[test]
fn try_resolve_model_credentials_swaps_first_party_own_key_under_kill_switch() {
use kigi_chat_state::AuthType;
let entry = test_model_entry(
"m",
"https://api.x.ai/v1",
Some("xai-model-key"),
None,
None,
);
let mut creds = resolve_credentials(&entry, Some("session-jwt"));
assert_eq!(
creds.auth_type,
AuthType::ApiKey,
"own key wins over session"
);
assert_eq!(creds.api_key.as_deref(), Some("xai-model-key"));
enforce_disable_api_key_auth(&mut creds, true, Some("session-jwt"));
assert_eq!(
creds.auth_type,
AuthType::SessionToken,
"swapped under switch"
);
assert_eq!(creds.api_key.as_deref(), Some("session-jwt"));
let byok = test_model_entry(
"b",
"https://api.example.com/v1",
Some("sk-byok"),
None,
None,
);
let mut byok_creds = resolve_credentials(&byok, Some("session-jwt"));
enforce_disable_api_key_auth(&mut byok_creds, true, Some("session-jwt"));
assert_eq!(byok_creds.auth_type, AuthType::ApiKey);
assert_eq!(byok_creds.api_key.as_deref(), Some("sk-byok"));
}
#[test]
fn x_api_key_auth_scheme_flows_from_config_to_sampler() {
let mut model = test_model_entry(
@@ -6631,124 +6422,16 @@ reasoning_effort = "low"
let info = ModelInfo::from_config(&entry);
assert_eq!(info.inference_idle_timeout_secs, Some(120));
}
/// The `[auth]` alias and the explicit `[kimi_code_config]` table both
/// deserialize (the auth block currently carries no per-deployment
/// options; the alias machinery is retained for future knobs).
#[test]
fn auth_alias_maps_to_grok_com_config() {
let raw: toml::Value = toml::from_str(
r#"
[auth.oidc]
issuer = "https://example.okta.com"
client_id = "test-id"
"#,
)
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse");
let oidc = cfg.grok_com_config.oidc.expect("oidc should be set");
assert_eq!(oidc.issuer, "https://example.okta.com");
assert_eq!(oidc.client_id, "test-id");
}
#[test]
fn grok_com_config_still_works() {
let raw: toml::Value = toml::from_str(
r#"
[grok_com_config.oidc]
issuer = "https://example.okta.com"
client_id = "test-id"
"#,
)
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse");
let oidc = cfg.grok_com_config.oidc.expect("oidc should be set");
assert_eq!(oidc.issuer, "https://example.okta.com");
}
/// `disable_api_key_auth` plumbs through the `[auth]` alias, and absent
/// means None (opt-in knob, zero impact by default).
#[test]
fn disable_api_key_auth_parses_from_auth_alias() {
let absent = Config::new_from_toml_cfg(&toml::from_str("").unwrap()).unwrap();
assert_eq!(absent.grok_com_config.disable_api_key_auth, None);
let raw: toml::Value = toml::from_str(
r#"
[auth]
disable_api_key_auth = true
"#,
)
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse");
assert_eq!(cfg.grok_com_config.disable_api_key_auth, Some(true));
}
/// `force_login_team_uuid` parses a string (pin), array (any-of), or `[]`
/// (fail closed); absent => None.
#[test]
fn force_login_team_uuid_parses_string_and_array() {
use crate::auth::ForceLoginTeam;
let absent = Config::new_from_toml_cfg(&toml::from_str("").unwrap()).unwrap();
assert_eq!(absent.grok_com_config.force_login_team_uuid, None);
let raw: toml::Value = toml::from_str(
r#"
[auth]
force_login_team_uuid = "team-abc"
"#,
)
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse");
assert_eq!(
cfg.grok_com_config.force_login_team_uuid,
Some(ForceLoginTeam::Single("team-abc".into())),
);
let raw: toml::Value = toml::from_str(
r#"
[grok_com_config]
force_login_team_uuid = ["team-a", "team-b"]
"#,
)
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse");
assert_eq!(
cfg.grok_com_config.force_login_team_uuid,
Some(ForceLoginTeam::AnyOf(vec![
"team-a".into(),
"team-b".into()
])),
);
let raw: toml::Value = toml::from_str(
r#"
[auth]
force_login_team_uuid = []
"#,
)
.unwrap();
let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse");
assert_eq!(
cfg.grok_com_config.force_login_team_uuid,
Some(ForceLoginTeam::AnyOf(vec![])),
);
}
/// Pinning a team via `force_login_team_uuid` implies API-key auth is
/// disabled even without an explicit `disable_api_key_auth` (team
/// membership can't be verified from a bare API key, so it needs IdP login).
#[test]
fn force_login_team_uuid_implies_api_key_auth_disabled() {
use crate::auth::{ForceLoginTeam, GrokComConfig};
let base = GrokComConfig {
disable_api_key_auth: None,
force_login_team_uuid: None,
..GrokComConfig::default()
};
assert!(!base.api_key_auth_disabled());
assert!(
GrokComConfig {
disable_api_key_auth: Some(true),
..base.clone()
}
.api_key_auth_disabled()
);
assert!(
GrokComConfig {
force_login_team_uuid: Some(ForceLoginTeam::Single("team-x".into())),
..base
}
.api_key_auth_disabled()
);
fn auth_alias_and_kimi_code_config_tables_parse() {
for body in ["[auth]\n", "[kimi_code_config]\n"] {
let raw: toml::Value = toml::from_str(body).unwrap();
let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse");
assert_eq!(cfg.kimi_code_config.auth_scope(), "oauth/kimi-code");
}
}
fn resolve_models_from_toml(
toml_str: &str,
@@ -8662,11 +8345,7 @@ agent_type = "cursor"
persistent_shell = true
[shortcuts]
ctrl_k = "search"
[grok_com_config]
token_header = "test"
[auth.oidc]
issuer = "https://sso.corp.com"
client_id = "abc123"
[kimi_code_config]
[storage]
cleanup_ttl_days = 7
[permission]
crates/codegen/kigi-shell/src/agent/feedback_client.rs
+26 -32
View File
@@ -320,14 +320,14 @@ pub struct FeedbackClient {
http: reqwest::Client,
client: reqwest_middleware::ClientWithMiddleware,
base_url: String,
credentials: crate::util::grok_auth_credentials::GrokAuthCredentials,
credentials: crate::util::kigi_auth_credentials::KigiAuthCredentials,
session_id: Option<String>,
}
impl FeedbackClient {
pub fn new(base_url: impl Into<String>, user_token: Option<String>) -> Self {
let http = crate::http::shared_client();
let credentials = crate::util::grok_auth_credentials::GrokAuthCredentials::new(user_token);
let credentials = crate::util::kigi_auth_credentials::KigiAuthCredentials::new(user_token);
let client = Self::build_middleware_client(&http, &credentials);
Self {
http,
@@ -361,7 +361,7 @@ impl FeedbackClient {
base_url: impl Into<String>,
user_token: Option<String>,
) -> Self {
let credentials = crate::util::grok_auth_credentials::GrokAuthCredentials::new(user_token);
let credentials = crate::util::kigi_auth_credentials::KigiAuthCredentials::new(user_token);
let client = Self::build_middleware_client(&http, &credentials);
Self {
http,
@@ -398,7 +398,7 @@ impl FeedbackClient {
fn build_middleware_client(
http: &reqwest::Client,
credentials: &crate::util::grok_auth_credentials::GrokAuthCredentials,
credentials: &crate::util::kigi_auth_credentials::KigiAuthCredentials,
) -> reqwest_middleware::ClientWithMiddleware {
let provider = Self::make_auth_provider(credentials);
// max_retries=0: the middleware stamps the auth header but does NOT
@@ -415,7 +415,7 @@ impl FeedbackClient {
}
fn make_auth_provider(
credentials: &crate::util::grok_auth_credentials::GrokAuthCredentials,
credentials: &crate::util::kigi_auth_credentials::KigiAuthCredentials,
) -> Arc<dyn kigi_auth::AuthCredentialProvider> {
if let Some(am) = credentials.auth_manager() {
Arc::new(
@@ -1118,7 +1118,7 @@ mod forbidden_tests {
#[cfg(test)]
mod auth_refresh_tests {
use super::*;
use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig};
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
use axum::{Router, routing::get};
use chrono::{Duration, Utc};
use std::net::SocketAddr;
@@ -1156,14 +1156,14 @@ mod auth_refresh_tests {
let (addr, _server) = start_server(router).await;
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
am.hot_swap(KimiAuth {
key: "fresh-from-auth-manager".into(),
auth_mode: AuthMode::ApiKey,
create_time: Utc::now(),
user_id: "user-42".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
let client = FeedbackClient::new(
@@ -1208,14 +1208,14 @@ mod auth_refresh_tests {
let (addr, _server) = start_server(router).await;
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let fresh = GrokAuth {
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
let fresh = KimiAuth {
key: "fresh-from-auth-manager".into(),
auth_mode: AuthMode::ApiKey,
create_time: Utc::now(),
user_id: "user-42".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
};
am.hot_swap(fresh);
@@ -1247,16 +1247,14 @@ mod auth_refresh_tests {
_reason: crate::auth::refresh::RefreshReason,
) -> crate::auth::refresh::RefreshOutcome {
self.calls.fetch_add(1, Ordering::SeqCst);
crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth {
crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth {
key: "fresh-from-refresher".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
create_time: Utc::now(),
user_id: "user-42".into(),
refresh_token: Some("rt-fresh".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
oidc_issuer: Some("https://issuer.example".into()),
oidc_client_id: Some("test-client".into()),
..GrokAuth::test_default()
..KimiAuth::test_default()
}))
}
}
@@ -1266,34 +1264,30 @@ mod auth_refresh_tests {
#[tokio::test]
async fn try_refresh_credentials_picks_up_disk_rotation_without_hitting_idp() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig::default();
let cfg = KimiCodeConfig::default();
let scope = cfg.auth_scope();
let am = Arc::new(AuthManager::new(dir.path(), cfg));
// In-memory: stale token (the one the server rejected).
am.hot_swap(GrokAuth {
am.hot_swap(KimiAuth {
key: "stale-rejected".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
create_time: Utc::now() - Duration::hours(2),
user_id: "user-42".into(),
refresh_token: Some("rt-stale".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
oidc_issuer: Some("https://issuer.example".into()),
oidc_client_id: Some("test-client".into()),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
// Disk: a sibling already rotated to a fresh token.
let disk_auth = GrokAuth {
let disk_auth = KimiAuth {
key: "fresh-from-sibling-on-disk".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
create_time: Utc::now(),
user_id: "user-42".into(),
refresh_token: Some("rt-fresh".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
oidc_issuer: Some("https://issuer.example".into()),
oidc_client_id: Some("test-client".into()),
..GrokAuth::test_default()
..KimiAuth::test_default()
};
let mut store = std::collections::BTreeMap::new();
store.insert(scope, disk_auth);
@@ -1329,15 +1323,15 @@ mod auth_refresh_tests {
#[tokio::test]
async fn try_refresh_credentials_returns_false_on_terminal_failure() {
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
// LegacySession: no refresh_token, no recovery possible.
am.hot_swap(GrokAuth {
am.hot_swap(KimiAuth {
key: "legacy-rejected".into(),
auth_mode: AuthMode::WebLogin,
auth_mode: AuthMode::OAuth,
create_time: Utc::now() - Duration::days(60),
user_id: "user-42".into(),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
let client = FeedbackClient::new("http://example/v1", Some("legacy-rejected".into()))
crates/codegen/kigi-shell/src/agent/init.rs
+3 -3
View File
@@ -79,7 +79,7 @@ fn resolve_config(cfg: &AgentConfig, auth_manager: &AuthManager) -> AgentConfig
// thread the result into `cfg.remote_settings` skip this entirely.
if cfg.remote_settings.is_none()
&& let Some(handle) =
crate::agent::models::start_early_prefetch(Some(cfg.grok_com_config.clone()))
crate::agent::models::start_early_prefetch(Some(cfg.kimi_code_config.clone()))
{
match handle.join() {
Ok(result) => {
@@ -103,9 +103,9 @@ fn resolve_config(cfg: &AgentConfig, auth_manager: &AuthManager) -> AgentConfig
{
cfg.storage_mode = StorageMode::resolve(None, cfg.remote_settings.as_ref());
}
// Writeback talks to the code backend; requires grok.com auth.
// Writeback talks to the code backend; requires a Kimi Code session.
if cfg.storage_mode == StorageMode::Writeback
&& !auth_manager.current().is_some_and(|a| a.is_xai_auth())
&& !auth_manager.current().is_some_and(|a| a.is_session_auth())
{
tracing::info!("Writeback is disabled: requires auth with grok.com");
cfg.storage_mode = StorageMode::Local;
crates/codegen/kigi-shell/src/agent/mod.rs
-1
View File
@@ -18,7 +18,6 @@ pub mod server;
pub mod session_config;
pub mod session_registry_client;
pub(crate) mod subagent;
pub(crate) mod subscription_check;
pub(crate) mod update_chunk_merge;
pub use mvp_agent::MvpAgent;
crates/codegen/kigi-shell/src/agent/models.rs
+20 -18
View File
@@ -10,7 +10,7 @@ use chrono::{DateTime, Duration as ChronoDuration, Utc};
use indexmap::IndexMap;
use crate::agent::config::{self, ModelEntry, resolve_credentials, sampling_config_for_model};
use crate::auth::{AuthManager, GrokAuth, GrokComConfig};
use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig};
use crate::remote::{FetchModelsResult, fetch_models_blocking};
use crate::sampling::SamplerConfig as SamplingConfig;
use globset::{Glob, GlobSet, GlobSetBuilder};
@@ -151,7 +151,7 @@ struct Inner {
impl Default for ModelsManager {
fn default() -> Self {
let kigi_home = crate::util::kigi_home::kigi_home();
let auth_manager = Arc::new(AuthManager::new(&kigi_home, GrokComConfig::default()));
let auth_manager = Arc::new(AuthManager::new(&kigi_home, KimiCodeConfig::default()));
Self::new(
None,
IndexMap::new(),
@@ -1397,7 +1397,7 @@ fn build_prefetched_map(
/// Fetch remote models. Checks disk cache first; persists after fetch.
pub(crate) fn prefetch_models_blocking(
endpoints: &config::EndpointsConfig,
auth: Option<&GrokAuth>,
auth: Option<&KimiAuth>,
fetch_auth: ModelFetchAuth,
) -> Option<IndexMap<String, ModelEntry>> {
prefetch_models_blocking_gated(
@@ -1414,7 +1414,7 @@ pub(crate) fn prefetch_models_blocking(
/// decisions cannot disagree mid-startup.
pub(crate) fn prefetch_models_and_settings_blocking(
endpoints: &config::EndpointsConfig,
auth: Option<&GrokAuth>,
auth: Option<&KimiAuth>,
fetch_auth: ModelFetchAuth,
) -> (
Option<IndexMap<String, ModelEntry>>,
@@ -1441,7 +1441,7 @@ pub(crate) fn prefetch_models_and_settings_blocking(
/// knob once for both halves.
fn prefetch_models_blocking_gated(
endpoints: &config::EndpointsConfig,
auth: Option<&GrokAuth>,
auth: Option<&KimiAuth>,
fetch_auth: ModelFetchAuth,
remote_fetch_enabled: bool,
) -> Option<IndexMap<String, ModelEntry>> {
@@ -1499,12 +1499,12 @@ pub struct EarlyPrefetchResult {
pub type EarlyPrefetchHandle = std::thread::JoinHandle<EarlyPrefetchResult>;
struct PrefetchEnv {
auth: Option<GrokAuth>,
auth: Option<KimiAuth>,
endpoints: config::EndpointsConfig,
model_fetch_auth: ModelFetchAuth,
}
fn resolve_prefetch_env_with_auth(auth: Option<GrokAuth>) -> Option<PrefetchEnv> {
fn resolve_prefetch_env_with_auth(auth: Option<KimiAuth>) -> Option<PrefetchEnv> {
let _timer = crate::instrumentation_timer!("startup.early_prefetch_launch");
// Config-aware (not env-only) so the prefetch can't leak the bearer to api.x.ai.
let mut endpoints = config::EndpointsConfig::from_effective_config();
@@ -1529,7 +1529,7 @@ fn resolve_prefetch_env_with_auth(auth: Option<GrokAuth>) -> Option<PrefetchEnv>
/// `deployment_key` would re-arm the prefetch — and with it the `/v1/settings`
/// fetch and the deployment-config sync on the prefetch thread.
fn resolve_prefetch_env_from_parts(
auth: Option<GrokAuth>,
auth: Option<KimiAuth>,
endpoints: config::EndpointsConfig,
remote_fetch_enabled: bool,
) -> Option<PrefetchEnv> {
@@ -1554,9 +1554,9 @@ fn resolve_prefetch_env_from_parts(
})
}
fn resolve_prefetch_env(grok_com_config: Option<GrokComConfig>) -> Option<PrefetchEnv> {
fn resolve_prefetch_env(kimi_code_config: Option<KimiCodeConfig>) -> Option<PrefetchEnv> {
let kigi_home = crate::util::kigi_home::kigi_home();
let auth_manager = AuthManager::new(&kigi_home, grok_com_config.unwrap_or_default());
let auth_manager = AuthManager::new(&kigi_home, kimi_code_config.unwrap_or_default());
let auth = auth_manager.current();
resolve_prefetch_env_with_auth(auth)
}
@@ -1566,7 +1566,7 @@ fn resolve_prefetch_env(grok_com_config: Option<GrokComConfig>) -> Option<Prefet
/// When the caller has already obtained valid credentials (e.g. via
/// `try_ensure_fresh_auth`), pass them here to avoid re-reading stale cached
/// credentials from disk.
pub fn start_early_prefetch_with_auth(auth: Option<GrokAuth>) -> Option<EarlyPrefetchHandle> {
pub fn start_early_prefetch_with_auth(auth: Option<KimiAuth>) -> Option<EarlyPrefetchHandle> {
let env = resolve_prefetch_env_with_auth(auth)?;
Some(spawn_prefetch_thread(env))
}
@@ -1575,8 +1575,10 @@ pub fn start_early_prefetch_with_auth(auth: Option<GrokAuth>) -> Option<EarlyPre
///
/// Convenience wrapper that reads cached auth from disk. Prefer
/// `start_early_prefetch_with_auth` when you have pre-resolved credentials.
pub fn start_early_prefetch(grok_com_config: Option<GrokComConfig>) -> Option<EarlyPrefetchHandle> {
let env = resolve_prefetch_env(grok_com_config)?;
pub fn start_early_prefetch(
kimi_code_config: Option<KimiCodeConfig>,
) -> Option<EarlyPrefetchHandle> {
let env = resolve_prefetch_env(kimi_code_config)?;
Some(spawn_prefetch_thread(env))
}
@@ -1969,7 +1971,7 @@ pub(crate) fn validate_selectable(
/// Async wrapper around `prefetch_models_blocking`.
pub(crate) async fn fetch_models_async(
endpoints: config::EndpointsConfig,
auth: Option<GrokAuth>,
auth: Option<KimiAuth>,
fetch_auth: ModelFetchAuth,
) -> Option<IndexMap<String, ModelEntry>> {
tokio::task::spawn_blocking(move || {
@@ -1991,7 +1993,7 @@ mod tests {
// Use a temp dir so AuthManager finds no credentials — ensures
// refresh_async bails at the auth check without needing a tokio runtime.
let tmp = std::env::temp_dir().join("grok-test-models-manager");
let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default()));
let auth_manager = Arc::new(AuthManager::new(&tmp, KimiCodeConfig::default()));
ModelsManager::new(
None,
IndexMap::new(),
@@ -2233,7 +2235,7 @@ mod tests {
#[test]
fn current_reasoning_effort_seeded_from_config() {
let tmp = std::env::temp_dir().join("grok-test-models-manager-seed");
let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default()));
let auth_manager = Arc::new(AuthManager::new(&tmp, KimiCodeConfig::default()));
let mut cfg = config::Config::default();
cfg.models.default_reasoning_effort = Some(ReasoningEffort::Xhigh);
let mgr = ModelsManager::new(
@@ -2401,7 +2403,7 @@ mod tests {
// The internal getters read those derived fields.
let tmp = std::env::temp_dir().join("grok-test-models-manager-menu-only");
let auth_manager = Arc::new(AuthManager::new(&tmp, GrokComConfig::default()));
let auth_manager = Arc::new(AuthManager::new(&tmp, KimiCodeConfig::default()));
let mgr = ModelsManager::new(
None,
catalog,
@@ -3222,7 +3224,7 @@ mod tests {
};
assert!(
resolve_prefetch_env_from_parts(
Some(GrokAuth::test_default()),
Some(KimiAuth::test_default()),
endpoints.clone(),
false,
)
crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs
+18 -269
View File
@@ -48,31 +48,6 @@ impl acp::Agent for MvpAgent {
);
});
kigi_workspace::trust::migrate_legacy_hook_trust();
if let Some(auth) = self.auth_manager.current() {
let user_id = auth.user_id.trim();
let needs_user_info = user_id.is_empty()
|| user_id.eq_ignore_ascii_case("unknown");
kigi_log::unified_log::info(
"auth init user_info check",
None,
Some(
serde_json::json!(
{ "user_id" : user_id, "needs_user_info" : needs_user_info,
"key_prefix" : crate ::auth::token_suffix(& auth.key),
"rt_prefix" : auth.refresh_token.as_deref().map(crate
::auth::token_suffix), }
),
),
);
if needs_user_info && let Err(e) = self.auth_manager.update(auth).await {
tracing::warn!(
"Failed to refresh user info from proxy during new_session: {}", e
);
}
}
if !self.tier_allowed.get() && let Some(auth) = self.auth_manager.current() {
self.enforce_grok_code_access(&auth).await;
}
self.maybe_sync_bundle_in_background(false);
let mut client_type = arguments
.meta
@@ -186,8 +161,7 @@ impl acp::Agent for MvpAgent {
),
),
);
if !self.cfg.borrow().grok_com_config.api_key_auth_disabled()
&& auth_method::read_xai_api_key_env().is_err()
if auth_method::read_xai_api_key_env().is_err()
&& let Some(api_key) = crate::auth::read_api_key(
&crate::util::kigi_home::kigi_home(),
)
@@ -200,33 +174,8 @@ impl acp::Agent for MvpAgent {
None,
);
}
let disable_api_key_auth = self
.cfg
.borrow()
.grok_com_config
.api_key_auth_disabled();
{
let cfg = self.cfg.borrow();
let gc = &cfg.grok_com_config;
if disable_api_key_auth || gc.force_login_team_uuid.is_some() {
kigi_log::unified_log::info(
"auth: enterprise login policy active",
None,
Some(
serde_json::json!(
{ "force_login_team_uuid" : gc.force_login_team_uuid.as_ref()
.map(| t | format!("{t:?}")), "disable_api_key_auth_knob" :
gc.disable_api_key_auth, "api_key_auth_disabled" :
disable_api_key_auth, }
),
),
);
}
}
let has_external_api_key = auth_method::should_advertise_xai_api_key(
disable_api_key_auth,
self.models_manager.models().values(),
);
let has_external_api_key =
auth_method::should_advertise_xai_api_key(self.models_manager.models().values());
let init_has_current = self.auth_manager.current().is_some();
let init_is_expired = self.auth_manager.is_expired();
kigi_log::unified_log::info(
@@ -267,58 +216,11 @@ impl acp::Agent for MvpAgent {
);
}
}
let (
login_label,
has_auth_provider,
has_enterprise_oidc,
enterprise_oidc_issuer,
) = {
let cfg = self.cfg.borrow();
let issuer = cfg.grok_com_config.oidc.as_ref().map(|o| o.issuer.clone());
(
cfg.grok_com_config.auth_provider_label.clone(),
cfg.grok_com_config.auth_provider_command.is_some(),
cfg.grok_com_config.oidc.is_some(),
issuer,
)
};
if has_enterprise_oidc {
let issuer = enterprise_oidc_issuer
.as_deref()
.expect(
"enterprise_oidc_issuer must be Some when has_enterprise_oidc is true",
);
tracing::info!(
issuer = % issuer, "auth: advertising enterprise OIDC auth method",
);
kigi_log::unified_log::info(
"auth: advertising enterprise OIDC auth method",
None,
Some(serde_json::json!({ "issuer" : issuer })),
);
} else {
tracing::info!(
label = ? login_label, has_auth_provider,
"auth: advertising grok.com auth method",
);
}
let preferred_method = self.cfg.borrow().grok_com_config.preferred_method;
let has_external_api_key = match preferred_method {
Some(crate::auth::PreferredAuthMethod::Oidc) => false,
_ => has_external_api_key,
};
let has_cached_token = match preferred_method {
Some(crate::auth::PreferredAuthMethod::ApiKey) => false,
_ => has_cached_token,
};
tracing::info!("auth: advertising Kimi Code device login auth method");
let built = auth_method::build_auth_methods(auth_method::AuthMethodsBuildInputs {
has_external_api_key,
has_cached_token,
has_enterprise_oidc,
enterprise_oidc_issuer: enterprise_oidc_issuer.as_deref(),
login_label: login_label.as_deref(),
has_auth_provider_command: has_auth_provider,
preferred_method,
login_label: None,
});
let auth_methods = built.methods;
kigi_log::unified_log::info(
@@ -329,9 +231,8 @@ impl acp::Agent for MvpAgent {
{ "kigi_home" : crate ::util::kigi_home::kigi_home().display()
.to_string(), "HOME" : std::env::var("HOME").unwrap_or_else(| _ |
"(unset)".into()), "has_external_api_key" : has_external_api_key,
"disable_api_key_auth" : disable_api_key_auth, "has_cached_token" :
has_cached_token, "has_enterprise_oidc" : has_enterprise_oidc,
"init_has_current" : init_has_current, "init_is_expired" :
"has_cached_token" :
has_cached_token, "init_has_current" : init_has_current, "init_is_expired" :
init_is_expired, "auth_mode" : self.auth_manager.current().map(| a |
format!("{:?}", a.auth_mode)), "methods" : auth_methods.iter().map(|
m | m.id().0.as_ref()).collect::< Vec < _ >> (),
@@ -438,39 +339,8 @@ impl acp::Agent for MvpAgent {
None,
Some(serde_json::json!({ "method" : arguments.method_id.0.as_ref() })),
);
if let Some(preferred) = self.cfg.borrow().grok_com_config.preferred_method {
let kind = auth_method::AuthMethodKind::from_id(&arguments.method_id);
let allowed = match preferred {
crate::auth::PreferredAuthMethod::ApiKey => kind.is_api_key(),
crate::auth::PreferredAuthMethod::Oidc => kind.is_session_based(),
};
if !allowed {
let msg = match preferred {
crate::auth::PreferredAuthMethod::ApiKey => {
auth_method::PREFERRED_API_KEY_UNAVAILABLE
}
crate::auth::PreferredAuthMethod::Oidc => {
"preferred_method=oidc; API-key auth is not allowed."
}
};
emit_login_span(
false,
arguments.method_id.0.as_ref(),
None,
Some("preferred_method_mismatch"),
);
return Err(acp::Error::auth_required().data(msg));
}
}
match arguments.method_id.0.as_ref() {
auth_method::XAI_API_KEY_METHOD_ID => {
if self.cfg.borrow().grok_com_config.api_key_auth_disabled() {
emit_login_span(false, "api_key", None, Some("disabled_by_admin"));
return Err(
acp::Error::auth_required()
.data("API-key auth is disabled by your administrator."),
);
}
let mut sampling_config = self.sampling_config.borrow_mut();
if sampling_config.api_key.is_none() {
if let Ok(api_key) = auth_method::read_xai_api_key_env() {
@@ -516,82 +386,23 @@ impl acp::Agent for MvpAgent {
return self
.authenticate(
acp::AuthenticateRequest::new(
acp::AuthMethodId::new(auth_method::OIDC_METHOD_ID),
acp::AuthMethodId::new(auth_method::KIGI_COM_METHOD_ID),
)
.meta(arguments.meta),
)
.await;
}
let current_auth = self.auth_manager.current();
let has_current = current_auth.is_some();
let has_current = self.auth_manager.current().is_some();
let is_expired = self.auth_manager.is_expired();
let is_devbox = crate::auth::devbox_login::is_devbox_environment();
let is_legacy = current_auth
.as_ref()
.is_some_and(|a| a.auth_mode == crate::auth::AuthMode::WebLogin);
kigi_log::unified_log::info(
"auth cached_token check",
None,
Some(
serde_json::json!(
{ "has_current" : has_current, "is_expired" : is_expired,
"is_devbox" : is_devbox, "is_legacy" : is_legacy, }
{ "has_current" : has_current, "is_expired" : is_expired, }
),
),
);
let pin_blocks_oidc_mint = matches!(
self.cfg.borrow().grok_com_config.preferred_method, Some(crate
::auth::PreferredAuthMethod::ApiKey)
);
if is_devbox && is_legacy && !pin_blocks_oidc_mint {
kigi_log::unified_log::info(
"auth cached_token: devbox legacy migration starting",
None,
None,
);
match crate::auth::devbox_login::mint_devbox_auth(&self.auth_manager)
.await
{
Ok(new_auth) => {
match self
.auth_manager
.save_without_enrichment(new_auth)
.await
{
Ok(_) => {
if let Err(e) = self
.auth_manager
.remove_scope(crate::auth::LEGACY_AUTH_SCOPE)
{
tracing::warn!(
error = ? e,
"auth: failed to remove legacy scope (non-fatal)"
);
}
kigi_log::unified_log::info(
"auth cached_token: devbox legacy migration succeeded",
None,
None,
);
}
Err(e) => {
kigi_log::unified_log::warn(
"auth cached_token: devbox migration save failed",
None,
Some(serde_json::json!({ "error" : e.to_string() })),
);
}
}
}
Err(e) => {
kigi_log::unified_log::warn(
"auth cached_token: devbox mint failed, will reject legacy token",
None,
Some(serde_json::json!({ "error" : format!("{e}") })),
);
}
}
}
let Some(auth) = self.auth_manager.current() else {
let message = if self.auth_manager.is_expired() {
"Session expired, re-authentication required"
@@ -610,34 +421,8 @@ impl acp::Agent for MvpAgent {
.authenticate_after_cached_token_unavailable(arguments)
.await;
};
if auth.auth_mode == crate::auth::AuthMode::WebLogin {
tracing::info!("auth: rejecting legacy WebLogin token");
kigi_log::unified_log::warn(
"auth cached_token legacy rejected",
None,
Some(
serde_json::json!(
{ "auth_mode" : format!("{:?}", auth.auth_mode) }
),
),
);
self.auth_manager.clear_in_memory();
if let Err(e) = self
.auth_manager
.remove_scope(crate::auth::LEGACY_AUTH_SCOPE)
{
tracing::warn!(
error = ? e,
"auth: failed to remove legacy scope during WebLogin rejection (non-fatal)"
);
}
return self
.authenticate_after_cached_token_unavailable(arguments)
.await;
}
self.refresh_remote_settings(&auth).await;
self.emit_settings_update_notification();
self.enforce_grok_code_access(&auth).await;
self.maybe_sync_bundle_in_background(false);
{
let mut sampling_config = self.sampling_config.borrow_mut();
@@ -660,13 +445,12 @@ impl acp::Agent for MvpAgent {
self.maybe_fetch_post_auth_settings().await;
Ok(self.auth_response_with_meta())
}
auth_method::KIGI_COM_METHOD_ID | auth_method::OIDC_METHOD_ID => {
let grok_ctx = self.auth_manager.grok_com_config();
auth_method::KIGI_COM_METHOD_ID => {
let kimi_ctx = self.auth_manager.kimi_code_config().clone();
let auth_meta = AuthRequestMeta::from_json(arguments.meta.as_ref());
tracing::info!(
method = arguments.method_id.0.as_ref(), headless = auth_meta
.headless, reauth = auth_meta.reauth, use_oauth = auth_meta
.use_oauth, "auth: inline auth flow",
.headless, reauth = auth_meta.reauth, "auth: inline auth flow",
);
kigi_log::unified_log::info(
"auth: inline auth flow",
@@ -674,31 +458,13 @@ impl acp::Agent for MvpAgent {
Some(
serde_json::json!(
{ "method" : arguments.method_id.0.as_ref(), "headless" :
auth_meta.headless, "reauth" : auth_meta.reauth, "use_oauth"
: auth_meta.use_oauth, }
auth_meta.headless, "reauth" : auth_meta.reauth, }
),
),
);
if auth_meta.reauth {
let _ = self.auth_manager.clear();
}
let cli_oauth = auth_meta.use_oauth.then_some(true);
let use_oidc = self.cfg.borrow().resolve_grok_oauth(cli_oauth);
tracing::debug!(
resolved = use_oidc.value, source = ? use_oidc.source,
"auth: method resolved"
);
kigi_log::unified_log::debug(
"auth: method resolved",
None,
Some(
serde_json::json!(
{ "use_oidc" : use_oidc.value, "source" : format!("{:?}",
use_oidc.source), }
),
),
);
let login_override = auth_meta.login_override();
let (auth, _did_auth) = if !auth_meta.headless {
let (url_tx, url_rx) = tokio::sync::oneshot::channel();
let (code_tx, code_rx) = tokio::sync::mpsc::channel(1);
@@ -706,14 +472,13 @@ impl acp::Agent for MvpAgent {
*self.auth_url_rx.borrow_mut() = Some(url_rx);
let result = crate::auth::run_auth_flow_with_stderr_bridge(
&self.auth_manager,
grok_ctx,
&kimi_ctx,
crate::auth::AuthChannels {
url_tx: Some(url_tx),
code_rx,
},
auth_meta.reauth,
auth_meta.force_interactive,
login_override,
)
.await;
*self.auth_code_tx.borrow_mut() = None;
@@ -722,12 +487,9 @@ impl acp::Agent for MvpAgent {
} else {
crate::auth::run_auth_flow(
&self.auth_manager,
grok_ctx,
&kimi_ctx,
auth_meta.reauth,
None,
None,
None,
login_override,
)
.await
}
@@ -757,11 +519,7 @@ impl acp::Agent for MvpAgent {
self.auth_manager.hot_swap(auth.clone());
self.refresh_remote_settings(&auth).await;
self.emit_settings_update_notification();
self.enforce_grok_code_access(&auth).await;
self.maybe_sync_bundle_in_background(false);
tokio::task::spawn_local(
crate::managed_config::post_login_sync(Some(auth.clone())),
);
self.set_auth_method(arguments.method_id.clone());
self.models_manager.on_auth_changed().await;
if crate::agent::chat_modes::process_chat_mode_enabled() {
@@ -959,10 +717,7 @@ impl acp::Agent for MvpAgent {
.session_registry_client()
.map(|client| crate::session::persistence::RegistryGeneratedTitleSync {
client,
suppress_for_zdr: self
.auth_manager
.current_or_expired()
.is_some_and(|a| a.is_zdr_team()),
suppress_for_zdr: false,
});
crate::session::persistence::new(
&session_info,
@@ -1239,10 +994,7 @@ impl acp::Agent for MvpAgent {
.session_registry_client()
.map(|client| crate::session::persistence::RegistryGeneratedTitleSync {
client,
suppress_for_zdr: self
.auth_manager
.current_or_expired()
.is_some_and(|a| a.is_zdr_team()),
suppress_for_zdr: false,
});
let (persistence_info, persistence) = crate::session::persistence::load_light(
&session_info,
@@ -2576,9 +2328,6 @@ impl acp::Agent for MvpAgent {
crate::extensions::billing::handle(self, &args).await
}
"x.ai/share_session" => crate::extensions::share::handle(self, &args).await,
"x.ai/privacy/setCodingDataRetention" => {
crate::extensions::privacy::handle(self, &args).await
}
"x.ai/rollout/survey" => {
crate::extensions::rollout::handle(self, &args).await
}
crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs
+30 -129
View File
@@ -28,10 +28,9 @@ impl MvpAgent {
let session_key = self.auth_manager.current_or_expired().map(|a| a.key.clone());
let models = self.models_manager.models();
let endpoints = self.models_manager.endpoints();
let (disable_api_key_auth, alpha_test_key, client_version) = {
let (alpha_test_key, client_version) = {
let cfg = self.cfg.borrow();
(
cfg.grok_com_config.api_key_auth_disabled(),
cfg.endpoints.alpha_test_key.clone(),
cfg.client_version.clone(),
)
@@ -41,7 +40,6 @@ impl MvpAgent {
&models,
&endpoints,
session_key.as_deref(),
disable_api_key_auth,
alpha_test_key,
client_version,
) {
@@ -64,7 +62,7 @@ impl MvpAgent {
}
fn has_proxy_credentials(&self) -> bool {
self.cfg.borrow().endpoints.deployment_key.is_some()
|| self.auth_manager.current_or_expired().is_some_and(|a| a.is_xai_auth())
|| self.auth_manager.current_or_expired().is_some_and(|a| a.is_session_auth())
}
/// `true` for session-based ACP auth methods.
fn is_session_based_auth(&self) -> bool {
@@ -79,7 +77,7 @@ impl MvpAgent {
self.auth_method_id.store(Some(std::sync::Arc::new(id)));
}
/// Return auth for sync config construction.
pub(super) fn current_or_buffered_auth(&self) -> Option<crate::auth::GrokAuth> {
pub(super) fn current_or_buffered_auth(&self) -> Option<crate::auth::KimiAuth> {
self.auth_manager
.current()
.or_else(|| {
@@ -101,7 +99,7 @@ impl MvpAgent {
fn has_managed_mcp_auth(&self) -> bool {
self.auth_manager
.current_or_expired()
.is_some_and(|a| a.is_managed_mcp_eligible())
.is_some_and(|a| a.is_session_auth())
}
/// Requires feature flag AND xAI authentication (OIDC or legacy WebLogin).
pub(super) fn can_fetch_managed_mcps(&self) -> bool {
@@ -195,7 +193,7 @@ impl MvpAgent {
.or_else(|| auth_manager.current_or_expired().map(|a| a.key));
if !auth_manager
.current_or_expired()
.is_some_and(|a| a.is_managed_mcp_eligible())
.is_some_and(|a| a.is_session_auth())
{
cache.lock().await.disable_gateway_tools();
for tx in session_txs {
@@ -401,7 +399,7 @@ impl MvpAgent {
let user_token = self
.auth_manager
.current_or_expired()
.filter(|a| a.is_xai_auth())
.filter(|a| a.is_session_auth())
.map(|a| a.key.clone());
let cfg = self.cfg.borrow();
let base_url = cfg.endpoints.resolve_feedback_base_url();
@@ -434,7 +432,7 @@ impl MvpAgent {
return None;
}
let auth = self.auth_manager.current_or_expired()?;
if !auth.is_xai_auth() {
if !auth.is_session_auth() {
return None;
}
let key = auth.key.clone();
@@ -498,12 +496,6 @@ impl MvpAgent {
..crate::session::slash_commands::CommandAvailability::default()
}
}
/// `true` when data collection should be suppressed (team ZDR or
/// coding-data-retention opt-out). Delegates to
/// [`AuthManager::is_data_collection_disabled`].
pub(crate) fn is_data_collection_disabled(&self) -> bool {
self.auth_manager.is_data_collection_disabled()
}
/// Current client type as set by the most recent `initialize()` call.
pub(crate) fn client_type(&self) -> ClientType {
*self.client_type.borrow()
@@ -513,8 +505,8 @@ impl MvpAgent {
pub(crate) fn session_turn_number(&self, sid: &acp::SessionId) -> Option<u64> {
self.session_turn_numbers.borrow().get(sid).copied()
}
/// Return the current GrokAuth credentials, if authenticated and not expired.
pub(crate) fn current_auth(&self) -> Option<crate::auth::GrokAuth> {
/// Return the current KimiAuth credentials, if authenticated and not expired.
pub(crate) fn current_auth(&self) -> Option<crate::auth::KimiAuth> {
self.auth_manager.current()
}
/// Shared plugin registry handle used by extensions for snapshot/reload.
@@ -613,55 +605,21 @@ impl MvpAgent {
}
}
/// When `cached_token` cannot proceed, prefer non-interactive `xai.api_key`
/// iff `should_advertise_xai_api_key`; otherwise `grok.com`. Returns `None`
/// when `preferred_method` is pinned (fail-closed — no cross-method fallthrough).
pub(super) fn cached_token_fallthrough_method_id(
&self,
) -> Option<acp::AuthMethodId> {
let preferred = self.cfg.borrow().grok_com_config.preferred_method;
/// iff `should_advertise_xai_api_key`; otherwise the interactive device
/// login.
pub(super) fn cached_token_fallthrough_method_id(&self) -> acp::AuthMethodId {
let id = auth_method::method_id_after_cached_token_unavailable(
auth_method::should_advertise_xai_api_key(
self.cfg.borrow().grok_com_config.api_key_auth_disabled(),
self.models_manager.models().values(),
),
preferred,
)?;
Some(acp::AuthMethodId::new(id))
auth_method::should_advertise_xai_api_key(self.models_manager.models().values()),
);
acp::AuthMethodId::new(id)
}
/// Shared exit for missing/expired/legacy `cached_token`: fall through with
/// `use_oauth` only when the target is interactive `grok.com`. When
/// `preferred_method` is pinned, fail instead of falling through.
/// Shared exit for missing/expired `cached_token`.
pub(super) async fn authenticate_after_cached_token_unavailable(
&self,
arguments: acp::AuthenticateRequest,
) -> Result<AuthenticateResponse, acp::Error> {
let Some(method_id) = self.cached_token_fallthrough_method_id() else {
let preferred = self.cfg.borrow().grok_com_config.preferred_method;
let msg = match preferred {
Some(crate::auth::PreferredAuthMethod::ApiKey) => {
auth_method::PREFERRED_API_KEY_UNAVAILABLE
}
_ => auth_method::PREFERRED_OIDC_UNAVAILABLE,
};
tracing::info!(
% msg, "cached_token unavailable; preferred_method forbids fallthrough"
);
kigi_log::unified_log::warn(
"auth cached_token fallthrough blocked by preferred_method",
None,
Some(
serde_json::json!(
{ "preferred_method" : preferred.map(| p | format!("{p:?}")), }
),
),
);
return Err(acp::Error::auth_required().data(msg));
};
let meta = if method_id.0.as_ref() == auth_method::KIGI_COM_METHOD_ID {
serde_json::json!({ "use_oauth" : true }).as_object().cloned()
} else {
arguments.meta
};
let method_id = self.cached_token_fallthrough_method_id();
let meta = arguments.meta;
tracing::info!(fallback = % method_id.0, "cached_token fallthrough");
kigi_log::unified_log::warn(
"auth cached_token fallthrough",
@@ -693,7 +651,7 @@ impl MvpAgent {
/// Agent-level fields materialised at startup (`worktree_type`,
/// `restore_code`) are NOT re-resolved here; that requires a
/// broader refactor of the init path.
pub(super) async fn refresh_remote_settings(&self, auth: &crate::auth::GrokAuth) {
pub(super) async fn refresh_remote_settings(&self, auth: &crate::auth::KimiAuth) {
if !crate::util::config::resolve_remote_fetch_enabled() {
tracing::debug!("post-auth settings refresh skipped: remote_fetch disabled");
return;
@@ -722,7 +680,7 @@ impl MvpAgent {
/// In-flight sessions are unaffected — they snapshot config at creation.
pub(super) async fn refresh_settings_and_reapply(
&self,
auth: &crate::auth::GrokAuth,
auth: &crate::auth::KimiAuth,
) {
self.refresh_remote_settings(auth).await;
let cwd = std::env::current_dir().ok();
@@ -746,7 +704,7 @@ impl MvpAgent {
/// Callers own their miss logging.
pub(super) async fn fetch_remote_settings(
&self,
auth: crate::auth::GrokAuth,
auth: crate::auth::KimiAuth,
) -> Option<crate::util::config::RemoteSettings> {
if !crate::util::config::resolve_remote_fetch_enabled() {
tracing::debug!("settings fetch skipped: remote_fetch disabled");
@@ -828,29 +786,16 @@ impl MvpAgent {
model: &ModelEntry,
origin_client: Option<crate::http::OriginClientInfo>,
) -> SamplingConfig {
let preferred = self.cfg.borrow().grok_com_config.preferred_method;
let session = match preferred {
Some(crate::auth::PreferredAuthMethod::ApiKey) => None,
_ if self.is_session_based_auth() => self.auth_manager.current_or_expired(),
_ => None,
let session = if self.is_session_based_auth() {
self.auth_manager.current_or_expired()
} else {
None
};
let has_session_key = session.is_some();
let mut credentials = resolve_credentials(
model,
session.as_ref().map(|a| a.key.as_str()),
);
if matches!(preferred, Some(crate ::auth::PreferredAuthMethod::Oidc))
&& !model.has_own_credentials()
&& credentials.auth_type == kigi_chat_state::AuthType::ApiKey
{
credentials.api_key = None;
credentials.auth_type = kigi_chat_state::AuthType::SessionToken;
}
crate::agent::config::enforce_disable_api_key_auth(
&mut credentials,
self.cfg.borrow().grok_com_config.api_key_auth_disabled(),
session.as_ref().map(|a| a.key.as_str()),
);
if !has_session_key && credentials.auth_type == kigi_chat_state::AuthType::ApiKey
&& !model.has_own_credentials() && self.is_session_based_auth()
{
@@ -893,7 +838,7 @@ impl MvpAgent {
let user_id = self
.auth_manager
.current_or_expired()
.filter(|a| a.is_xai_auth())
.filter(|a| a.is_session_auth())
.map(|a| a.user_id);
let mut config = crate::agent::config::sampling_config_for_model(
model,
@@ -945,39 +890,6 @@ impl MvpAgent {
);
(id.clone(), new_config)
}
/// Whether the current session is a personal grok.com account on a gated
/// tier (free / X Basic). The Imagine tools stay advertised to the model but
/// are flagged tier-restricted so they short-circuit at call time with the
/// SuperGrok upsell prose (see `ImageGenConfig`/`VideoGenConfig`'s
/// `tier_restricted`).
///
/// Fails **open** (returns `false`) whenever we can't positively confirm a
/// restricted personal tier — no auth yet, BYOK / API-key sessions, team
/// accounts, and an unknown/absent tier all pass. The server
/// authoritatively zero-limits Imagine for free & X Basic (429), so this
/// client gate is a UX optimization (a clean in-chat upsell instead of a
/// doomed request), never the security boundary — under-restricting is safe,
/// over-restricting would wrongly disable a paid feature.
///
/// Mirrors the pager's cosmetic slash-command gate
/// ([`crate::tier::is_restricted_tier_name`]); the only difference is the
/// absent-tier policy (the pager hides on `None`, we fail open on `None`).
fn is_tier_restricted_capability(&self) -> bool {
let Some(auth) = self.auth_manager.current() else {
return false;
};
if !auth.is_xai_auth() || auth.team_id.is_some() {
return false;
}
let tier = self
.cfg
.borrow()
.remote_settings
.as_ref()
.and_then(|rs| rs.subscription_tier_display.clone())
.or_else(|| jwt_tier_claim(&auth.key));
tier.as_deref().is_some_and(crate::tier::is_restricted_tier_name)
}
/// Build image generation config.
///
/// Both BYOK and session (OAuth) users go direct to `xai_api_base_url`.
@@ -992,7 +904,6 @@ impl MvpAgent {
let Some(ref api_key) = sampling_config.api_key else {
return ImageGenConfig::Disabled;
};
let tier_restricted = self.is_tier_restricted_capability();
let cfg = self.cfg.borrow();
let base_url = cfg.endpoints.xai_api_base_url.clone();
let version = cfg
@@ -1015,7 +926,7 @@ impl MvpAgent {
image_gen_enabled: cfg.resolve_image_gen().value,
image_edit_enabled: cfg.resolve_image_edit().value,
model_override: cfg.resolve_image_gen_model_override(),
tier_restricted,
tier_restricted: false,
}
}
/// Build deploy-service config. The tool talks directly to the deployer service.
@@ -1033,7 +944,6 @@ impl MvpAgent {
let Some(api_key) = self.sampling_config.borrow().api_key.clone() else {
return VideoGenConfig::Disabled;
};
let tier_restricted = self.is_tier_restricted_capability();
let cfg = self.cfg.borrow();
let zdr_video_output_s3 = cfg
.disable_zdr_incompatible_tools
@@ -1063,7 +973,7 @@ impl MvpAgent {
base_url,
extra_headers: headers,
zdr_video_output_s3: zdr_video_output_s3.map(Box::new),
tier_restricted,
tier_restricted: false,
}
}
pub(super) fn prepare_web_search_sampling_config(&self) -> Option<SamplingConfig> {
@@ -1076,7 +986,6 @@ impl MvpAgent {
&model_id,
&models,
session.as_ref().map(|a| a.key.as_str()),
self.cfg.borrow().grok_com_config.api_key_auth_disabled(),
alpha_test_key.clone(),
client_version,
&self.cfg.borrow().endpoints,
@@ -1224,7 +1133,6 @@ impl MvpAgent {
interactive_trust_prompted: Rc::new(
RefCell::new(std::collections::HashSet::new()),
),
tier_allowed: std::cell::Cell::new(true),
storage_mode,
default_yolo_mode,
default_auto_mode,
@@ -1252,9 +1160,6 @@ impl MvpAgent {
subagent_coordinator: RefCell::new(subagent_coordinator),
monitor_event_buffer: kigi_tools::implementations::grok_build::task::types::MonitorEventBuffer::default(),
bundle_sync_in_flight: Arc::new(std::sync::atomic::AtomicBool::new(false)),
post_unblock_jwt_retry_in_flight: Arc::new(
std::sync::atomic::AtomicBool::new(false),
),
workspace_ops: RefCell::new(None),
require_gateway_sessions: Rc::new(
RefCell::new(std::collections::HashSet::new()),
@@ -1268,11 +1173,7 @@ impl MvpAgent {
#[cfg(test)]
supervisor_spawn_count: std::cell::Cell::new(0),
};
instance
.auth_manager
.configure_refresher(
instance.cfg.borrow().grok_com_config.auth_provider_command.clone(),
);
instance.auth_manager.configure_refresher();
instance
}
/// Handle `x.ai/internal/evict_sessions` — the leader server tells us a
@@ -2189,7 +2090,7 @@ impl MvpAgent {
}
None => (kigi_hunk_tracker::HunkTrackerHandle::noop(), None),
};
let has_xai_auth = self.auth_manager.current().is_some_and(|a| a.is_xai_auth());
let has_xai_auth = self.auth_manager.current().is_some_and(|a| a.is_session_auth());
let loc_tracking_enabled = hunk_tracking_enabled && has_xai_auth
&& (self
.cfg
crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs
+10 -479
View File
@@ -98,75 +98,6 @@ pub(crate) fn reject_direct_hub_cloud_meta(
}
Ok(())
}
/// Marks a notification's meta field with `isReplay: true` for replayed session updates.
/// If `persist_data` is provided, it will be included in the meta under `x.ai/persist`.
/// Extract the numeric `tier` claim from a JWT access token (no signature
/// verification). Maps the `prod_auth.SubscriptionTier` proto enum values
/// to display-style strings that `normalize_tier` in the telemetry crate
/// will canonicalize for Mixpanel.
pub(crate) fn jwt_tier_claim(jwt: &str) -> Option<String> {
use base64::Engine;
let payload_b64 = jwt.split('.').nth(1)?;
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(payload_b64)
.ok()?;
let claims: serde_json::Value = serde_json::from_slice(&payload).ok()?;
let tier = claims.get("tier")?.as_u64()?;
Some(
match tier {
1 => "supergrok",
2 => "x_basic",
3 => "x_premium",
4 => "x_premium_plus",
5 => "supergrok_heavy",
6 => "supergrok_lite",
0 => "free",
_ => return Some(tier.to_string()),
}
.to_string(),
)
}
/// Resolve Mixpanel / AuthMeta `subscription_tier`.
///
/// Precedence:
/// 1. CCP `/settings` `subscription_tier_display` (when present and non-empty)
/// 2. [`AuthMode::ApiKey`] → `"api_key"` (never free)
/// 3. JWT `tier` claim via [`jwt_tier_claim`] (OAuth free → `"free"`)
pub(crate) fn resolve_subscription_tier_for_telemetry(
display: Option<String>,
auth: Option<&crate::auth::GrokAuth>,
) -> Option<String> {
if let Some(t) = display.filter(|s| !s.trim().is_empty()) {
return Some(t);
}
let auth = auth?;
if auth.auth_mode == crate::auth::AuthMode::ApiKey {
return Some("api_key".into());
}
jwt_tier_claim(&auth.key)
}
/// Whether a JWT `tier` claim (from [`jwt_tier_claim`]) reflects the live
/// `/user?include=subscription` tier string (from the subscription API / QUALIFYING_TIERS).
///
/// Post-unblock catalog refresh must not treat *any* present claim as enough:
/// an older paid claim (e.g. `x_basic`) can remain on the access token while
/// `/user` already reports a newly qualifying tier (e.g. `SuperGrokPro`). In
/// that case `/v1/models` would still be targeted at the stale level (the
/// "stale JWT tier skips retry" bug).
pub(crate) fn jwt_claim_matches_user_subscription_tier(
jwt_claim: &str,
user_subscription_tier: &str,
) -> bool {
match user_subscription_tier {
"GrokPro" => jwt_claim == "supergrok",
"XBasic" => jwt_claim == "x_basic",
"XPremium" => jwt_claim == "x_premium",
"XPremiumPlus" => jwt_claim == "x_premium_plus",
"SuperGrokPro" => jwt_claim == "supergrok_heavy",
"SuperGrokLite" => jwt_claim == "supergrok_lite",
_ => false,
}
}
fn parse_session_computer_sessions(_meta: Option<&acp::Meta>) -> Option<Vec<()>> {
None
}
@@ -617,11 +548,6 @@ pub struct MvpAgent {
/// into the detached prompt task; cleared for a workspace on GUI untrust
/// (`execute_hooks_action`) so a later re-open can re-prompt.
interactive_trust_prompted: Rc<RefCell<std::collections::HashSet<PathBuf>>>,
/// Whether the user's subscription tier is in the remote settings `allowed_tiers`
/// list. Set by `enforce_grok_code_access`; defaults to `true` (API-key and
/// external-auth users bypass the check). When `false`, the pager shows a
/// gate CTA instead of the prompt.
tier_allowed: std::cell::Cell<bool>,
/// Storage mode - determines whether to sync to backend (writeback) or local only
storage_mode: StorageMode,
/// Default YOLO mode - when true, sessions start with auto-approve enabled.
@@ -760,17 +686,6 @@ pub struct MvpAgent {
/// on completion without re-borrowing `&self`. `Send` is required
/// because the inner `sync_bundle_to_root` now uses `spawn_blocking`.
bundle_sync_in_flight: Arc<std::sync::atomic::AtomicBool>,
/// Single-flight guard for [`spawn_post_unblock_jwt_and_catalog_retry`].
///
/// After free→paid unblock the JWT may still lack a `tier` claim for
/// several seconds. Overlapping `CheckSubscription` RPCs (watch debounce,
/// paywall ticks, concurrent in-flight checks) would each otherwise spawn
/// another five-attempt `refresh_chain` backoff loop — multiplying IdP
/// traffic and redundant catalog work.
///
/// Cleared by [`PostUnblockJwtRetryInFlightGuard`] on task exit (including
/// panic/abort), not only on the normal post-backoff path.
post_unblock_jwt_retry_in_flight: Arc<std::sync::atomic::AtomicBool>,
/// Local workspace ops, built lazily via [`Self::ensure_local_workspace_ops`].
/// The agent never opens Computer Hub as a harness/client; remote cloud
/// sandboxes are gateway-owned (`gateway_bridge` / `computer_sessions`).
@@ -1013,26 +928,14 @@ struct AuthRequestMeta {
headless: bool,
#[serde(default)]
reauth: bool,
/// `--oauth`: force loopback. The only transport override sent over ACP
/// (loopback is the default; device is opt-in via env/config).
#[serde(default)]
use_oauth: bool,
/// When true, skip cached tokens and force the interactive browser login
/// flow. Used by the `/login` slash command for mid-session re-auth.
/// Unlike `reauth`, this does NOT clear existing credentials — if the
/// user abandons the browser flow, the current session continues.
/// When true, skip cached tokens and force the interactive login flow.
/// Used by the `/login` slash command for mid-session re-auth. Unlike
/// `reauth`, this does NOT clear existing credentials — if the user
/// abandons the device flow, the current session continues.
#[serde(default)]
force_interactive: bool,
}
impl AuthRequestMeta {
/// `--oauth` → force loopback; otherwise default (loopback).
fn login_override(&self) -> crate::auth::LoginTransportOverride {
if self.use_oauth {
crate::auth::LoginTransportOverride::ForceLoopback
} else {
crate::auth::LoginTransportOverride::None
}
}
fn from_json(meta: Option<&acp::Meta>) -> Self {
meta.cloned()
.and_then(|value| {
@@ -1654,239 +1557,21 @@ impl MvpAgent {
}
result
}
/// Check whether the user has access via remote settings `allow_access`.
///
/// Non-xAI auth (API keys, enterprise) always passes. For xAI OAuth2
/// users, reads `allow_access` from remote settings. Defaults to
/// `false` (blocked) when remote settings are unavailable.
pub(super) async fn enforce_grok_code_access(&self, auth: &crate::auth::GrokAuth) {
if !auth.is_xai_auth() {
self.tier_allowed.set(true);
return;
}
let allow = settings_allow_access(self.cfg.borrow().remote_settings.as_ref());
self.tier_allowed.set(allow);
if !allow {
tracing::info!(
"auth: user blocked by allow_access (remote settings grok_build_access_gate)"
);
self.retry_subscription_check().await;
}
}
/// Single-shot subscription check called by the pager's "Check
/// subscription" button (`x.ai/auth/check_subscription`). The pager
/// calls this every 5s while the paywall is shown, acting as the poller.
///
/// Queries `/user?include=subscription` for the live tier from the
/// subscription API. If a qualifying tier is found, does a best-effort
/// JWT refresh and settings re-fetch, lifts the gate, then — when the
/// access token's `tier` claim **matches** that live tier
/// ([`jwt_claim_matches_user_subscription_tier`]; bare `refresh_chain`
/// Ok or any older paid claim is not enough) — fire-and-forgets an
/// explicit model catalog refresh (`ModelsManager::on_auth_changed`) so
/// tier-targeted models appear without restart.
/// Catalog refresh is not awaited so gate lift / auth meta are not
/// blocked on `/v1/models`. Without a matching claim, defers to
/// `spawn_post_unblock_jwt_and_catalog_retry`.
pub(crate) async fn retry_subscription_check(&self) {
let (proxy_base_url, alpha_test_key) = {
let cfg = self.cfg.borrow();
(cfg.endpoints.proxy_url(), cfg.endpoints.alpha_test_key.clone())
};
let user_id = self
.auth_manager
.current()
.map(|a| a.user_id.clone())
.unwrap_or_default();
let result = super::subscription_check::single_check(
self.auth_manager.clone(),
&proxy_base_url,
alpha_test_key.as_deref(),
&user_id,
)
.await;
if let Some(unblocked) = result {
tracing::info!(
new_tier = % unblocked.new_tier, "subscription detected, lifting gate"
);
kigi_log::unified_log::info(
"paywall_check_gate_lifting",
None,
Some(
serde_json::json!(
{ "user_id" : user_id, "new_tier" : unblocked.new_tier, }
),
),
);
if let Some(settings) = unblocked.settings {
{
let mut cfg = self.cfg.borrow_mut();
cfg.remote_settings = Some(settings);
crate::agent::config::apply_remote_settings_side_effects(
cfg.remote_settings.as_ref(),
);
}
}
if crate::util::config::resolve_remote_fetch_enabled()
&& !settings_allow_access(self.cfg.borrow().remote_settings.as_ref())
{
tracing::info!(
new_tier = % unblocked.new_tier,
"subscription detected but allow_access still false, keeping gate"
);
kigi_log::unified_log::warn(
"paywall_check_gate_kept_allow_access_false",
None,
Some(
serde_json::json!(
{ "user_id" : user_id, "new_tier" : unblocked.new_tier, }
),
),
);
return;
}
self.tier_allowed.set(true);
let refresh_ok = match self
.auth_manager
.refresh_chain(
crate::auth::token_type::TokenType::OidcSession,
crate::auth::manager::RefreshReason::ServerRejected,
)
.await
{
Ok(_) => {
tracing::info!("post-unblock: JWT refresh_chain succeeded");
kigi_log::unified_log::info(
"paywall_check_jwt_refreshed",
None,
Some(serde_json::json!({ "user_id" : user_id })),
);
true
}
Err(e) => {
tracing::warn!(
error = % e,
"post-unblock: JWT refresh failed, user may need to re-login on next restart"
);
kigi_log::unified_log::warn(
"paywall_check_error",
None,
Some(
serde_json::json!(
{ "user_id" : user_id, "kind" :
"post_unblock_refresh_failed", "detail" : e.to_string(), }
),
),
);
false
}
};
let jwt_claim = self
.auth_manager
.current_or_expired()
.and_then(|auth| jwt_tier_claim(&auth.key));
let jwt_matches_new_tier = jwt_claim
.as_ref()
.is_some_and(|claim| jwt_claim_matches_user_subscription_tier(
claim,
&unblocked.new_tier,
));
if jwt_matches_new_tier {
let models_manager = self.models_manager.clone();
let user_id_log = user_id.clone();
let new_tier = unblocked.new_tier.clone();
let jwt_claim_log = jwt_claim.clone();
tokio::task::spawn(async move {
kigi_log::unified_log::info(
"model catalog: post_subscription_unblock refresh",
None,
Some(
serde_json::json!(
{ "user_id" : user_id_log, "new_tier" : new_tier,
"refresh_ok" : refresh_ok, "jwt_claim" : jwt_claim_log,
"jwt_matches_new_tier" : true, }
),
),
);
models_manager.on_auth_changed().await;
});
} else {
tracing::warn!(
refresh_ok, jwt_claim = ? jwt_claim, new_tier = % unblocked.new_tier,
"post-unblock: JWT tier claim missing or stale vs live tier; deferring model catalog refresh with retry"
);
kigi_log::unified_log::warn(
"model catalog: post_subscription_unblock deferred (jwt tier missing or stale)",
None,
Some(
serde_json::json!(
{ "user_id" : user_id, "new_tier" : unblocked.new_tier,
"refresh_ok" : refresh_ok, "jwt_claim" : jwt_claim, }
),
),
);
spawn_post_unblock_jwt_and_catalog_retry(
self.auth_manager.clone(),
self.models_manager.clone(),
self.post_unblock_jwt_retry_in_flight.clone(),
user_id.clone(),
unblocked.new_tier.clone(),
);
}
} else {
kigi_log::unified_log::info(
"paywall_check_no_subscription",
None,
Some(serde_json::json!({ "user_id" : user_id, })),
);
}
}
pub(crate) fn auth_response_with_meta(&self) -> AuthenticateResponse {
let (show_resolved_model, gate, subscription_tier) = {
let show_resolved_model = {
let cfg = self.cfg.borrow();
let rs = cfg.remote_settings.as_ref();
let gate = rs
.and_then(|s| s.gate_message.as_ref())
.filter(|m| !m.is_empty())
.map(|message| crate::auth::GateInfo {
message: message.clone(),
url: rs.and_then(|s| s.gate_url.clone()),
label: rs.and_then(|s| s.gate_label.clone()),
});
let subscription_tier = rs.and_then(|s| s.subscription_tier_display.clone());
(rs.and_then(|s| s.show_resolved_model), gate, subscription_tier)
cfg.remote_settings
.as_ref()
.and_then(|s| s.show_resolved_model)
};
let subscription_tier = resolve_subscription_tier_for_telemetry(
subscription_tier,
self.auth_manager.current_or_expired().as_ref(),
);
let meta = self
.auth_manager
.current()
.map(|auth| {
let gate = if !self.tier_allowed.get() && gate.is_none() {
let message = "A subscription is required.".to_string();
Some(crate::auth::GateInfo {
message,
url: Some(
"https://grok.com/supergrok?referrer=grok-build".to_string(),
),
label: Some("Subscribe".to_string()),
})
} else {
gate
};
let auth_meta = crate::auth::AuthMeta {
email: auth.email.clone(),
auth_mode: Some(format!("{:?}", auth.auth_mode)),
team_id: auth.team_id.clone(),
team_name: auth.team_name.clone(),
is_zdr: auth.is_zdr_team(),
team_role: auth.team_role.clone(),
coding_data_retention_opt_out: auth.coding_data_retention_opt_out,
show_resolved_model,
gate,
subscription_tier,
};
serde_json::to_value(auth_meta)
.ok()
@@ -1904,7 +1589,7 @@ impl MvpAgent {
let Some(auth) = self.auth_manager.current() else {
return;
};
let is_xai_auth = auth.is_xai_auth();
let is_session_auth = auth.is_session_auth();
let Some(settings) = self.fetch_remote_settings(auth).await else {
return;
};
@@ -1922,7 +1607,7 @@ impl MvpAgent {
None,
cfg.remote_settings.as_ref(),
);
if cfg.storage_mode == StorageMode::Writeback && !is_xai_auth {
if cfg.storage_mode == StorageMode::Writeback && !is_session_auth {
cfg.storage_mode = StorageMode::Local;
}
}
@@ -2107,160 +1792,6 @@ impl MvpAgent {
});
}
}
/// Clears [`MvpAgent::post_unblock_jwt_retry_in_flight`] on scope exit —
/// success, exhaustion, cancel/abort, or panic — so the single-flight flag
/// cannot wedge `true` for the rest of the process.
struct PostUnblockJwtRetryInFlightGuard {
flag: Arc<std::sync::atomic::AtomicBool>,
}
impl Drop for PostUnblockJwtRetryInFlightGuard {
fn drop(&mut self) {
self.flag.store(false, std::sync::atomic::Ordering::Release);
}
}
/// Background retry when post-unblock JWT lacks a tier claim that matches
/// the live `/user` tier. Re-attempts `refresh_chain` and only treats an
/// attempt as success when [`jwt_claim_matches_user_subscription_tier`]
/// holds (bare refresh Ok, free token, or a *stale older* paid claim are
/// all misses). Then refreshes the model catalog.
///
/// Gate lift already happened; this only recovers the tier-targeted catalog.
///
/// Single-flight: concurrent unblocks (overlapping `CheckSubscription`
/// RPCs while the JWT is still free/stale-targeted) share one backoff loop
/// via `in_flight`. A second spawn while a loop is running is a no-op.
/// The flag is released by [`PostUnblockJwtRetryInFlightGuard`] (Drop), not
/// only on the happy path after `execute_with_backoff`.
fn spawn_post_unblock_jwt_and_catalog_retry(
auth_manager: std::sync::Arc<crate::auth::AuthManager>,
models_manager: crate::agent::models::ModelsManager,
in_flight: Arc<std::sync::atomic::AtomicBool>,
user_id: String,
new_tier: String,
) {
use std::sync::atomic::Ordering;
if in_flight
.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_err()
{
tracing::debug!(
"post-unblock JWT/catalog retry already in flight, skipping duplicate spawn"
);
kigi_log::unified_log::info(
"model catalog: post_subscription_unblock jwt retry skipped (already in flight)",
None,
Some(serde_json::json!({ "user_id" : user_id, "new_tier" : new_tier, })),
);
return;
}
tokio::task::spawn(async move {
let _in_flight_guard = PostUnblockJwtRetryInFlightGuard {
flag: in_flight,
};
let backoff = crate::tools::retry::BackoffConfig::new(5, 2_000, 30_000);
let result = crate::tools::retry::execute_with_backoff(
&backoff,
|| {
let auth_manager = auth_manager.clone();
let new_tier = new_tier.clone();
async move {
let refresh_result = auth_manager
.refresh_chain(
crate::auth::token_type::TokenType::OidcSession,
crate::auth::manager::RefreshReason::ServerRejected,
)
.await;
let jwt_claim = auth_manager
.current_or_expired()
.and_then(|auth| jwt_tier_claim(&auth.key));
let matches = jwt_claim
.as_ref()
.is_some_and(|claim| jwt_claim_matches_user_subscription_tier(
claim,
&new_tier,
));
if matches {
Ok(())
} else {
let detail = match (&refresh_result, &jwt_claim) {
(Ok(_), None) => "refresh_ok but no tier claim".to_string(),
(Ok(_), Some(c)) => {
format!(
"refresh_ok but stale tier claim={c} (want {new_tier})"
)
}
(Err(e), Some(c)) => {
format!(
"refresh_err={e}; stale tier claim={c} (want {new_tier})"
)
}
(Err(e), None) => e.to_string(),
};
Err(format!("jwt tier not current: {detail}"))
}
}
},
|attempt, max_retries, delay| {
let user_id = user_id.clone();
let new_tier = new_tier.clone();
async move {
kigi_log::unified_log::warn(
"model catalog: post_subscription_unblock jwt retry scheduled",
None,
Some(
serde_json::json!(
{ "user_id" : user_id, "new_tier" : new_tier, "attempt" :
attempt, "max_retries" : max_retries, "delay_ms" : delay
.as_millis() as u64, }
),
),
);
}
},
)
.await;
match result {
Ok(()) => {
kigi_log::unified_log::info(
"model catalog: post_subscription_unblock refresh (after jwt retry)",
None,
Some(
serde_json::json!(
{ "user_id" : user_id, "new_tier" : new_tier, }
),
),
);
models_manager.on_auth_changed().await;
}
Err(e) => {
kigi_log::unified_log::warn(
"model catalog: post_subscription_unblock jwt retry exhausted",
None,
Some(
serde_json::json!(
{ "user_id" : user_id, "new_tier" : new_tier, "error" : e
.to_string(), }
),
),
);
}
}
});
}
/// Resolve `allow_access` from remote settings.
///
/// Returns `true` only when remote settings explicitly set `allow_access: true`.
/// Defaults to `false` (blocked) when settings are `None` or the field is
/// absent — matching the `grok_build_access_gate` flag's server-side default.
///
/// Used by both `enforce_grok_code_access` (initial login gate) and
/// `retry_subscription_check` (poller gate lift) to keep the decision in
/// one place.
pub(crate) fn settings_allow_access(
rs: Option<&crate::util::config::RemoteSettings>,
) -> bool {
rs.and_then(|s| s.allow_access).unwrap_or(false)
}
/// Parse `_meta.agentProfile` as a JSON object or string name.
/// Returns `None` if absent or invalid.
pub(crate) fn parse_agent_profile_from_meta(
crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs
+25 -386
View File
@@ -1,159 +1,4 @@
use super::*;
/// Build an unsigned JWT with a `tier` claim (header.payload.sig base64url).
fn jwt_with_tier(tier: u64) -> String {
use base64::Engine;
let enc = base64::engine::general_purpose::URL_SAFE_NO_PAD;
let header = enc.encode(br#"{"alg":"none"}"#);
let payload = enc.encode(format!(r#"{{"tier":{tier}}}"#).as_bytes());
format!("{header}.{payload}.sig")
}
#[test]
fn jwt_tier_claim_maps_free_and_paid() {
assert_eq!(jwt_tier_claim(&jwt_with_tier(0)).as_deref(), Some("free"));
assert_eq!(
jwt_tier_claim(&jwt_with_tier(1)).as_deref(),
Some("supergrok")
);
assert_eq!(
jwt_tier_claim(&jwt_with_tier(2)).as_deref(),
Some("x_basic")
);
assert_eq!(
jwt_tier_claim(&jwt_with_tier(3)).as_deref(),
Some("x_premium")
);
assert_eq!(
jwt_tier_claim(&jwt_with_tier(4)).as_deref(),
Some("x_premium_plus")
);
assert_eq!(
jwt_tier_claim(&jwt_with_tier(5)).as_deref(),
Some("supergrok_heavy")
);
assert_eq!(
jwt_tier_claim(&jwt_with_tier(6)).as_deref(),
Some("supergrok_lite")
);
assert_eq!(jwt_tier_claim(&jwt_with_tier(99)).as_deref(), Some("99"));
}
fn auth_with_mode(mode: crate::auth::AuthMode, key: &str) -> crate::auth::GrokAuth {
crate::auth::GrokAuth {
key: key.into(),
auth_mode: mode,
create_time: chrono::Utc::now(),
user_id: "u".into(),
email: None,
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: None,
team_name: None,
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token: None,
expires_at: None,
oidc_issuer: None,
oidc_client_id: None,
}
}
#[test]
fn resolve_subscription_tier_prefers_display_then_api_key_then_jwt() {
assert_eq!(
resolve_subscription_tier_for_telemetry(Some("Free".into()), None).as_deref(),
Some("Free")
);
let api = auth_with_mode(crate::auth::AuthMode::ApiKey, "xai-not-a-jwt");
assert_eq!(
resolve_subscription_tier_for_telemetry(Some(" ".into()), Some(&api)).as_deref(),
Some("api_key")
);
assert_eq!(
resolve_subscription_tier_for_telemetry(None, Some(&api)).as_deref(),
Some("api_key")
);
let oauth = auth_with_mode(crate::auth::AuthMode::Oidc, &jwt_with_tier(0));
assert_eq!(
resolve_subscription_tier_for_telemetry(None, Some(&oauth)).as_deref(),
Some("free")
);
assert_ne!(
resolve_subscription_tier_for_telemetry(None, Some(&api)).as_deref(),
Some("free")
);
}
/// JWT claim ↔ `/user` tier mapping used to gate post-unblock catalog refresh
/// (a stale older paid claim must not skip retry).
#[test]
fn jwt_claim_matches_user_subscription_tier_known_pairs() {
let cases = [
("supergrok", "GrokPro"),
("x_basic", "XBasic"),
("x_premium", "XPremium"),
("x_premium_plus", "XPremiumPlus"),
("supergrok_heavy", "SuperGrokPro"),
("supergrok_lite", "SuperGrokLite"),
];
for (claim, user_tier) in cases {
assert!(
jwt_claim_matches_user_subscription_tier(claim, user_tier),
"{claim} should match {user_tier}"
);
}
}
#[test]
fn jwt_claim_matches_user_subscription_tier_rejects_stale_and_unknown() {
assert!(!jwt_claim_matches_user_subscription_tier(
"x_basic",
"SuperGrokPro"
));
assert!(!jwt_claim_matches_user_subscription_tier(
"supergrok",
"SuperGrokPro"
));
assert!(!jwt_claim_matches_user_subscription_tier("free", "GrokPro"));
assert!(!jwt_claim_matches_user_subscription_tier("", "XPremium"));
assert!(!jwt_claim_matches_user_subscription_tier(
"supergrok_heavy",
"EnterpriseMystery"
));
}
/// Single-flight flag must clear on Drop even if the retry task panics /
/// aborts mid-backoff (guards against the flag stuck true forever).
#[test]
fn post_unblock_jwt_retry_in_flight_guard_clears_on_drop() {
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
let flag = Arc::new(AtomicBool::new(true));
{
let _guard = PostUnblockJwtRetryInFlightGuard { flag: flag.clone() };
assert!(flag.load(Ordering::Acquire));
}
assert!(
!flag.load(Ordering::Acquire),
"Drop must release post_unblock_jwt_retry_in_flight"
);
let flag = Arc::new(AtomicBool::new(true));
let flag_for_catch = flag.clone();
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let _guard = PostUnblockJwtRetryInFlightGuard {
flag: flag_for_catch,
};
panic!("simulate retry task panic");
}));
assert!(result.is_err());
assert!(
!flag.load(Ordering::Acquire),
"Drop must release flag on panic unwind"
);
}
mod hunk_tracking_mode {
use super::super::{plan_hunk_tracking, resolve_hunk_tracking_mode};
use kigi_hunk_tracker::TrackingMode;
@@ -354,42 +199,6 @@ fn trace_turn_to_i32_saturates_at_max() {
let result = i32::try_from(boundary).unwrap_or(i32::MAX);
assert_eq!(result, i32::MAX);
}
/// When remote settings are absent (`None`), default to blocked.
#[test]
fn settings_allow_access_none_settings_is_blocked() {
assert!(!settings_allow_access(None));
}
/// When `allow_access` is `Some(true)`, user is allowed.
#[test]
fn settings_allow_access_true_is_allowed() {
let rs = crate::util::config::RemoteSettings {
allow_access: Some(true),
..Default::default()
};
assert!(settings_allow_access(Some(&rs)));
}
/// When `allow_access` is `Some(false)` (remote settings default / rule
/// disabled), user stays blocked — even if they hold a qualifying
/// subscription. This is the regression guard for the bug where
/// `retry_subscription_check` unconditionally lifted the gate.
#[test]
fn settings_allow_access_false_is_blocked() {
let rs = crate::util::config::RemoteSettings {
allow_access: Some(false),
..Default::default()
};
assert!(!settings_allow_access(Some(&rs)));
}
/// When `/settings` returned successfully but the field is absent
/// (`None`), default to blocked (conservative).
#[test]
fn settings_allow_access_field_absent_is_blocked() {
let rs = crate::util::config::RemoteSettings {
allow_access: None,
..Default::default()
};
assert!(!settings_allow_access(Some(&rs)));
}
/// After allocating a turn number, `session_turn_numbers` holds the next
/// value (current + 1). This is the value that must be persisted via
/// `SetNextTraceTurn` so the counter survives restarts.
@@ -1333,10 +1142,10 @@ async fn ext_method_routes_auth_cleared_and_refreshes_resident_sessions() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let agent = build_agent_with_auth(crate::auth::GrokAuth {
let agent = build_agent_with_auth(crate::auth::KimiAuth {
key: "eligible".into(),
auth_mode: crate::auth::AuthMode::WebLogin,
..crate::auth::GrokAuth::test_default()
auth_mode: crate::auth::AuthMode::OAuth,
..crate::auth::KimiAuth::test_default()
});
use acp::Agent as _;
agent.managed_mcp_cache.lock().await.enable_gateway_tools();
@@ -1363,22 +1172,22 @@ async fn ext_method_routes_auth_cleared_and_refreshes_resident_sessions() {
/// Build a minimal MvpAgent suitable for testing extension methods.
fn build_minimal_agent_for_tests() -> MvpAgent {
use crate::agent::config::Config as AgentConfig;
use crate::auth::{AuthManager, GrokComConfig};
use crate::auth::{AuthManager, KimiCodeConfig};
let temp_dir = tempfile::tempdir().unwrap();
let auth_manager =
std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default()));
std::sync::Arc::new(AuthManager::new(temp_dir.path(), KimiCodeConfig::default()));
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(tx);
let cfg = AgentConfig::default();
MvpAgent::new(gateway, &cfg, auth_manager, None).expect("valid test config")
}
/// Build a minimal MvpAgent with pre-loaded auth for gate tests.
fn build_agent_with_auth(auth: crate::auth::GrokAuth) -> MvpAgent {
fn build_agent_with_auth(auth: crate::auth::KimiAuth) -> MvpAgent {
use crate::agent::config::Config as AgentConfig;
use crate::auth::{AuthManager, GrokComConfig};
use crate::auth::{AuthManager, KimiCodeConfig};
let temp_dir = tempfile::tempdir().unwrap();
let auth_manager =
std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default()));
std::sync::Arc::new(AuthManager::new(temp_dir.path(), KimiCodeConfig::default()));
auth_manager.hot_swap(auth);
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(tx);
@@ -1395,7 +1204,7 @@ fn build_agent_with_auth(auth: crate::auth::GrokAuth) -> MvpAgent {
#[serial_test::serial]
async fn ensure_plugin_registry_lazily_populates_snapshot() {
use crate::agent::config::Config as AgentConfig;
use crate::auth::{AuthManager, GrokComConfig};
use crate::auth::{AuthManager, KimiCodeConfig};
use kigi_test_support::EnvGuard;
let kigi_home = tempfile::tempdir().unwrap();
let _env = EnvGuard::set("KIGI_SHARE_DIR", kigi_home.path());
@@ -1412,7 +1221,7 @@ async fn ensure_plugin_registry_lazily_populates_snapshot() {
.unwrap();
let auth_home = tempfile::tempdir().unwrap();
let auth_manager =
std::sync::Arc::new(AuthManager::new(auth_home.path(), GrokComConfig::default()));
std::sync::Arc::new(AuthManager::new(auth_home.path(), KimiCodeConfig::default()));
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(tx);
let mut cfg = AgentConfig::default();
@@ -1599,10 +1408,10 @@ fn drain_roster_changed(
async fn push_roster_activity_delta_broadcasts_overridden_activity() {
use crate::agent::config::Config as AgentConfig;
use crate::agent::roster::RosterActivity;
use crate::auth::{AuthManager, GrokComConfig};
use crate::auth::{AuthManager, KimiCodeConfig};
let temp_dir = tempfile::tempdir().unwrap();
let auth_manager =
std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default()));
std::sync::Arc::new(AuthManager::new(temp_dir.path(), KimiCodeConfig::default()));
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(tx);
let cfg = AgentConfig::default();
@@ -2043,7 +1852,6 @@ async fn auth_type_session_based_no_current_returns_session_token() {
for method_id in [
crate::agent::auth_method::CACHED_TOKEN_AUTH_METHOD_ID,
crate::agent::auth_method::KIGI_COM_METHOD_ID,
crate::agent::auth_method::OIDC_METHOD_ID,
] {
let agent = build_minimal_agent_for_tests();
agent.set_auth_method(acp::AuthMethodId::new(method_id));
@@ -2084,12 +1892,12 @@ async fn auth_type_xai_api_key_no_current_returns_api_key() {
/// common case during a healthy session.
#[tokio::test(flavor = "current_thread")]
async fn auth_type_session_based_with_current_returns_session_token() {
use crate::auth::GrokAuth;
use crate::auth::KimiAuth;
let agent = build_minimal_agent_for_tests();
agent.set_auth_method(acp::AuthMethodId::new(
crate::agent::auth_method::OIDC_METHOD_ID,
crate::agent::auth_method::KIGI_COM_METHOD_ID,
));
agent.auth_manager.hot_swap(GrokAuth::test_default());
agent.auth_manager.hot_swap(KimiAuth::test_default());
assert!(agent.auth_manager.current().is_some());
assert_eq!(agent.auth_type(), kigi_chat_state::AuthType::SessionToken,);
}
@@ -2112,27 +1920,13 @@ async fn auth_type_no_method_id_no_current_returns_api_key() {
/// here matches pre-fix behavior and keeps logging stable.
#[tokio::test(flavor = "current_thread")]
async fn auth_type_no_method_id_with_current_returns_session_token() {
use crate::auth::GrokAuth;
use crate::auth::KimiAuth;
let agent = build_minimal_agent_for_tests();
agent.auth_manager.hot_swap(GrokAuth::test_default());
agent.auth_manager.hot_swap(KimiAuth::test_default());
assert!(agent.auth_method_id.load().is_none());
assert!(agent.auth_manager.current().is_some());
assert_eq!(agent.auth_type(), kigi_chat_state::AuthType::SessionToken,);
}
/// Minimal agent whose `grok_com_config` engages the api-key kill switch
/// (`disable_api_key_auth = true`), mirroring a forced-IdP deployment.
fn build_agent_with_api_key_auth_disabled() -> MvpAgent {
use crate::agent::config::Config as AgentConfig;
use crate::auth::{AuthManager, GrokComConfig};
let temp_dir = tempfile::tempdir().unwrap();
let auth_manager =
std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default()));
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(tx);
let mut cfg = AgentConfig::default();
cfg.grok_com_config.disable_api_key_auth = Some(true);
MvpAgent::new(gateway, &cfg, auth_manager, None).expect("valid test config")
}
/// Deployment-key / managed-config user: `XAI_API_KEY` resolves and the kill
/// switch is off, so a dead `cached_token` MUST fall through to `xai.api_key`
/// (no browser). This is the exact regression the fallthrough fixes.
@@ -2145,36 +1939,12 @@ async fn cached_token_fallthrough_prefers_api_key_for_deployment_key() {
let _key = EnvGuard::set(XAI_API_KEY_ENV_VAR, "test-deployment-key");
let agent = build_minimal_agent_for_tests();
assert_eq!(
agent
.cached_token_fallthrough_method_id()
.as_ref()
.map(|id| id.0.as_ref()),
Some(XAI_API_KEY_METHOD_ID),
agent.cached_token_fallthrough_method_id().0.as_ref(),
XAI_API_KEY_METHOD_ID,
"deployment-key user (XAI_API_KEY set, no kill switch) must fall \
through to xai.api_key on a dead cached_token -- not interactive login",
);
}
/// Forced-IdP deployment: even with `XAI_API_KEY` present, the admin kill
/// switch keeps the fallthrough on interactive `grok.com` (api-key auth is
/// neither advertised nor an eligible fallthrough).
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn cached_token_fallthrough_respects_kill_switch() {
use crate::agent::auth_method::{KIGI_COM_METHOD_ID, XAI_API_KEY_ENV_VAR};
use kigi_test_support::EnvGuard;
let _lockdown = EnvGuard::unset("KIGI_DISABLE_API_KEY_AUTH");
let _key = EnvGuard::set(XAI_API_KEY_ENV_VAR, "test-deployment-key");
let agent = build_agent_with_api_key_auth_disabled();
assert_eq!(
agent
.cached_token_fallthrough_method_id()
.as_ref()
.map(|id| id.0.as_ref()),
Some(KIGI_COM_METHOD_ID),
"disable_api_key_auth must keep the cached_token fallthrough on \
interactive grok.com so XAI_API_KEY can't bypass forced IdP login",
);
}
/// No advertiseable credentials at all (no env key, no kill switch): the user
/// genuinely needs to log in, so the fallthrough is interactive `grok.com`.
#[tokio::test(flavor = "current_thread")]
@@ -2189,75 +1959,11 @@ async fn cached_token_fallthrough_falls_to_grok_com_without_credentials() {
let _legacy = EnvGuard::unset(LEGACY_XAI_API_KEY_ENV_VAR);
let agent = build_minimal_agent_for_tests();
assert_eq!(
agent
.cached_token_fallthrough_method_id()
.as_ref()
.map(|id| id.0.as_ref()),
Some(KIGI_COM_METHOD_ID),
agent.cached_token_fallthrough_method_id().0.as_ref(),
KIGI_COM_METHOD_ID,
"no API-key creds and no kill switch -> interactive grok.com login",
);
}
/// Verifies the 4-state matrix of `(disable_zdr_incompatible_tools, zdr_video_output_s3)`:
///
/// | ZDR flag | S3 config | Result |
/// |----------|-----------|---------------------------------------------|
/// | false | None | Enabled, no S3 (normal non-ZDR mode) |
/// | true | None | Disabled (ZDR with no escape hatch) |
/// | false | Some | Enabled, S3 **not** threaded (non-ZDR) |
/// | true | Some | Enabled, S3 threaded (ZDR with upload path) |
#[tokio::test(flavor = "current_thread")]
async fn prepare_video_gen_config_disabled_when_zdr_flag_set() {
use kigi_tools::implementations::grok_build::video_gen::{
S3AccessCredentials, VideoGenConfig, ZdrVideoOutputS3Config,
};
fn zdr_s3() -> ZdrVideoOutputS3Config {
ZdrVideoOutputS3Config {
bucket: "team-videos".into(),
endpoint: "https://s3.example.com".into(),
region: "us-east-1".into(),
key_prefix: "grok-videos/".into(),
expires_secs: 900,
read_write: S3AccessCredentials {
access_key_id: "AKIA...".into(),
secret_access_key: "secret".into(),
},
read_only: None,
}
}
let agent = build_minimal_agent_for_tests();
agent.sampling_config.borrow_mut().api_key = Some("test-key".to_string());
assert!(matches!(
agent.prepare_video_gen_config(),
VideoGenConfig::Enabled { .. }
));
agent.cfg.borrow_mut().disable_zdr_incompatible_tools = true;
assert!(matches!(
agent.prepare_video_gen_config(),
VideoGenConfig::Disabled
));
agent.cfg.borrow_mut().zdr_video_output_s3 = Some(zdr_s3());
agent.cfg.borrow_mut().disable_zdr_incompatible_tools = false;
let VideoGenConfig::Enabled {
zdr_video_output_s3: s3_when_non_zdr,
..
} = agent.prepare_video_gen_config()
else {
panic!("expected Enabled");
};
assert!(
s3_when_non_zdr.is_none(),
"S3 config must not be threaded when ZDR flag is off"
);
agent.cfg.borrow_mut().disable_zdr_incompatible_tools = true;
let VideoGenConfig::Enabled {
zdr_video_output_s3,
..
} = agent.prepare_video_gen_config()
else {
panic!("expected Enabled");
};
assert!(zdr_video_output_s3.as_ref().is_some_and(|c| c.is_valid()));
}
/// The imagine tier gate fails **open**: with no resolved auth we can't confirm
/// a restricted personal tier, so the tools stay advertised and un-flagged (the
/// server 429 remains the authoritative backstop). Guards against accidentally
@@ -2278,73 +1984,6 @@ async fn prepare_image_gen_config_fails_open_without_auth() {
"no resolved auth ⇒ fail open (tools not tier-restricted)"
);
}
#[tokio::test]
async fn data_collection_enabled_for_normal_user() {
let agent = build_agent_with_auth(crate::auth::GrokAuth::test_default());
assert!(
!agent.is_data_collection_disabled(),
"normal user must have data collection enabled"
);
}
#[tokio::test]
async fn data_collection_disabled_for_zdr_team() {
let agent = build_agent_with_auth(crate::auth::GrokAuth {
team_blocked_reasons: vec!["BLOCKED_REASON_NO_LOGS".into()],
..crate::auth::GrokAuth::test_default()
});
assert!(
agent.is_data_collection_disabled(),
"ZDR team must have data collection disabled"
);
}
#[tokio::test]
async fn data_collection_disabled_for_zdr_moderated_team() {
let agent = build_agent_with_auth(crate::auth::GrokAuth {
team_blocked_reasons: vec!["BLOCKED_REASON_NO_LOGS_MODERATED".into()],
..crate::auth::GrokAuth::test_default()
});
assert!(
agent.is_data_collection_disabled(),
"ZDR-moderated team must have data collection disabled"
);
}
#[tokio::test]
async fn data_collection_disabled_for_opted_out_team() {
let agent = build_agent_with_auth(crate::auth::GrokAuth {
coding_data_retention_opt_out: true,
..crate::auth::GrokAuth::test_default()
});
assert!(
agent.is_data_collection_disabled(),
"opted-out team must have data collection disabled"
);
}
#[tokio::test]
async fn data_collection_disabled_for_zdr_plus_opt_out() {
let agent = build_agent_with_auth(crate::auth::GrokAuth {
team_blocked_reasons: vec!["BLOCKED_REASON_NO_LOGS".into()],
coding_data_retention_opt_out: true,
..crate::auth::GrokAuth::test_default()
});
assert!(
agent.is_data_collection_disabled(),
"ZDR + opt-out must have data collection disabled"
);
}
#[tokio::test]
async fn data_collection_enabled_for_non_zdr_team_with_unrelated_blocks() {
let agent = build_agent_with_auth(crate::auth::GrokAuth {
team_blocked_reasons: vec![
"BLOCKED_REASON_BILLING".into(),
"BLOCKED_REASON_SUSPENDED".into(),
],
..crate::auth::GrokAuth::test_default()
});
assert!(
!agent.is_data_collection_disabled(),
"non-ZDR blocked reasons must not disable data collection"
);
}
/// `parse_session_kind` routes `session/load` to the gateway Chat path vs. the
/// disk-backed Build path. Anything but an explicit `kind: "chat"` is Build.
#[test]
@@ -3098,10 +2737,10 @@ fn build_agent_with_gateway_rx() -> (
tokio::sync::mpsc::UnboundedReceiver<kigi_acp_lib::AcpClientMessage>,
) {
use crate::agent::config::Config as AgentConfig;
use crate::auth::{AuthManager, GrokComConfig};
use crate::auth::{AuthManager, KimiCodeConfig};
let temp_dir = tempfile::tempdir().unwrap();
let auth_manager =
std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default()));
std::sync::Arc::new(AuthManager::new(temp_dir.path(), KimiCodeConfig::default()));
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(tx);
let cfg = AgentConfig::default();
@@ -3648,14 +3287,14 @@ mod soft_default_settings_emit {
#[tokio::test]
async fn emit_settings_update_carries_permission_mode_from_cfg() {
use crate::agent::config::Config as AgentConfig;
use crate::auth::{AuthManager, GrokComConfig};
use crate::auth::{AuthManager, KimiCodeConfig};
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let temp_dir = tempfile::tempdir().unwrap();
let auth_manager = std::sync::Arc::new(AuthManager::new(
temp_dir.path(),
GrokComConfig::default(),
KimiCodeConfig::default(),
));
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(tx);
crates/codegen/kigi-shell/src/agent/session_registry_client.rs
+6 -6
View File
@@ -141,7 +141,7 @@ pub struct SessionRegistryClient {
raw_client: reqwest::Client,
client: reqwest_middleware::ClientWithMiddleware,
base_url: String,
credentials: crate::util::grok_auth_credentials::GrokAuthCredentials,
credentials: crate::util::kigi_auth_credentials::KigiAuthCredentials,
session_id: Option<String>,
}
@@ -152,7 +152,7 @@ impl SessionRegistryClient {
raw_client: http_client.clone(),
client: reqwest_middleware::ClientBuilder::new(http_client).build(),
base_url: base_url.into(),
credentials: crate::util::grok_auth_credentials::GrokAuthCredentials::new(Some(
credentials: crate::util::kigi_auth_credentials::KigiAuthCredentials::new(Some(
user_token.into(),
)),
session_id: None,
@@ -549,7 +549,7 @@ mod tests {
/// Verify per-request auth resolve picks up rotated tokens.
#[tokio::test]
async fn session_registry_client_uses_active_auth_for_each_request() {
use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig};
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
use axum::{Router, response::IntoResponse, routing::post};
use chrono::{Duration, Utc};
use std::net::SocketAddr;
@@ -575,14 +575,14 @@ mod tests {
tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
am.hot_swap(KimiAuth {
key: "fresh-from-auth-manager".into(),
auth_mode: AuthMode::ApiKey,
create_time: Utc::now(),
user_id: "user-42".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
let client = SessionRegistryClient::new(format!("http://{addr}"), "STALE-build-time-token")
crates/codegen/kigi-shell/src/agent/subagent/mod.rs
+1 -1
View File
@@ -156,7 +156,7 @@ pub(crate) struct SubagentSpawnContext {
reason = "unused in production; remove expect when wired or delete the item"
)]
pub storage_mode: crate::config::StorageMode,
pub auth: Option<crate::auth::GrokAuth>,
pub auth: Option<crate::auth::KimiAuth>,
pub parent_cwd: PathBuf,
pub parent_session_id: String,
pub yolo_mode: bool,
crates/codegen/kigi-shell/src/agent/subscription_check.rs
-191
View File
@@ -1,191 +0,0 @@
//! Subscription check for paywall gate lift.
//!
//! Provides `single_check()` which queries `GET /user?include=subscription`
//! for the live subscription tier from the backend, independent of the JWT.
//! If a qualifying tier is detected, does a best-effort JWT refresh and
//! settings re-fetch, then returns an `UnblockResult` so the agent can
//! lift the gate.
//!
//! The pager drives the polling via `x.ai/auth/check_subscription`: the 5s
//! paywall chain, the free-tier watch, the refocus check, and
//! verify-before-paywall gate deferral (see the pager's `app::subscription`
//! module).
use crate::auth::AuthManager;
use crate::auth::UserInfo;
use crate::auth::manager::RefreshReason;
use crate::auth::token_type::TokenType;
use std::sync::Arc;
use std::time::Duration;
/// Subscription tiers that qualify for Grok Build access.
/// Any active subscription qualifies -- the access gate in remote settings
/// controls which tiers are actually allowed.
const QUALIFYING_TIERS: &[&str] = &[
"SuperGrokPro",
"GrokPro",
"SuperGrokLite",
"XPremiumPlus",
"XPremium",
"XBasic",
];
/// Successful subscription check result: confirmed qualifying tier +
/// optionally refreshed settings.
pub(crate) struct UnblockResult {
pub(crate) new_tier: String,
pub(crate) settings: Option<crate::util::config::RemoteSettings>,
}
/// Fetch `/user?include=subscription` and return the parsed `UserInfo`.
async fn fetch_user_info(
http_client: &reqwest::Client,
url: &str,
auth: &crate::auth::GrokAuth,
auth_manager: &AuthManager,
alpha_test_key: Option<&str>,
) -> Result<UserInfo, &'static str> {
let request = http_client
.get(url)
.timeout(Duration::from_secs(10))
.header("Authorization", format!("Bearer {}", auth.key))
.header(
"X-XAI-Token-Auth",
auth_manager.grok_com_config().token_header.as_str(),
)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
);
let _ = alpha_test_key;
match request.send().await {
Ok(resp) if resp.status().is_success() => {
resp.json::<UserInfo>().await.map_err(|_| "parse")
}
Ok(_resp) => Err("http_status"),
Err(e) if e.is_timeout() => Err("timeout"),
Err(_) => Err("transport"),
}
}
/// Single-shot subscription check. Called by the pager every 5s while
/// the paywall is shown (`x.ai/auth/check_subscription`).
///
/// Queries `/user?include=subscription` for the live tier. If a qualifying
/// tier is found, does a best-effort JWT refresh + settings re-fetch and
/// returns `Some(UnblockResult)`. Returns `None` if no qualifying
/// subscription exists or the request fails.
#[tracing::instrument(name = "paywall_check", skip_all, fields(user_id = %user_id))]
pub(crate) async fn single_check(
auth_manager: Arc<AuthManager>,
proxy_base_url: &str,
alpha_test_key: Option<&str>,
user_id: &str,
) -> Option<UnblockResult> {
let user_url = format!("{}/user?include=subscription", proxy_base_url);
let http_client = crate::http::shared_client();
let auth = auth_manager.current()?;
let user_info = match fetch_user_info(
&http_client,
&user_url,
&auth,
&auth_manager,
alpha_test_key,
)
.await
{
Ok(ui) => ui,
Err(kind) => {
kigi_log::unified_log::warn(
"paywall_check_error",
None,
Some(serde_json::json!({ "user_id" : user_id, "kind" : kind })),
);
return None;
}
};
kigi_log::unified_log::info(
"paywall_check_result",
None,
Some(serde_json::json!(
{ "user_id" : user_id, "subscription_tier" : user_info.subscription_tier,
}
)),
);
let new_tier = match &user_info.subscription_tier {
Some(tier) if !tier.is_empty() => tier.clone(),
_ => return None,
};
if !QUALIFYING_TIERS.contains(&new_tier.as_str()) {
return None;
}
kigi_log::unified_log::info(
"paywall_check_subscription_detected",
None,
Some(serde_json::json!({ "user_id" : user_id, "new_tier" : new_tier, })),
);
if let Err(e) = auth_manager
.refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected)
.await
{
kigi_log::unified_log::warn(
"paywall_check_error",
None,
Some(serde_json::json!(
{ "user_id" : user_id, "kind" : "refresh_failed", "detail" : e
.to_string(), }
)),
);
}
let settings = if crate::util::config::resolve_remote_fetch_enabled() {
let base_url = proxy_base_url.to_string();
let auth_for_settings = auth_manager.current().unwrap_or(auth);
let atk = alpha_test_key.map(str::to_string);
tokio::task::spawn_blocking(move || {
crate::remote::fetch_settings_blocking(&base_url, &auth_for_settings, atk.as_deref())
})
.await
.ok()
.flatten()
} else {
None
};
kigi_log::unified_log::info(
"paywall_check_unblocked",
None,
Some(serde_json::json!({ "user_id" : user_id, "new_tier" : new_tier })),
);
Some(UnblockResult { new_tier, settings })
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn qualifying_tiers_includes_all_paid_tiers() {
for tier in &[
"SuperGrokPro",
"GrokPro",
"SuperGrokLite",
"XPremiumPlus",
"XPremium",
"XBasic",
] {
assert!(
QUALIFYING_TIERS.contains(tier),
"{tier} must be in QUALIFYING_TIERS"
);
}
}
#[test]
fn free_tier_is_not_qualifying() {
assert!(!QUALIFYING_TIERS.contains(&"Free"));
}
#[test]
fn empty_tier_is_not_qualifying() {
assert!(!QUALIFYING_TIERS.contains(&""));
}
/// The subscription check only returns `Some` when `/user` reports a
/// qualifying tier. Verify the tier matching is exact (no prefix match).
#[test]
fn partial_tier_name_is_not_qualifying() {
assert!(!QUALIFYING_TIERS.contains(&"Super"));
assert!(!QUALIFYING_TIERS.contains(&"Grok"));
assert!(!QUALIFYING_TIERS.contains(&"XPremium+"));
}
}
crates/codegen/kigi-shell/src/auth/attribution.rs
+10 -10
View File
@@ -370,7 +370,7 @@ pub(crate) fn record_auth_401(
/// This function performs **exactly one** read-side acquisition of
/// [`AuthManager`]'s internal `RwLock` -- it calls
/// [`AuthManager::current`] once and derives both `current_key_prefix`
/// and the mint/expiry fields from the resulting `GrokAuth`.
/// and the mint/expiry fields from the resulting `KimiAuth`.
///
/// `is_stale_snapshot` is `true` only when the live `current()` token
/// differs from the bearer the client sent. When `current()` returns
@@ -390,7 +390,7 @@ fn compute_attribution_payload(
// query can break down on this).
let sent_prefix = sent_bearer.map(token_suffix).unwrap_or("");
// Single read-lock acquisition: pull the live `GrokAuth` (or
// Single read-lock acquisition: pull the live `KimiAuth` (or
// `None`) once and derive every other field from it.
let current_auth = auth_manager.current();
let current_prefix_owned: Option<String> = current_auth
@@ -411,7 +411,7 @@ fn compute_attribution_payload(
//
// TODO: mirror the full External-with-ttl branch from
// `AuthManager::is_token_expired` (uses
// `grok_com_config.auth_token_ttl` when `expires_at` is `None`
// `kimi_code_config.auth_token_ttl` when `expires_at` is `None`
// and `auth_mode == External`). The current 2-branch fallback
// (`expires_at` if Some else `create_time + TOKEN_TTL`) is good
// enough for diagnostic metadata; the External-ttl branch is
@@ -441,7 +441,7 @@ mod tests {
use chrono::{Duration, Utc};
use crate::auth::{AuthManager, GrokAuth, GrokComConfig};
use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig};
use super::*;
@@ -449,17 +449,17 @@ mod tests {
/// nothing from a developer's actual `~/.kigi/auth.json` leaks in.
fn empty_auth_manager() -> (tempfile::TempDir, AuthManager) {
let dir = tempfile::tempdir().expect("tempdir");
let cfg = GrokComConfig::default();
let cfg = KimiCodeConfig::default();
let am = AuthManager::new(dir.path(), cfg);
(dir, am)
}
fn fresh_auth(key: &str) -> GrokAuth {
GrokAuth {
fn fresh_auth(key: &str) -> KimiAuth {
KimiAuth {
key: key.to_string(),
create_time: Utc::now(),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
}
}
@@ -553,12 +553,12 @@ mod tests {
#[test]
fn legacy_token_uses_two_branch_fallback() {
let (_dir, am) = empty_auth_manager();
let auth = GrokAuth {
let auth = KimiAuth {
key: "k".into(),
create_time: Utc::now() - Duration::seconds(60),
// No expires_at => falls through to create_time + TOKEN_TTL
// (= 30 days).
..GrokAuth::test_default()
..KimiAuth::test_default()
};
am.hot_swap(auth);
crates/codegen/kigi-shell/src/auth/config.rs
+30 -410
View File
@@ -1,423 +1,43 @@
use super::model::TEAM_PRINCIPAL_TYPE;
//! Kimi Code auth configuration.
//!
//! The wire endpoints come from [`kigi_env`] (`oauth_host()`, overridable via
//! `KIGI_OAUTH_HOST`) and the client id is fixed
//! ([`crate::auth::kimi_oauth::KIMI_CODE_CLIENT_ID`]), so this config carries
//! no per-deployment OAuth knobs. The struct is kept (deserialized from the
//! agent config TOML) as the extension point for future auth options.
use serde::{Deserialize, Serialize};
// Transitional: the M1 auth rewrite (Kimi device flow) replaces this origin.
const AUTH_ORIGIN_DEFAULT: &str = "https://grok.com";
fn default_oidc_scopes() -> Vec<String> {
vec![
"openid".into(),
"profile".into(),
"email".into(),
"offline_access".into(),
"api:access".into(),
]
}
/// Default scopes for the xAI OAuth2 provider. Includes `grok-cli:access`
/// which authorizes the token for API proxy requests.
fn default_oauth2_scopes() -> Vec<String> {
vec![
"openid".into(),
"profile".into(),
"email".into(),
"offline_access".into(),
"grok-cli:access".into(),
"api:access".into(),
"conversations:read".into(),
"conversations:write".into(),
]
}
fn default_team_oauth2_scopes() -> Vec<String> {
vec![
"profile".into(),
"offline_access".into(),
"grok-cli:access".into(),
"api:access".into(),
"team:read".into(),
"conversations:read".into(),
"conversations:write".into(),
]
}
/// Pin automatic auth to one method (`[auth] preferred_method` in config.toml).
///
/// When set, only that method is used for automatic selection; if it is
/// unavailable, auth fails (no silent fallthrough to the other method).
/// Unset keeps today's multi-method fallthrough (session preferred when both
/// exist). Config-toml only — not remote settings, settings UI, or env.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PreferredAuthMethod {
/// `XAI_API_KEY` / auth.json `xai::api_key` / per-model BYOK (`xai.api_key`).
ApiKey,
/// OIDC / OAuth2 session (`cached_token`, interactive `grok.com` / `oidc`,
/// including devbox-minted OIDC).
Oidc,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
/// Persisted-credential scope key for the Kimi Code OAuth session — both the
/// auth.json map key and the system-keyring entry name (service `kigi`).
pub const KIMI_CODE_OAUTH_SCOPE: &str = "oauth/kimi-code";
/// Auth configuration block (`[kimi_code_config]` in the agent config).
/// Currently empty: the OAuth host and client id are fixed by the
/// environment crate.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct GrokComConfig {
/// Auth origin / login-host display (doubles as the legacy WS origin name).
pub grok_ws_origin: String,
pub token_header: String,
/// OIDC config for customer-provided IdPs. See [`OidcAuthConfig`].
#[serde(default, skip_serializing_if = "Option::is_none")]
pub oidc: Option<OidcAuthConfig>,
/// OAuth2 provider config. When set, preferred over the legacy relay flow.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub oauth2: Option<OAuth2ProviderConfig>,
/// External auth provider command (stdout = token, stderr = user UX, exit 0 = success).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub auth_provider_command: Option<String>,
/// Login button label (env: `KIGI_AUTH_PROVIDER_LABEL`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub auth_provider_label: Option<String>,
/// Token TTL in seconds for external auth providers that output bare
/// tokens without `expires_in`. Synthesizes `expires_at` so proactive
/// refresh works. Env: `KIGI_AUTH_TOKEN_TTL`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub auth_token_ttl: Option<u64>,
/// Admin kill switch: when `Some(true)`, the `xai.api_key` auth method is
/// neither advertised nor accepted, so `XAI_API_KEY`/per-model credentials
/// can't bypass the deployment's IdP login. Env: `KIGI_DISABLE_API_KEY_AUTH`.
/// Parity with common force-login-method admin knobs.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub disable_api_key_auth: Option<bool>,
/// Restrict login to a specific team — the login token's team principal must
/// equal this. Put in `requirements.toml` to enforce as non-overridable policy.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub force_login_team_uuid: Option<ForceLoginTeam>,
/// Pin automatic auth to `api_key` or `oidc`. When set and the chosen
/// method is unavailable, auth fails (no fallthrough). Unset keeps
/// multi-method fallthrough. Config.toml only (`[auth] preferred_method`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub preferred_method: Option<PreferredAuthMethod>,
}
/// Team login restriction. TOML string or array; an empty array fails closed.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ForceLoginTeam {
/// The only allowed team.
Single(String),
/// Allowed teams; empty = fail closed.
AnyOf(Vec<String>),
}
/// Customer OIDC Identity Provider configuration (`[grok_com_config.oidc]`).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OidcAuthConfig {
pub issuer: String,
pub client_id: String,
#[serde(default = "default_oidc_scopes")]
pub scopes: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub audience: Option<String>,
}
/// OAuth2 provider configuration (`KIGI_OAUTH2_ISSUER` / `KIGI_OAUTH2_CLIENT_ID`).
///
/// Uses the standard OAuth 2.1 Auth Code + PKCE flow via [`OidcAuthConfig`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OAuth2ProviderConfig {
pub issuer: String,
pub client_id: String,
#[serde(default = "default_oauth2_scopes")]
pub scopes: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub principal_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub principal_id: Option<String>,
/// Client-supplied referrer for OAuth usage-attribution analytics.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub referrer: Option<String>,
}
pub const XAI_OAUTH2_ISSUER: &str = "https://auth.x.ai";
/// Production accounts-app origin allowlist — the only origins builds without
/// non-production builds accept. Lives in its own const, referenced by both
/// profiles below, so the frozen-contract test (monorepo CI compiles with
/// that feature enabled) still pins this production-origin const.
const PROD_ACCOUNTS_APP_ORIGINS: &[&str] = &["https://accounts.x.ai"];
/// See the opt-in non-production feature variant above — builds without
/// the feature accept only the production accounts app.
pub fn allowed_accounts_app_origins() -> Vec<String> {
PROD_ACCOUNTS_APP_ORIGINS
.iter()
.map(|o| o.to_string())
.collect()
}
/// Build a CORS layer that accepts requests from the accounts-app deployments
/// listed in [`allowed_accounts_app_origins`] for the given HTTP method.
///
/// Callers can chain additional configuration (e.g. `.allow_headers(...)` or
/// `.allow_private_network(true)`) onto the returned layer.
pub fn accounts_app_cors_layer(method: axum::http::Method) -> tower_http::cors::CorsLayer {
tower_http::cors::CorsLayer::new()
.allow_origin(tower_http::cors::AllowOrigin::list(
allowed_accounts_app_origins()
.iter()
.filter_map(|origin| match origin.parse() {
Ok(value) => Some(value),
Err(_) => {
tracing::warn!(origin, "skipping malformed accounts-app CORS origin");
None
}
}),
))
.allow_methods([method])
}
/// Local-dev OAuth2 issuer (accounts-app running on localhost).
const XAI_OAUTH2_LOCAL_ISSUER: &str = "http://localhost:22255";
const DEFAULT_OAUTH2_REFERRER: &str = "grok-build";
/// Returns `true` when `KIGI_LOCAL_AUTH=1` is set,
/// indicating the local accounts-app should be used as the OAuth2 issuer.
pub fn use_local_auth() -> bool {
std::env::var("KIGI_LOCAL_AUTH")
.map(|v| !v.is_empty() && v != "0")
.unwrap_or(false)
}
/// Returns the active xAI OAuth2 issuer — the local-dev issuer when
/// `KIGI_LOCAL_AUTH=1` is set, otherwise the production issuer.
pub fn xai_oauth2_issuer() -> &'static str {
if use_local_auth() {
XAI_OAUTH2_LOCAL_ISSUER
} else {
XAI_OAUTH2_ISSUER
}
}
/// Returns `true` if `issuer` is a recognised xAI OAuth2 issuer
/// (production **or** local-dev). Use this instead of comparing against
/// [`XAI_OAUTH2_ISSUER`] directly so that local-dev sessions are still
/// treated as first-party xAI auth.
pub fn is_xai_oauth2_issuer(issuer: &str) -> bool {
issuer == XAI_OAUTH2_ISSUER || issuer == XAI_OAUTH2_LOCAL_ISSUER
}
/// auth.json scope key used by the pre-OIDC `grok login --legacy` flow.
/// Matches the key format produced by the original `accounts.x.ai` relay auth.
pub const LEGACY_AUTH_SCOPE: &str = "https://accounts.x.ai/sign-in";
impl GrokComConfig {
/// Whether `xai.api_key` auth is disabled. Pinning a team
/// (`force_login_team_uuid`) implies this — team membership can't be verified
/// from a bare API key, so it must go through IdP login. The
/// `KIGI_DISABLE_API_KEY_AUTH` env lockdown is sticky: because the env value
/// seeds `default()` (the merge base), a lower-trust user `config.toml` could
/// otherwise set `disable_api_key_auth = false` and override it — so the env
/// is OR-ed in here and cannot be turned back off by a user layer. Trusted
/// `requirements.toml` already wins over `config.toml` via layer precedence.
pub fn api_key_auth_disabled(&self) -> bool {
self.disable_api_key_auth == Some(true)
|| self.force_login_team_uuid.is_some()
|| env_lockdown_forced()
}
/// When `preferred_method = api_key`, automatic OIDC paths (devbox mint,
/// interactive browser login, external auth provider) must not run — the
/// pin is fail-closed. Explicit `grok login --devbox` / `--api-key` bypass
/// this by not consulting automatic flow helpers.
pub fn blocks_automatic_oidc(&self) -> bool {
matches!(self.preferred_method, Some(PreferredAuthMethod::ApiKey))
}
/// The auth.json scope key for this config.
pub struct KimiCodeConfig {}
impl KimiCodeConfig {
/// The persisted-credential scope key for this configuration.
pub fn auth_scope(&self) -> String {
if let Some(ref oidc) = self.oidc {
format!("{}::{}", oidc.issuer.trim_end_matches('/'), oidc.client_id)
} else if let Some(ref oauth2) = self.oauth2 {
oauth2.auth_scope()
} else {
unreachable!("oauth2 config is always present (xAI default or env override)")
}
}
}
impl OAuth2ProviderConfig {
pub fn is_team_principal(&self) -> bool {
self.principal_type.as_deref() == Some(TEAM_PRINCIPAL_TYPE)
}
pub fn from_env() -> Option<Self> {
let issuer = std::env::var("KIGI_OAUTH2_ISSUER").ok()?;
let client_id = std::env::var("KIGI_OAUTH2_CLIENT_ID").ok()?;
let principal_type = std::env::var("KIGI_OAUTH2_PRINCIPAL_TYPE").ok();
let principal_id = std::env::var("KIGI_OAUTH2_PRINCIPAL_ID").ok();
let default_scopes = match principal_type.as_deref() {
Some(TEAM_PRINCIPAL_TYPE) => default_team_oauth2_scopes(),
_ => default_oauth2_scopes(),
};
Some(Self {
issuer,
client_id,
scopes: std::env::var("KIGI_OAUTH2_SCOPES")
.map(|s| s.split(',').map(|s| s.trim().to_owned()).collect())
.unwrap_or(default_scopes),
principal_type,
principal_id,
referrer: Some(
std::env::var("KIGI_OAUTH2_REFERRER")
.unwrap_or_else(|_| DEFAULT_OAUTH2_REFERRER.to_owned()),
),
})
}
/// Convert to [`OidcAuthConfig`] to reuse the OIDC login flow.
pub fn as_oidc(&self) -> OidcAuthConfig {
OidcAuthConfig {
issuer: self.issuer.clone(),
client_id: self.client_id.clone(),
scopes: self.scopes.clone(),
audience: None,
}
}
pub fn base_auth_scope(&self) -> String {
format!("{}::{}", self.issuer.trim_end_matches('/'), self.client_id)
}
pub fn auth_scope(&self) -> String {
self.base_auth_scope()
}
}
impl Default for GrokComConfig {
fn default() -> Self {
let oidc = OidcAuthConfig::from_env();
let oauth2 = if oidc.is_some() {
None
} else {
Some(
OAuth2ProviderConfig::from_env().unwrap_or_else(|| OAuth2ProviderConfig {
issuer: xai_oauth2_issuer().to_owned(),
client_id: obfstr::obfstr!("b1a00492-073a-47ea-816f-4c329264a828").to_owned(),
scopes: default_oauth2_scopes(),
principal_type: None,
principal_id: None,
referrer: Some(DEFAULT_OAUTH2_REFERRER.to_owned()),
}),
)
};
Self {
grok_ws_origin: std::env::var("KIGI_WS_ORIGIN")
.unwrap_or_else(|_| AUTH_ORIGIN_DEFAULT.to_owned()),
token_header: "xai-grok-cli".to_owned(),
oidc,
oauth2,
auth_provider_command: std::env::var("KIGI_AUTH_PROVIDER_COMMAND").ok(),
auth_provider_label: std::env::var("KIGI_AUTH_PROVIDER_LABEL").ok(),
auth_token_ttl: std::env::var("KIGI_AUTH_TOKEN_TTL")
.ok()
.and_then(|v| v.parse().ok()),
disable_api_key_auth: std::env::var("KIGI_DISABLE_API_KEY_AUTH")
.ok()
.map(|v| env_flag_enabled(&v)),
force_login_team_uuid: None,
preferred_method: None,
}
}
}
/// Parse a boolean env-var value for grok's on/off flags. A bare presence
/// enables the flag, but the common falsy spellings (`0`, `false`, `off`,
/// `no`, empty) count as disabled — so e.g. `KIGI_DISABLE_API_KEY_AUTH=false`
/// does NOT turn the kill switch on.
fn env_flag_enabled(value: &str) -> bool {
!matches!(
value.trim().to_ascii_lowercase().as_str(),
"" | "0" | "false" | "off" | "no"
)
}
/// True when the admin has set `KIGI_DISABLE_API_KEY_AUTH` to a truthy value in
/// the process environment. Read live (call-time) and OR-ed into
/// `api_key_auth_disabled()` so the env lockdown is non-overridable by a
/// user-layer `config.toml`.
fn env_lockdown_forced() -> bool {
std::env::var("KIGI_DISABLE_API_KEY_AUTH")
.ok()
.is_some_and(|v| env_flag_enabled(&v))
}
impl OidcAuthConfig {
pub fn from_env() -> Option<Self> {
let issuer = std::env::var("KIGI_OIDC_ISSUER").ok()?;
let client_id = std::env::var("KIGI_OIDC_CLIENT_ID").ok()?;
Some(Self {
issuer,
client_id,
scopes: std::env::var("KIGI_OIDC_SCOPES")
.map(|s| s.split(',').map(|s| s.trim().to_owned()).collect())
.unwrap_or_else(|_| default_oidc_scopes()),
audience: std::env::var("KIGI_OIDC_AUDIENCE").ok(),
})
KIMI_CODE_OAUTH_SCOPE.to_owned()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn team_auth_scope_is_base_scope() {
let cfg = OAuth2ProviderConfig {
issuer: "https://auth.x.ai".into(),
client_id: "client-123".into(),
scopes: default_team_oauth2_scopes(),
principal_type: Some("Team".into()),
principal_id: Some("team-abc".into()),
referrer: Some("grok-build".into()),
};
assert_eq!(cfg.auth_scope(), "https://auth.x.ai::client-123");
fn auth_scope_is_the_kimi_code_key() {
assert_eq!(KimiCodeConfig::default().auth_scope(), "oauth/kimi-code");
}
#[test]
fn env_flag_enabled_treats_falsy_spellings_as_off() {
for off in ["", " ", "0", "false", "FALSE", "off", "No", " false "] {
assert!(!env_flag_enabled(off), "{off:?} should be off");
}
for on in ["1", "true", "yes", "on", "enabled"] {
assert!(env_flag_enabled(on), "{on:?} should be on");
}
}
#[test]
fn personal_auth_scope_is_base_scope() {
let cfg = OAuth2ProviderConfig {
issuer: "https://auth.x.ai".into(),
client_id: "client-123".into(),
scopes: default_oauth2_scopes(),
principal_type: None,
principal_id: None,
referrer: Some("grok-build".into()),
};
assert_eq!(cfg.auth_scope(), "https://auth.x.ai::client-123");
}
/// FROZEN loopback contract: the accounts-app origins the CLI's loopback
/// callback server accepts cross-origin requests from. The consent page
/// (served from accounts.x.ai) delivers the code via `fetch(..., cors)`, so
/// removing an origin breaks loopback delivery for already-installed CLIs.
/// Keep in sync with the oauth2-provider / accounts-app deployments.
/// Non-production / local-dev origins are opt-in only.
#[test]
fn allowed_accounts_app_origins_are_frozen() {
assert_eq!(PROD_ACCOUNTS_APP_ORIGINS, &["https://accounts.x.ai"]);
assert_eq!(allowed_accounts_app_origins(), PROD_ACCOUNTS_APP_ORIGINS);
}
/// FROZEN client contract: the 8 scopes the xAI OAuth2 client requests.
/// The server must keep accepting all of them; existing tokens carry
/// exactly this set. Frozen OAuth client scope contract.
#[test]
fn default_oauth2_scopes_are_frozen() {
let scopes = default_oauth2_scopes();
let scopes: Vec<&str> = scopes.iter().map(String::as_str).collect();
assert_eq!(
scopes,
[
"openid",
"profile",
"email",
"offline_access",
"grok-cli:access",
"api:access",
"conversations:read",
"conversations:write",
]
);
}
#[test]
fn preferred_method_deserializes_from_toml() {
let cfg: GrokComConfig = toml::from_str(
r#"
preferred_method = "api_key"
"#,
)
.expect("parse");
assert_eq!(cfg.preferred_method, Some(PreferredAuthMethod::ApiKey));
let cfg: GrokComConfig = toml::from_str(
r#"
preferred_method = "oidc"
"#,
)
.expect("parse");
assert_eq!(cfg.preferred_method, Some(PreferredAuthMethod::Oidc));
let cfg: GrokComConfig = toml::from_str("").expect("parse empty");
assert_eq!(cfg.preferred_method, None);
fn deserializes_from_empty_toml() {
let cfg: KimiCodeConfig = toml::from_str("").expect("empty config parses");
assert_eq!(cfg.auth_scope(), KIMI_CODE_OAUTH_SCOPE);
}
}
crates/codegen/kigi-shell/src/auth/credential_provider.rs
+26 -28
View File
@@ -1,5 +1,5 @@
use crate::auth::AuthManager;
use crate::util::grok_auth_credentials::GrokAuthCredentials;
use crate::util::kigi_auth_credentials::KigiAuthCredentials;
use kigi_auth::{
AuthCredentialProvider, CredentialSnapshot, HttpAuth, StaticAuthCredentialProvider,
};
@@ -7,7 +7,7 @@ use reqwest::RequestBuilder;
use std::sync::Arc;
/// `api_key.id` for the active credential: hash the stable API key, never the
/// OIDC bearer (which rotates). `None` for non-API-key auth.
fn api_key_id_for(auth: Option<&crate::auth::GrokAuth>) -> Option<String> {
fn api_key_id_for(auth: Option<&crate::auth::KimiAuth>) -> Option<String> {
auth.filter(|a| matches!(a.auth_mode, crate::auth::AuthMode::ApiKey))
.map(|a| crate::agent::config::deployment_id_from_key(&a.key))
}
@@ -15,7 +15,7 @@ fn api_key_id_for(auth: Option<&crate::auth::GrokAuth>) -> Option<String> {
/// delegates to `AuthManager::unauthorized_recovery`.
pub struct ShellAuthCredentialProvider {
auth_manager: Arc<AuthManager>,
static_credentials: GrokAuthCredentials,
static_credentials: KigiAuthCredentials,
}
impl ShellAuthCredentialProvider {
pub(crate) fn new(
@@ -23,7 +23,7 @@ impl ShellAuthCredentialProvider {
deployment_key: Option<String>,
alpha_test_key: Option<String>,
) -> Self {
let mut static_credentials = GrokAuthCredentials::new(None);
let mut static_credentials = KigiAuthCredentials::new(None);
static_credentials.deployment_key = deployment_key;
static_credentials.alpha_test_key = alpha_test_key;
Self {
@@ -61,18 +61,19 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider {
};
}
let auth = self.auth_manager.current_or_expired();
let user_id = auth.as_ref().map(|a| a.user_id.clone());
let team_id = auth.as_ref().and_then(|a| a.team_id.clone());
let organization_id = auth.as_ref().and_then(|a| a.organization_id.clone());
// The Kimi token response carries no account info; `user_id` stays
// empty until a later feature surfaces it.
let user_id = auth
.as_ref()
.map(|a| a.user_id.clone())
.filter(|id| !id.is_empty());
let api_key_id = api_key_id_for(auth.as_ref());
let token = auth.map(|a| a.key);
CredentialSnapshot {
token,
user_id,
team_id,
deployment_id: None,
api_key_id,
organization_id,
}
}
async fn refresh_after_unauthorized(&self) -> bool {
@@ -81,15 +82,12 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider {
}
self.auth_manager.try_recover_unauthorized().await
}
fn needs_token_auth_header(&self) -> bool {
self.static_credentials.deployment_key.is_none()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::GrokAuth;
use crate::auth::GrokComConfig;
use crate::auth::KimiAuth;
use crate::auth::KimiCodeConfig;
use crate::auth::manager::AuthManager;
use chrono::{Duration as ChronoDuration, Utc};
use kigi_auth::AuthCredentialProvider;
@@ -128,19 +126,19 @@ mod tests {
}
}
}
fn make_auth(key: &str, expires_in: ChronoDuration) -> GrokAuth {
GrokAuth {
fn make_auth(key: &str, expires_in: ChronoDuration) -> KimiAuth {
KimiAuth {
key: key.to_string(),
user_id: "test-user".to_string(),
create_time: Utc::now(),
expires_at: Some(Utc::now() + expires_in),
..GrokAuth::test_default()
..KimiAuth::test_default()
}
}
/// Build an `AuthManager` rooted at `dir`. Caller keeps `dir` alive for
/// the duration of the test so the `TempDir` `Drop` actually cleans up.
fn make_manager(dir: &tempfile::TempDir, initial: Option<GrokAuth>) -> Arc<AuthManager> {
let mgr = AuthManager::new(dir.path(), GrokComConfig::default());
fn make_manager(dir: &tempfile::TempDir, initial: Option<KimiAuth>) -> Arc<AuthManager> {
let mgr = AuthManager::new(dir.path(), KimiCodeConfig::default());
if let Some(auth) = initial {
mgr.hot_swap(auth);
}
@@ -210,16 +208,16 @@ mod tests {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(
dir.path(),
crate::auth::GrokComConfig::default(),
crate::auth::KimiCodeConfig::default(),
));
mgr.hot_swap(GrokAuth {
mgr.hot_swap(KimiAuth {
key: "stale".into(),
auth_mode: crate::auth::AuthMode::Oidc,
auth_mode: crate::auth::AuthMode::OAuth,
create_time: chrono::Utc::now() - ChronoDuration::hours(2),
user_id: "u".into(),
refresh_token: Some("rt-stale".into()),
expires_at: Some(chrono::Utc::now() - ChronoDuration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
struct OkRefresher {
calls: Arc<std::sync::atomic::AtomicU32>,
@@ -231,14 +229,14 @@ mod tests {
_r: crate::auth::manager::RefreshReason,
) -> crate::auth::refresh::RefreshOutcome {
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth {
crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth {
key: "fresh".into(),
auth_mode: crate::auth::AuthMode::Oidc,
auth_mode: crate::auth::AuthMode::OAuth,
create_time: chrono::Utc::now(),
user_id: "u".into(),
refresh_token: Some("rt-new".into()),
expires_at: Some(chrono::Utc::now() + ChronoDuration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
}))
}
}
@@ -282,11 +280,11 @@ mod tests {
Some(deployment_id_from_key("xai-token-EX").as_str())
);
assert!(dep.api_key_id.is_none());
let api_auth = GrokAuth {
let api_auth = KimiAuth {
key: "sk-apikey-xyz".into(),
auth_mode: crate::auth::AuthMode::ApiKey,
expires_at: Some(Utc::now() + ChronoDuration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
};
let api = ShellAuthCredentialProvider::new(make_manager(&dir, Some(api_auth)), None, None)
.snapshot();
crates/codegen/kigi-shell/src/auth/devbox_login_stub.rs
-37
View File
@@ -1,37 +0,0 @@
//! Stub for builds without the devbox auth feature.
//!
//! Compiled instead of `devbox_login.rs` when the devbox auth feature is
//! off, so the remote devbox login helper is not reached. The API
//! mirrors the real module: `is_devbox_environment()` is always `false`, which
//! short-circuits every auto-recovery/migration call site, and the entry
//! points that can still be reached directly (`grok login --devbox`) return a
//! descriptive error.
use super::manager::AuthManager;
use super::model::GrokAuth;
const UNAVAILABLE: &str =
"devbox login is not available in this build (compiled without the `devbox-login` feature)";
/// Always `false` without the devbox auth feature; callers treat the
/// process as running outside a devbox environment.
pub(crate) fn is_devbox_environment() -> bool {
false
}
/// Unreachable in practice (guarded by [`is_devbox_environment`]); errors
/// defensively if called.
pub(crate) async fn mint_devbox_auth(_auth_manager: &AuthManager) -> anyhow::Result<GrokAuth> {
anyhow::bail!(UNAVAILABLE)
}
/// Unreachable in practice (guarded by [`is_devbox_environment`]); errors
/// defensively if called.
pub(super) async fn mint_devbox_auth_raw() -> anyhow::Result<GrokAuth> {
anyhow::bail!(UNAVAILABLE)
}
/// `grok login --devbox` entry point: always errors in this build.
pub async fn run_devbox_login(_config: &crate::agent::config::Config) -> anyhow::Result<GrokAuth> {
anyhow::bail!(UNAVAILABLE)
}
crates/codegen/kigi-shell/src/auth/device.rs
+297
View File
@@ -0,0 +1,297 @@
//! Device identity headers for the Kimi Code OAuth endpoints.
//!
//! Every OAuth call (device authorization, token poll, refresh) carries three
//! headers identifying this installation (PRD F1):
//!
//! - `X-Msh-Device-Name` — the local hostname
//! - `X-Msh-Device-Model` — an honest local OS/arch string (e.g.
//! "macOS 15.5 arm64"), ported from kimi-cli's `_device_model()`
//! - `X-Msh-Device-Id` — a uuid4 hex persisted at `~/.kigi/device_id`
//! (owner-only), created on first use
//!
//! All values are ASCII-sanitized (ported from kimi-cli's
//! `_ascii_header_value`) since HTTP header values must be ASCII.
use std::path::PathBuf;
use std::sync::OnceLock;
use anyhow::Context as _;
/// Sanitize a header value to ASCII: non-ASCII bytes are dropped; an empty
/// result falls back to `"unknown"`. Port of kimi-cli `_ascii_header_value`.
pub(crate) fn ascii_header_value(value: &str) -> String {
let sanitized: String = value.chars().filter(char::is_ascii).collect();
let trimmed = sanitized.trim();
if trimmed.is_empty() {
"unknown".to_owned()
} else {
trimmed.to_owned()
}
}
/// The three device-identity headers sent on every OAuth call.
///
/// Errors when the persistent device id cannot be created (e.g. read-only
/// `~/.kigi`): the OAuth endpoints require `X-Msh-Device-Id`, so login cannot
/// proceed without it.
pub(crate) fn device_headers() -> anyhow::Result<[(&'static str, String); 3]> {
Ok([
("X-Msh-Device-Name", ascii_header_value(&device_name())),
("X-Msh-Device-Model", ascii_header_value(device_model())),
("X-Msh-Device-Id", ascii_header_value(&device_id()?)),
])
}
/// Local hostname (kimi-cli: `platform.node() or socket.gethostname()`).
fn device_name() -> String {
#[cfg(unix)]
{
let mut buf = [0u8; 256];
// SAFETY: buf is a valid writable buffer of the passed length.
let rc = unsafe { libc::gethostname(buf.as_mut_ptr().cast(), buf.len()) };
if rc == 0 {
let end = buf.iter().position(|&b| b == 0).unwrap_or(buf.len());
let name = String::from_utf8_lossy(&buf[..end]).into_owned();
if !name.trim().is_empty() {
return name;
}
}
"unknown".to_owned()
}
#[cfg(windows)]
{
std::env::var("COMPUTERNAME").unwrap_or_else(|_| "unknown".to_owned())
}
#[cfg(not(any(unix, windows)))]
{
"unknown".to_owned()
}
}
/// Honest local device-model string, computed once per process. Port of
/// kimi-cli `_device_model()`:
/// - macOS → `macOS {product_version} {arch}` (e.g. "macOS 15.5 arm64")
/// - Windows → `Windows {10|11} {arch}` (build ≥ 22000 reports 11)
/// - other → `{sysname} {kernel_release} {machine}`
pub(crate) fn device_model() -> &'static str {
static MODEL: OnceLock<String> = OnceLock::new();
MODEL.get_or_init(compute_device_model)
}
fn compute_device_model() -> String {
#[cfg(target_os = "macos")]
{
// Match Python's platform.machine() spelling on macOS.
let arch = match std::env::consts::ARCH {
"aarch64" => "arm64",
other => other,
};
match macos_product_version() {
Some(version) => format!("macOS {version} {arch}"),
None => format!("macOS {arch}"),
}
}
#[cfg(windows)]
{
let arch = std::env::consts::ARCH;
match windows_release() {
Some(release) => format!("Windows {release} {arch}"),
None => format!("Windows {arch}"),
}
}
#[cfg(not(any(target_os = "macos", windows)))]
{
let (sysname, release, machine) = uname_fields();
match (release, machine) {
(Some(r), Some(m)) => format!("{sysname} {r} {m}"),
(Some(r), None) => format!("{sysname} {r}"),
(None, Some(m)) => format!("{sysname} {m}"),
(None, None) => sysname,
}
}
}
/// macOS product version (e.g. "15.5") from the SystemVersion plist — the
/// same source Python's `platform.mac_ver()` reads.
#[cfg(target_os = "macos")]
fn macos_product_version() -> Option<String> {
let plist = std::fs::read_to_string("/System/Library/CoreServices/SystemVersion.plist").ok()?;
plist_string_value(&plist, "ProductVersion")
}
/// Extract `<key>{key}</key><string>value</string>` from a plist XML body.
#[cfg(target_os = "macos")]
fn plist_string_value(plist: &str, key: &str) -> Option<String> {
let key_tag = format!("<key>{key}</key>");
let after_key = &plist[plist.find(&key_tag)? + key_tag.len()..];
let start = after_key.find("<string>")? + "<string>".len();
let end = after_key.find("</string>")?;
(start <= end).then(|| after_key[start..end].trim().to_owned())
}
/// Windows major release ("10" or "11"), from the build number reported by
/// `cmd /c ver` (kimi-cli: `sys.getwindowsversion().build >= 22000` → 11).
#[cfg(windows)]
fn windows_release() -> Option<String> {
let output = std::process::Command::new("cmd")
.args(["/c", "ver"])
.output()
.ok()?;
let text = String::from_utf8_lossy(&output.stdout);
// "Microsoft Windows [Version 10.0.22631.3155]"
let version = text.split("Version").nth(1)?.trim();
let mut parts = version.trim_end_matches(']').split('.');
let major = parts.next()?.trim().to_owned();
let _minor = parts.next()?;
let build: u32 = parts.next()?.trim().parse().ok()?;
if major == "10" && build >= 22000 {
Some("11".to_owned())
} else {
Some(major)
}
}
/// `uname(2)` sysname / release / machine for Linux and other Unix.
#[cfg(all(unix, not(target_os = "macos")))]
fn uname_fields() -> (String, Option<String>, Option<String>) {
// SAFETY: utsname is a plain-old-data struct; uname fills it in.
let mut uts: libc::utsname = unsafe { std::mem::zeroed() };
if unsafe { libc::uname(&mut uts) } != 0 {
return (std::env::consts::OS.to_owned(), None, None);
}
fn field(raw: &[libc::c_char]) -> Option<String> {
let bytes: Vec<u8> = raw
.iter()
.take_while(|&&c| c != 0)
.map(|&c| c as u8)
.collect();
let s = String::from_utf8_lossy(&bytes).trim().to_owned();
(!s.is_empty()).then_some(s)
}
(
field(&uts.sysname).unwrap_or_else(|| std::env::consts::OS.to_owned()),
field(&uts.release),
field(&uts.machine),
)
}
#[cfg(not(unix))]
#[cfg(not(windows))]
fn uname_fields() -> (String, Option<String>, Option<String>) {
(std::env::consts::OS.to_owned(), None, None)
}
/// Path of the persistent device id: `{kigi_home}/device_id`.
fn device_id_path() -> PathBuf {
kigi_config::kigi_home().join("device_id")
}
/// Persistent uuid4-hex device id, created (owner-only, 0o600) on first use
/// and cached for the process lifetime.
pub(crate) fn device_id() -> anyhow::Result<String> {
static DEVICE_ID: OnceLock<String> = OnceLock::new();
if let Some(id) = DEVICE_ID.get() {
return Ok(id.clone());
}
let id = load_or_create_device_id(&device_id_path())?;
Ok(DEVICE_ID.get_or_init(|| id).clone())
}
/// Read `path`, or mint a uuid4 hex and persist it owner-only.
fn load_or_create_device_id(path: &std::path::Path) -> anyhow::Result<String> {
if let Ok(existing) = std::fs::read_to_string(path) {
let trimmed = existing.trim();
if !trimmed.is_empty() {
return Ok(trimmed.to_owned());
}
}
let id = uuid::Uuid::new_v4().simple().to_string();
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.with_context(|| format!("creating {} for device_id", parent.display()))?;
}
std::fs::write(path, &id)
.with_context(|| format!("writing device id to {}", path.display()))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
.with_context(|| format!("chmod 600 {}", path.display()))?;
}
tracing::info!(path = %path.display(), "auth: created persistent device id");
Ok(id)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn ascii_header_value_passes_ascii_through() {
assert_eq!(ascii_header_value("macOS 15.5 arm64"), "macOS 15.5 arm64");
assert_eq!(ascii_header_value(" padded "), "padded");
}
#[test]
fn ascii_header_value_strips_non_ascii() {
assert_eq!(ascii_header_value("café-host"), "caf-host");
assert_eq!(ascii_header_value("机器"), "unknown");
assert_eq!(ascii_header_value(" "), "unknown");
}
#[test]
fn device_model_is_nonempty_ascii() {
let model = device_model();
assert!(!model.is_empty());
assert!(model.is_ascii(), "device model must be ASCII: {model:?}");
// The honest local OS name must lead the string.
#[cfg(target_os = "macos")]
assert!(model.starts_with("macOS "), "got {model:?}");
#[cfg(windows)]
assert!(model.starts_with("Windows"), "got {model:?}");
}
#[test]
fn load_or_create_device_id_roundtrips_and_is_owner_only() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("device_id");
let created = load_or_create_device_id(&path).unwrap();
assert_eq!(created.len(), 32, "uuid4 hex is 32 chars: {created:?}");
assert!(created.chars().all(|c| c.is_ascii_hexdigit()));
// Second call reads the same id back.
let reread = load_or_create_device_id(&path).unwrap();
assert_eq!(created, reread);
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mode = std::fs::metadata(&path).unwrap().permissions().mode();
assert_eq!(mode & 0o777, 0o600, "device_id must be owner-only");
}
}
#[test]
fn load_or_create_device_id_ignores_empty_file() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("device_id");
std::fs::write(&path, " \n").unwrap();
let created = load_or_create_device_id(&path).unwrap();
assert_eq!(created.len(), 32);
}
#[cfg(target_os = "macos")]
#[test]
fn plist_string_value_extracts_product_version() {
let plist = r#"<?xml version="1.0"?>
<dict>
<key>ProductBuildVersion</key>
<string>24F74</string>
<key>ProductVersion</key>
<string>15.5</string>
</dict>"#;
assert_eq!(
plist_string_value(plist, "ProductVersion").as_deref(),
Some("15.5")
);
assert_eq!(plist_string_value(plist, "Missing"), None);
}
}
crates/codegen/kigi-shell/src/auth/device_code.rs
+282 -812
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/error.rs
+20 -45
View File
@@ -3,30 +3,21 @@ use thiserror::Error;
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum AuthError {
#[error("Not logged in. Run `grok login`.")]
#[error("Not logged in. Run `kigi login`.")]
NotLoggedIn,
/// Token expired and no refresh authority available.
#[error("Token expired. Run `grok login` to re-authenticate.")]
#[error("Token expired. Run `kigi login` to re-authenticate.")]
TokenExpiredNoRefresh,
/// Server rejected the token (401) with no recovery path.
#[error("Authentication rejected by server. Run `grok login` to re-authenticate.")]
#[error("Authentication rejected by server. Run `kigi login` to re-authenticate.")]
ServerRejectedNoRecovery,
/// All recovery strategies exhausted.
#[error("Auth recovery exhausted; re-authentication required.")]
RecoveryExhausted,
/// A session's team principal violates the `force_login_team_uuid` pin.
/// `message` states which team is required vs. returned.
#[error("{message} Run `grok login` to sign in with the required team.")]
PinnedTeamMismatch { message: String },
/// Cached API-key session rejected because API-key auth is disabled.
#[error("API-key auth is disabled by your administrator. Run `grok login` to authenticate.")]
ApiKeyAuthDisabled,
/// Outcome of a refresh-authority attempt. Recoverability (and, for
/// permanent failures, the reason) lives in [`RefreshTokenError`].
#[error(transparent)]
@@ -38,7 +29,7 @@ pub enum AuthError {
/// caller must make, so a future third state should break consumers loudly.
#[derive(Debug, Error)]
pub enum RefreshTokenError {
/// The credential is dead; the user must re-authenticate.
/// The credential was rejected; the tombstone cooldown gates re-attempts.
#[error(transparent)]
Permanent(#[from] RefreshTokenFailedError),
/// Network / 5xx / unknown blip; safe to retry later. Carries the cause.
@@ -48,12 +39,10 @@ pub enum RefreshTokenError {
/// A retryable refresh failure, wrapping its cause. No public `From`:
/// construct only via [`AuthError::transient`] /
/// [`AuthError::transient_source`], so a stray `?` on some error can't silently
/// classify a permanent failure as retryable (mirrors the dedicated
/// [`RefreshTokenFailedError`] on the permanent arm). Display frames the cause
/// as an auth-refresh failure so internal messages (lock timeout, sleep defer)
/// don't surface bare; the permanent arm derives its copy from
/// [`RefreshTokenFailedReason::user_message`] and is not prefixed.
/// [`AuthError::transient_source`], so a stray `?` on some error can't
/// silently classify a permanent failure as retryable. Display frames the
/// cause as an auth-refresh failure so internal messages (lock timeout,
/// sleep defer) don't surface bare.
#[derive(Debug, Error)]
#[error("auth refresh failed: {0}")]
pub struct RefreshTransientError(#[source] Box<dyn std::error::Error + Send + Sync>);
@@ -74,45 +63,31 @@ impl From<RefreshTokenFailedReason> for RefreshTokenFailedError {
}
}
/// Why a token refresh terminally failed, grounded in the OAuth2 error codes
/// our IdP actually emits.
/// Why a token refresh terminally failed. Both reasons carry the same
/// tombstone semantics (PRD F1): a 300s cooldown scoped to the rejected
/// refresh token, auto-cleared when the persisted refresh token differs
/// (another process rotated) or a fresh login lands.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum RefreshTokenFailedReason {
/// `invalid_grant` — the refresh token is no longer valid (expired, reused,
/// or revoked; the IdP does not distinguish these).
/// The OAuth host answered 401/403 — the refresh token is no longer
/// valid (expired, reused, or revoked).
RefreshTokenRejected,
/// `invalid_client` — the client/app credential was rejected.
ClientRejected,
/// Escalation from repeated transient failures (OIDC) or a single
/// external-binary failure. Never a raw IdP code: an unrecognized terminal
/// code is classified transient, not `Other` (see `classify_terminal`).
/// Non-retryable terminal failure that isn't an explicit rejection
/// (malformed payload, unexpected 4xx).
Other,
}
impl RefreshTokenFailedReason {
/// Sticky until the credential changes (never ages out): a revoked refresh
/// token never self-heals, whereas client rotation / transient escalation
/// recover, so those age out past the TTL.
pub(crate) fn is_sticky(self) -> bool {
match self {
Self::RefreshTokenRejected => true,
Self::ClientRejected | Self::Other => false,
}
}
/// User-facing copy for a terminal refresh failure; the raw IdP code stays
/// in logs.
/// User-facing copy for a terminal refresh failure; the raw wire detail
/// stays in logs.
pub(crate) fn user_message(self) -> &'static str {
match self {
Self::RefreshTokenRejected => {
"Your session has expired. Run `grok login` to sign in again."
}
Self::ClientRejected => {
"Authentication is temporarily unavailable. Run `grok login` if this persists."
"Your session has expired. Run `kigi login` to sign in again."
}
Self::Other => {
"Authentication could not be refreshed. Run `grok login` to sign in again."
"Authentication could not be refreshed. Run `kigi login` to sign in again."
}
}
}
crates/codegen/kigi-shell/src/auth/external_auth.rs
-283
View File
@@ -1,283 +0,0 @@
use crate::auth::{AuthMode, GrokAuth};
#[derive(serde::Deserialize)]
pub(crate) struct ExternalAuthOutput {
pub access_token: String,
#[serde(default)]
pub refresh_token: Option<String>,
#[serde(default)]
pub expires_in: Option<u64>,
/// Token issuer. An xAI issuer marks the credential as first-party;
/// see [`GrokAuth::is_xai_auth`].
#[serde(default)]
pub issuer: Option<String>,
}
/// Parse process output (stdout) into a `GrokAuth`. Accepts bare token or JSON.
pub(crate) fn parse_output(output: &std::process::Output) -> anyhow::Result<GrokAuth> {
if !output.status.success() {
anyhow::bail!("exited with {}", output.status);
}
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_owned();
if stdout.is_empty() {
anyhow::bail!("produced no output on stdout");
}
let (token, refresh_token, expires_at, issuer) =
if let Ok(parsed) = serde_json::from_str::<ExternalAuthOutput>(&stdout) {
tracing::debug!(
has_refresh_token = parsed.refresh_token.is_some(),
expires_in = ?parsed.expires_in,
issuer = ?parsed.issuer,
"auth: parsed external provider output as JSON"
);
let expires_at = parsed
.expires_in
.map(|secs| chrono::Utc::now() + chrono::Duration::seconds(secs as i64));
let issuer = parsed
.issuer
.map(|i| i.trim().to_owned())
.filter(|i| !i.is_empty());
(
parsed.access_token,
parsed.refresh_token,
expires_at,
issuer,
)
} else {
tracing::debug!(
stdout_len = stdout.len(),
"auth: treating output as bare token"
);
(stdout, None, None, None)
};
Ok(GrokAuth {
key: token,
auth_mode: AuthMode::External,
create_time: chrono::Utc::now(),
user_id: String::new(),
email: None,
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: None,
team_name: None,
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token,
expires_at,
oidc_issuer: issuer,
oidc_client_id: None,
})
}
/// Sync version for mid-session refresh. 5s timeout for refresh, 60s for initial.
pub(crate) fn run_external_auth_sync(command: &str, is_refresh: bool) -> Option<GrokAuth> {
use std::process::{Command, Stdio};
let timeout_secs = if is_refresh { 5 } else { 60 };
tracing::info!(cmd = %command, is_refresh, timeout_secs, "auth: running external auth provider (sync)");
let mut cmd = Command::new("sh");
cmd.args(["-c", command])
.stdin(Stdio::null())
.stdout(Stdio::piped())
// Pipe stderr — inherit would corrupt the TUI alternate screen.
.stderr(Stdio::piped());
if is_refresh {
cmd.env("KIGI_AUTH_EXPIRED", "1");
}
kigi_tools::util::detach_std_command(&mut cmd);
cmd.envs(kigi_tools::util::pager_env());
let mut child = cmd.spawn()
.map_err(|e| {
tracing::warn!(error = %e, cmd = %command, "auth: failed to start external auth provider");
e
})
.ok()?;
let timeout = std::time::Duration::from_secs(timeout_secs);
let start = std::time::Instant::now();
loop {
match child.try_wait() {
Ok(Some(_status)) => break,
Ok(None) => {
if start.elapsed() > timeout {
tracing::warn!(
cmd = %command,
timeout_secs,
"auth: external auth provider timed out (likely needs interactive auth), killing"
);
let _ = child.kill();
let _ = child.wait();
return None;
}
std::thread::sleep(std::time::Duration::from_millis(100));
}
Err(e) => {
tracing::warn!(error = %e, "auth: error waiting for external auth provider");
return None;
}
}
}
let output = child
.wait_with_output()
.map_err(|e| {
tracing::warn!(error = %e, "auth: failed to read external auth provider output");
e
})
.ok()?;
match parse_output(&output) {
Ok(auth) => {
tracing::info!("auth: external auth provider returned fresh token");
Some(auth)
}
Err(e) => {
tracing::warn!(error = %e, "auth: external auth provider failed");
None
}
}
}
/// Run external auth provider, carrying forward `/user`-derived fields from previous auth.
pub(crate) fn refresh_with_command(command: &str, prev_auth: &GrokAuth) -> Option<GrokAuth> {
let mut auth = run_external_auth_sync(command, true)?;
auth.carry_user_profile_from(prev_auth);
Some(auth)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_output_nonzero_exit_is_err() {
let output = std::process::Output {
status: std::process::Command::new("false").status().unwrap(),
stdout: b"token".to_vec(),
stderr: vec![],
};
assert!(parse_output(&output).is_err());
}
#[test]
fn parse_output_empty_stdout_is_err() {
let output = std::process::Output {
status: std::process::Command::new("true").status().unwrap(),
stdout: b" \n".to_vec(),
stderr: vec![],
};
assert!(parse_output(&output).is_err());
}
#[test]
fn parse_output_issuer_claim_enables_xai_auth() {
let ok = |stdout: &str| std::process::Output {
status: std::process::Command::new("true").status().unwrap(),
stdout: stdout.as_bytes().to_vec(),
stderr: vec![],
};
// x.ai issuer claim → first-party session (relay-eligible).
let auth = parse_output(&ok(
r#"{"access_token":"t","expires_in":900,"issuer":"https://auth.x.ai"}"#,
))
.unwrap();
assert_eq!(auth.oidc_issuer.as_deref(), Some("https://auth.x.ai"));
assert!(auth.is_xai_auth());
// Non-x.ai issuer is stored but stays third-party.
let auth = parse_output(&ok(
r#"{"access_token":"t","issuer":"https://idp.acme.example"}"#,
))
.unwrap();
assert_eq!(
auth.oidc_issuer.as_deref(),
Some("https://idp.acme.example")
);
assert!(!auth.is_xai_auth());
// Missing / empty / whitespace issuer → None.
let auth = parse_output(&ok(r#"{"access_token":"t"}"#)).unwrap();
assert_eq!(auth.oidc_issuer, None);
assert!(!auth.is_xai_auth());
let auth = parse_output(&ok(r#"{"access_token":"t","issuer":" "}"#)).unwrap();
assert_eq!(auth.oidc_issuer, None);
// Bare-token output never carries an issuer.
let auth = parse_output(&ok("bare-token")).unwrap();
assert_eq!(auth.oidc_issuer, None);
assert!(!auth.is_xai_auth());
}
#[test]
fn parse_output_malformed_json_falls_back_to_bare() {
let output = std::process::Output {
status: std::process::Command::new("true").status().unwrap(),
stdout: b"{not valid json}".to_vec(),
stderr: vec![],
};
let auth = parse_output(&output).unwrap();
assert_eq!(auth.key, "{not valid json}");
}
#[test]
fn sync_spawn_failure_returns_none() {
assert!(run_external_auth_sync("/nonexistent/binary", false).is_none());
}
#[test]
fn sync_sets_grok_auth_expired_env_on_refresh() {
let auth = run_external_auth_sync("echo $KIGI_AUTH_EXPIRED", true).unwrap();
assert_eq!(auth.key, "1");
}
#[test]
fn refresh_carries_zdr_flags_forward() {
let prev = GrokAuth {
user_blocked_reason: Some("BLOCKED_REASON_OTHER".into()),
team_blocked_reasons: vec!["BLOCKED_REASON_NO_LOGS".into()],
coding_data_retention_opt_out: true,
organization_id: Some("org-1".into()),
..GrokAuth::test_default()
};
let auth = refresh_with_command("echo fresh-token", &prev).unwrap();
assert_eq!(auth.key, "fresh-token");
assert!(auth.is_zdr_team(), "ZDR flag must survive refresh");
assert!(auth.coding_data_retention_opt_out);
assert_eq!(
auth.user_blocked_reason.as_deref(),
Some("BLOCKED_REASON_OTHER")
);
assert_eq!(auth.user_id, "test-user", "profile must survive refresh");
assert_eq!(auth.organization_id.as_deref(), Some("org-1"));
}
#[test]
fn sync_refresh_interactive_times_out() {
// Binary writes link to stderr then blocks — 5s refresh timeout kills it.
let cmd = r#"echo 'Visit http://example.com/auth' >&2; sleep 20; echo token"#;
let start = std::time::Instant::now();
let result = run_external_auth_sync(cmd, true);
let elapsed = start.elapsed();
assert!(result.is_none(), "should timeout and return None");
assert!(
elapsed.as_secs() < 10,
"refresh should use 5s timeout, not 60s (took {}s)",
elapsed.as_secs()
);
}
}
crates/codegen/kigi-shell/src/auth/flow.rs
+163 -1595
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/jwt.rs
-45
View File
@@ -1,45 +0,0 @@
//! JWT expiration detection. Returns `None`/`false` for non-JWT tokens.
use chrono::{DateTime, Duration, Utc};
use serde::Deserialize;
#[derive(Deserialize)]
struct Claims {
exp: Option<i64>,
}
pub fn parse_jwt_expiration(token: &str) -> Option<DateTime<Utc>> {
jsonwebtoken::dangerous::insecure_decode::<Claims>(token)
.ok()
.and_then(|data| data.claims.exp)
.and_then(|ts| DateTime::from_timestamp(ts, 0))
}
pub fn is_jwt_expired_or_near(token: &str, threshold: Duration) -> bool {
parse_jwt_expiration(token)
.map(|exp| exp <= Utc::now() + threshold)
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
/// Tokens with an `aud` claim must parse successfully.
/// `jsonwebtoken::Validation::default()` enables audience validation which
/// silently rejects these tokens unless `validate_aud = false` is set.
#[test]
fn parses_jwt_with_aud_claim() {
let token = build_test_jwt(r#"{"aud":["some-audience"],"exp":1772575524}"#);
let exp = parse_jwt_expiration(&token);
assert_eq!(exp.unwrap().timestamp(), 1772575524);
}
fn build_test_jwt(payload_json: &str) -> String {
use base64::Engine;
let enc = base64::engine::general_purpose::URL_SAFE_NO_PAD;
let header = enc.encode(r#"{"alg":"RS256","typ":"JWT"}"#);
let payload = enc.encode(payload_json);
format!("{header}.{payload}.fake-signature")
}
}
crates/codegen/kigi-shell/src/auth/kimi_oauth.rs
+626
View File
@@ -0,0 +1,626 @@
//! Kimi Code OAuth wire protocol (PRD F1).
//!
//! Three calls against `{host}` (= `kigi_env::oauth_host()`), all
//! `application/x-www-form-urlencoded` POSTs carrying the device-identity
//! headers from [`super::device`]:
//!
//! - `POST /api/oauth/device_authorization` — form `client_id`
//! - `POST /api/oauth/token` (poll) — form `client_id` + `device_code` +
//! `grant_type=urn:ietf:params:oauth:grant-type:device_code`
//! - `POST /api/oauth/token` (refresh) — form `client_id` +
//! `grant_type=refresh_token` + `refresh_token`, with exponential backoff
//! over the retryable statuses {429, 500, 502, 503, 504} (3 tries) and
//! 401/403 mapped to [`RefreshError::Unauthorized`].
//!
//! Ported from kimi-cli `auth/oauth.py` (the authoritative reference).
use chrono::{Duration, Utc};
use serde::Deserialize;
use super::device::device_headers;
use super::model::{AuthMode, KimiAuth};
/// Kimi Code OAuth client id (fixed for the official device-flow client).
pub(crate) const KIMI_CODE_CLIENT_ID: &str = "17e5f671-d194-4dfb-9706-5516cb48c098";
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];
/// Result of `POST /api/oauth/device_authorization`.
#[derive(Debug, Clone)]
pub struct DeviceAuthorization {
pub user_code: String,
pub device_code: String,
/// Bare verification page (may be absent; the complete URI is required).
pub verification_uri: Option<String>,
/// Verification page with the user code pre-filled — what we display
/// and open in the browser.
pub verification_uri_complete: String,
/// Device-code lifetime; `None` when the server omits it.
pub expires_in: Option<i64>,
/// Poll interval in seconds (server default 5; floored at 1 by callers).
pub interval: i64,
}
#[derive(Deserialize)]
struct DeviceAuthorizationResponse {
user_code: String,
device_code: String,
#[serde(default)]
verification_uri: Option<String>,
verification_uri_complete: String,
#[serde(default)]
expires_in: Option<i64>,
#[serde(default)]
interval: Option<i64>,
}
/// Successful token payload (device grant and refresh grant share it).
#[derive(Debug, Deserialize)]
pub(crate) struct TokenResponse {
pub access_token: String,
pub refresh_token: String,
pub expires_in: i64,
#[serde(default)]
pub scope: Option<String>,
#[serde(default)]
pub token_type: Option<String>,
}
impl TokenResponse {
/// Materialize the credential: `expires_at = now + expires_in`.
pub(crate) fn into_auth(self) -> KimiAuth {
let now = Utc::now();
KimiAuth {
key: self.access_token,
auth_mode: AuthMode::OAuth,
create_time: now,
user_id: String::new(),
email: None,
refresh_token: Some(self.refresh_token),
expires_at: Some(now + Duration::seconds(self.expires_in)),
expires_in: Some(self.expires_in),
scope: self.scope,
token_type: self.token_type,
}
}
}
#[derive(Deserialize, Default)]
struct OAuthErrorBody {
#[serde(default)]
error: Option<String>,
#[serde(default)]
error_description: Option<String>,
}
/// One poll tick against the token endpoint.
#[derive(Debug)]
pub(crate) enum DevicePollResult {
/// 200 with an access token — login complete.
Success(Box<KimiAuth>),
/// `error == "expired_token"` — restart the whole device authorization.
Expired,
/// Any other non-200 outcome (`authorization_pending`, `slow_down`,
/// unknown errors) — wait and poll again. `slow_down` additionally bumps
/// the caller's interval.
Pending {
error: String,
description: Option<String>,
},
}
fn oauth_url(host: &str, path: &str) -> String {
format!("{}{path}", host.trim_end_matches('/'))
}
/// Attach the device-identity headers to a request.
fn with_device_headers(
mut builder: reqwest::RequestBuilder,
) -> anyhow::Result<reqwest::RequestBuilder> {
for (name, value) in device_headers()? {
builder = builder.header(name, value);
}
Ok(builder)
}
/// Defend against control characters / non-https redirects from a
/// compromised or mis-configured OAuth host.
fn validate_verification_uri(uri: &str) -> anyhow::Result<()> {
if uri.chars().any(|c| c.is_ascii_control()) {
anyhow::bail!("Server returned invalid verification URI");
}
let parsed = url::Url::parse(uri)
.map_err(|_| anyhow::anyhow!("Server returned invalid verification URI"))?;
match parsed.scheme() {
"https" => Ok(()),
"http" if matches!(parsed.host_str(), Some("localhost") | Some("127.0.0.1")) => Ok(()),
_ => anyhow::bail!("Server returned unsupported verification URI scheme"),
}
}
/// `POST {host}/api/oauth/device_authorization` — start a device login.
pub(crate) async fn request_device_authorization(
host: &str,
) -> anyhow::Result<DeviceAuthorization> {
let url = oauth_url(host, "/api/oauth/device_authorization");
tracing::info!(url = %url, "auth: requesting device authorization");
let resp = with_device_headers(crate::http::shared_client().post(&url))?
.form(&[("client_id", KIMI_CODE_CLIENT_ID)])
.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");
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-])");
}
validate_verification_uri(&parsed.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"
);
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: parsed.verification_uri_complete,
expires_in: parsed.expires_in.filter(|&e| e > 0),
interval: parsed.interval.unwrap_or(5),
})
}
/// One poll of `POST {host}/api/oauth/token` with the device grant.
///
/// 5xx and network/decode failures are errors (kimi-cli parity: the login
/// loop surfaces them); everything else maps onto [`DevicePollResult`].
pub(crate) async fn poll_device_token(
host: &str,
device_code: &str,
) -> anyhow::Result<DevicePollResult> {
let url = oauth_url(host, "/api/oauth/token");
let resp = with_device_headers(crate::http::shared_client().post(&url))?
.form(&[
("client_id", KIMI_CODE_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");
return Ok(DevicePollResult::Success(Box::new(tokens.into_auth())));
}
// 200 without an access token: treat as still-pending (kimi-cli
// requires "access_token" in the payload before accepting).
tracing::warn!("auth: device poll returned 200 without access_token; continuing");
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");
return Ok(DevicePollResult::Expired);
}
tracing::debug!(error = %error, "auth: device poll pending");
Ok(DevicePollResult::Pending {
error,
description: err.error_description,
})
}
/// Why a refresh call terminally or transiently failed.
#[derive(Debug, thiserror::Error)]
pub(crate) enum RefreshError {
/// 401/403 — the refresh token was rejected. Triggers the tombstone
/// cooldown in the manager.
#[error("token refresh unauthorized (HTTP {status}): {description}")]
Unauthorized { status: u16, description: String },
/// Non-retryable non-200 status.
#[error("token refresh failed (HTTP {status}): {description}")]
Fatal { status: u16, description: String },
/// Retry budget exhausted over retryable statuses / network blips.
#[error("token refresh failed after {MAX_REFRESH_RETRIES} attempts: {last_error}")]
Exhausted { last_error: String },
/// Local failure before the wire (e.g. device-id creation failed).
#[error(transparent)]
Local(#[from] anyhow::Error),
}
/// `POST {host}/api/oauth/token` with `grant_type=refresh_token`.
///
/// Retries the retryable statuses and network errors with exponential
/// backoff (`2^attempt` seconds); 401/403 returns immediately as
/// [`RefreshError::Unauthorized`].
pub(crate) async fn refresh_token(
host: &str,
refresh_token: &str,
) -> Result<KimiAuth, RefreshError> {
let url = oauth_url(host, "/api/oauth/token");
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"
);
tokio::time::sleep(backoff).await;
}
tracing::info!(attempt, "auth: token refresh attempt");
let send_result = with_device_headers(crate::http::shared_client().post(&url))?
.form(&[
("client_id", KIMI_CODE_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 wiremock::matchers::{body_string_contains, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
fn token_json(access: &str, refresh: &str) -> serde_json::Value {
serde_json::json!({
"access_token": access,
"refresh_token": refresh,
"expires_in": 3600,
"scope": "kimi-code",
"token_type": "bearer",
})
}
#[tokio::test]
async fn device_authorization_parses_wire_payload() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/device_authorization"))
.and(body_string_contains(format!(
"client_id={KIMI_CODE_CLIENT_ID}"
)))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "ABCD-1234",
"device_code": "dev-code-1",
"verification_uri": "https://auth.kimi.com/device",
"verification_uri_complete": "https://auth.kimi.com/device?code=ABCD-1234",
"expires_in": 600,
"interval": 7,
})))
.expect(1)
.mount(&server)
.await;
let auth = request_device_authorization(&server.uri()).await.unwrap();
assert_eq!(auth.user_code, "ABCD-1234");
assert_eq!(auth.device_code, "dev-code-1");
assert_eq!(auth.interval, 7);
assert_eq!(auth.expires_in, Some(600));
assert_eq!(
auth.verification_uri_complete,
"https://auth.kimi.com/device?code=ABCD-1234"
);
}
#[tokio::test]
async fn device_authorization_sends_device_headers() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/device_authorization"))
.and(wiremock::matchers::header_exists("X-Msh-Device-Name"))
.and(wiremock::matchers::header_exists("X-Msh-Device-Model"))
.and(wiremock::matchers::header_exists("X-Msh-Device-Id"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "AAAA",
"device_code": "d",
"verification_uri_complete": "https://auth.kimi.com/device?code=AAAA",
"interval": 5,
})))
.expect(1)
.mount(&server)
.await;
request_device_authorization(&server.uri()).await.unwrap();
}
#[tokio::test]
async fn device_authorization_defaults_interval_to_five() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/device_authorization"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "AAAA",
"device_code": "d",
"verification_uri_complete": "https://auth.kimi.com/device?code=AAAA",
})))
.mount(&server)
.await;
let auth = request_device_authorization(&server.uri()).await.unwrap();
assert_eq!(auth.interval, 5);
assert_eq!(auth.expires_in, None);
}
#[tokio::test]
async fn device_authorization_rejects_bad_verification_uri() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/device_authorization"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "AAAA",
"device_code": "d",
"verification_uri_complete": "javascript:alert(1)",
})))
.mount(&server)
.await;
let err = request_device_authorization(&server.uri())
.await
.unwrap_err();
assert!(err.to_string().contains("verification URI"), "{err}");
}
#[tokio::test]
async fn device_authorization_surfaces_http_errors() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/device_authorization"))
.respond_with(ResponseTemplate::new(400).set_body_string("nope"))
.mount(&server)
.await;
let err = request_device_authorization(&server.uri())
.await
.unwrap_err();
assert!(err.to_string().contains("HTTP 400"), "{err}");
}
#[tokio::test]
async fn poll_success_builds_auth_with_expiry() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.and(body_string_contains("grant_type=urn"))
.and(body_string_contains("device_code=dev-1"))
.respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-1", "rt-1")))
.expect(1)
.mount(&server)
.await;
let result = poll_device_token(&server.uri(), "dev-1").await.unwrap();
let DevicePollResult::Success(auth) = result else {
panic!("expected success, got {result:?}");
};
assert_eq!(auth.key, "at-1");
assert_eq!(auth.refresh_token.as_deref(), Some("rt-1"));
assert_eq!(auth.expires_in, Some(3600));
let remaining = auth.expires_at.unwrap() - chrono::Utc::now();
assert!(
(3590..=3600).contains(&remaining.num_seconds()),
"expires_at must be ~now+expires_in, got {remaining:?}"
);
assert_eq!(auth.scope.as_deref(), Some("kimi-code"));
assert_eq!(auth.token_type.as_deref(), Some("bearer"));
}
#[tokio::test]
async fn poll_maps_expired_token_to_restart() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(
ResponseTemplate::new(400)
.set_body_json(serde_json::json!({ "error": "expired_token" })),
)
.mount(&server)
.await;
let result = poll_device_token(&server.uri(), "dev-1").await.unwrap();
assert!(matches!(result, DevicePollResult::Expired), "{result:?}");
}
#[tokio::test]
async fn poll_maps_pending_and_unknown_errors_to_pending() {
let server = MockServer::start().await;
for error in ["authorization_pending", "slow_down", "surprise_error"] {
server.reset().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(
ResponseTemplate::new(400).set_body_json(serde_json::json!({ "error": error })),
)
.mount(&server)
.await;
let result = poll_device_token(&server.uri(), "dev-1").await.unwrap();
match result {
DevicePollResult::Pending { error: got, .. } => assert_eq!(got, error),
other => panic!("expected pending for {error}, got {other:?}"),
}
}
}
#[tokio::test]
async fn poll_server_error_is_an_error() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(502))
.mount(&server)
.await;
let err = poll_device_token(&server.uri(), "dev-1").await.unwrap_err();
assert!(err.to_string().contains("server error"), "{err}");
}
#[tokio::test]
async fn refresh_success_round_trip() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.and(body_string_contains("grant_type=refresh_token"))
.and(body_string_contains("refresh_token=rt-old"))
.respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-new", "rt-new")))
.expect(1)
.mount(&server)
.await;
let auth = refresh_token(&server.uri(), "rt-old").await.unwrap();
assert_eq!(auth.key, "at-new");
assert_eq!(auth.refresh_token.as_deref(), Some("rt-new"));
}
#[tokio::test]
async fn refresh_401_maps_to_unauthorized_without_retry() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(
ResponseTemplate::new(401).set_body_json(
serde_json::json!({ "error_description": "refresh token revoked" }),
),
)
.expect(1)
.mount(&server)
.await;
let err = refresh_token(&server.uri(), "rt-dead").await.unwrap_err();
match err {
RefreshError::Unauthorized {
status,
description,
} => {
assert_eq!(status, 401);
assert_eq!(description, "refresh token revoked");
}
other => panic!("expected Unauthorized, got {other:?}"),
}
}
#[tokio::test]
async fn refresh_retries_retryable_status_then_succeeds() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(503))
.up_to_n_times(1)
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-2", "rt-2")))
.expect(1)
.mount(&server)
.await;
let auth = refresh_token(&server.uri(), "rt-old").await.unwrap();
assert_eq!(auth.key, "at-2");
}
#[tokio::test]
async fn refresh_exhausts_after_three_retryable_failures() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(500))
.expect(3)
.mount(&server)
.await;
let err = refresh_token(&server.uri(), "rt-old").await.unwrap_err();
assert!(matches!(err, RefreshError::Exhausted { .. }), "{err:?}");
}
#[tokio::test]
async fn refresh_non_retryable_status_is_fatal_without_retry() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(
ResponseTemplate::new(400)
.set_body_json(serde_json::json!({ "error_description": "bad request" })),
)
.expect(1)
.mount(&server)
.await;
let err = refresh_token(&server.uri(), "rt-old").await.unwrap_err();
match err {
RefreshError::Fatal {
status,
description,
} => {
assert_eq!(status, 400);
assert_eq!(description, "bad request");
}
other => panic!("expected Fatal, got {other:?}"),
}
}
}
crates/codegen/kigi-shell/src/auth/manager.rs
+383 -710
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/manager/enrichment.rs
-256
View File
@@ -1,256 +0,0 @@
//! Background `/user` enrichment spawned by `AuthManager::update()`.
use std::sync::Arc;
use std::time::Duration as StdDuration;
use super::AuthManager;
use super::lock::try_lock_auth_file_async;
use crate::auth::manager::AUTH_LOCK_TIMEOUT;
use crate::auth::model::{GrokAuth, UserInfo, lookup_auth};
use crate::auth::storage::{read_auth_json, write_auth_json};
/// `/user` fetch budget, shared by the inline (login) and background paths.
const USER_FETCH_TIMEOUT: StdDuration = StdDuration::from_secs(10);
/// Logs `auth update enrichment dropped` if the task is cancelled
/// mid-flight. Disarmed on normal completion.
pub(super) struct EnrichmentExitGuard {
pub(super) started: std::time::Instant,
pub(super) armed: bool,
}
impl EnrichmentExitGuard {
pub(super) fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for EnrichmentExitGuard {
fn drop(&mut self) {
if !self.armed {
return;
}
kigi_log::unified_log::warn(
"auth update enrichment dropped",
None,
Some(serde_json::json!({
"elapsed_ms": self.started.elapsed().as_millis() as u64,
})),
);
}
}
pub(super) fn spawn(manager: Arc<AuthManager>, auth: GrokAuth) {
tokio::spawn(async move {
let mut exit_guard = EnrichmentExitGuard {
started: std::time::Instant::now(),
armed: true,
};
run_user_info_enrichment(&manager, auth).await;
exit_guard.disarm();
});
}
async fn fetch_user_info(manager: &AuthManager, key: &str, log_label: &str) -> Option<UserInfo> {
let user_url = format!("{}/user", manager.proxy_base_url);
let token_header = &manager.grok_com_config.token_header;
let started = std::time::Instant::now();
let http_client = crate::http::shared_client();
let response = http_client
.get(&user_url)
.timeout(USER_FETCH_TIMEOUT)
.header("Authorization", format!("Bearer {}", key))
.header("X-XAI-Token-Auth", token_header.as_str())
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
)
.send()
.await;
match response {
Ok(resp) if resp.status().is_success() => match resp.json::<UserInfo>().await {
Ok(ui) if !ui.user_id.is_empty() => Some(ui),
Ok(_) => {
kigi_log::unified_log::warn(
&format!("{log_label} skipped"),
None,
Some(serde_json::json!({
"reason": "empty_user_id",
"elapsed_ms": started.elapsed().as_millis() as u64,
})),
);
None
}
Err(e) => {
kigi_log::unified_log::warn(
&format!("{log_label} failed"),
None,
Some(serde_json::json!({
"reason": "parse",
"error": e.to_string(),
"elapsed_ms": started.elapsed().as_millis() as u64,
})),
);
None
}
},
Ok(resp) => {
kigi_log::unified_log::warn(
&format!("{log_label} failed"),
None,
Some(serde_json::json!({
"reason": "http_status",
"http_status": resp.status().as_u16(),
"elapsed_ms": started.elapsed().as_millis() as u64,
})),
);
None
}
Err(e) => {
kigi_log::unified_log::warn(
&format!("{log_label} failed"),
None,
Some(serde_json::json!({
"reason": if e.is_timeout() { "timeout" } else { "transport" },
"error": e.to_string(),
"elapsed_ms": started.elapsed().as_millis() as u64,
})),
);
None
}
}
}
/// Blocking login-time enrichment: merge `/user` fields before the first save.
pub(super) async fn enrich_inline(manager: &AuthManager, auth: &mut GrokAuth) {
let Some(ui) = fetch_user_info(manager, &auth.key, "auth login enrichment").await else {
return;
};
apply_user_info_enrichment(auth, ui);
}
async fn run_user_info_enrichment(manager: &AuthManager, auth: GrokAuth) {
let started = std::time::Instant::now();
let Some(user_info) = fetch_user_info(manager, &auth.key, "auth update enrichment").await
else {
return;
};
let user_elapsed_ms = started.elapsed().as_millis() as u64;
// R-M-W file lock. On timeout, fall through to an unlocked write
// rather than drop the enrichment.
let lock_started = std::time::Instant::now();
let lock_guard = try_lock_auth_file_async(&manager.path, AUTH_LOCK_TIMEOUT).await;
let lock_wait_ms = lock_started.elapsed().as_millis() as u64;
if lock_guard.is_none() {
tracing::warn!("auth: enrichment proceeding without auth.json.lock");
}
let Ok(mut map) = read_auth_json(&manager.path) else {
kigi_log::unified_log::warn(
"auth update enrichment skipped",
None,
Some(serde_json::json!({ "reason": "read_disk_failed" })),
);
return;
};
let Some(mut disk) = lookup_auth(&map, &manager.scope) else {
kigi_log::unified_log::info(
"auth update enrichment skipped",
None,
Some(serde_json::json!({ "reason": "no_disk_auth" })),
);
return;
};
// Sibling-stomp guard. If either the access token or refresh
// token on disk differs from the one we wrote, a sibling process
// rotated tokens since our update(). Skip enrichment to avoid
// writing stale profile data over the sibling's fresher entry.
//
// OR logic (not AND): a single-field rotation (key changes, RT
// stays) is the common case during concurrent refresh. The old
// AND logic required ALL three fields to differ, letting
// single-field rotations through.
//
// Team-login transitions (placeholder→real user_id) don't rotate
// tokens, so OR correctly allows enrichment for that case.
if disk.key != auth.key || disk.refresh_token != auth.refresh_token {
kigi_log::unified_log::info(
"auth update enrichment skipped",
None,
Some(serde_json::json!({
"reason": "sibling_rotated",
"written_key_prefix": crate::auth::token_suffix(&auth.key),
"disk_key_prefix": crate::auth::token_suffix(&disk.key),
})),
);
return;
}
apply_user_info_enrichment(&mut disk, user_info);
map.insert(manager.scope.clone(), disk.clone());
let write_started = std::time::Instant::now();
if let Err(e) = write_auth_json(&manager.path, &map) {
kigi_log::unified_log::error(
"auth update enrichment write failed",
None,
Some(serde_json::json!({
"error": e.to_string(),
"user_ms": user_elapsed_ms,
"lock_wait_ms": lock_wait_ms,
"write_ms": write_started.elapsed().as_millis() as u64,
})),
);
return;
}
manager.with_inner_write(|inner| *inner = Some(disk));
kigi_log::unified_log::info(
"auth update enrichment done",
None,
Some(serde_json::json!({
"user_ms": user_elapsed_ms,
"lock_wait_ms": lock_wait_ms,
"write_ms": write_started.elapsed().as_millis() as u64,
"total_ms": started.elapsed().as_millis() as u64,
})),
);
}
/// Merge enrichment fields into disk auth. Does not touch token fields.
pub(super) fn apply_user_info_enrichment(disk: &mut GrokAuth, user_info: UserInfo) {
disk.user_id = user_info.user_id;
disk.first_name = user_info.first_name.or(disk.first_name.take());
disk.last_name = user_info.last_name.or(disk.last_name.take());
disk.profile_image_asset_id = user_info
.profile_image_asset_id
.or(disk.profile_image_asset_id.take());
disk.principal_type = user_info.principal_type.or(disk.principal_type.take());
disk.principal_id = user_info.principal_id.or(disk.principal_id.take());
disk.team_id = user_info.team_id.or(disk.team_id.take());
disk.team_name = user_info.team_name.or(disk.team_name.take());
disk.team_role = user_info.team_role.or(disk.team_role.take());
disk.organization_id = user_info.organization_id.or(disk.organization_id.take());
disk.organization_name = user_info
.organization_name
.or(disk.organization_name.take());
disk.organization_role = user_info
.organization_role
.or(disk.organization_role.take());
disk.user_blocked_reason = user_info
.user_blocked_reason
.or(disk.user_blocked_reason.take());
if let Some(reasons) = user_info.team_blocked_reasons {
disk.team_blocked_reasons = reasons;
}
if let Some(opt_out) = user_info.coding_data_retention_opt_out {
disk.coding_data_retention_opt_out = opt_out;
}
if let Some(ref email) = user_info.email
&& !email.is_empty()
{
disk.email = user_info.email;
}
}
crates/codegen/kigi-shell/src/auth/manager_tests.rs
+463 -3802
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/meta.rs
+3 -19
View File
@@ -1,6 +1,8 @@
use serde::{Deserialize, Serialize};
/// Access gate from `grok_build_access_gate`.
/// Access-gate copy resolved from remote settings (message + optional CTA).
/// Auth no longer produces gates (tier gating was an xAI concept); the pager
/// still renders one when remote settings carry a gate message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GateInfo {
pub message: String,
@@ -17,24 +19,6 @@ pub struct AuthMeta {
pub email: Option<String>,
#[serde(default)]
pub auth_mode: Option<String>,
/// Team principal UUID when the session is a team login (`None` for personal).
#[serde(default)]
pub team_id: Option<String>,
#[serde(default)]
pub team_name: Option<String>,
#[serde(default)]
pub is_zdr: bool,
#[serde(default)]
pub team_role: Option<String>,
#[serde(default)]
pub coding_data_retention_opt_out: bool,
#[serde(default)]
pub show_resolved_model: Option<bool>,
/// `Some` = user is blocked; `None` = user has access.
#[serde(default)]
pub gate: Option<GateInfo>,
/// User-friendly display name for the current subscription tier
/// (e.g. "SuperGrok Heavy", "X Premium", "Free"). From CCP `/settings`.
#[serde(default)]
pub subscription_tier: Option<String>,
}
crates/codegen/kigi-shell/src/auth/mod.rs
+9 -21
View File
@@ -1,42 +1,30 @@
pub(crate) mod attribution;
mod config;
pub mod credential_provider;
#[path = "devbox_login_stub.rs"]
pub(crate) mod devbox_login;
pub(crate) mod device;
pub mod device_code;
pub mod error;
mod external_auth;
mod flow;
mod jwt;
pub(crate) mod kimi_oauth;
pub(crate) mod manager;
mod model;
pub mod oidc;
pub(crate) mod recovery;
pub(crate) mod refresh;
mod storage;
pub(crate) mod token_type;
pub(crate) use config::LEGACY_AUTH_SCOPE;
pub use config::{
ForceLoginTeam, GrokComConfig, OAuth2ProviderConfig, OidcAuthConfig, PreferredAuthMethod,
XAI_OAUTH2_ISSUER, is_xai_oauth2_issuer, xai_oauth2_issuer,
};
pub(crate) use external_auth::{parse_output, refresh_with_command};
pub(crate) use flow::{
AuthChannels, run_auth_flow, run_auth_flow_with_stderr_bridge,
try_ensure_session_noninteractive,
};
pub use config::{KIMI_CODE_OAUTH_SCOPE, KimiCodeConfig};
pub(crate) use flow::try_ensure_session_noninteractive;
pub use flow::{
AuthUrlInfo, AuthUrlMode, LoginTransportOverride, LogoutResult, ensure_authenticated,
ensure_authenticated_or_noninteractive, ensure_authenticated_with_override, perform_logout,
run_cli_login, run_cli_logout, try_ensure_fresh_auth,
AuthChannels, AuthUrlInfo, AuthUrlMode, LogoutResult, ensure_authenticated,
ensure_authenticated_or_noninteractive, perform_logout, run_auth_flow,
run_auth_flow_with_stderr_bridge, run_cli_login, run_cli_logout, try_ensure_fresh_auth,
};
pub use jwt::{is_jwt_expired_or_near, parse_jwt_expiration};
mod meta;
pub use error::{AuthError, RefreshTokenError, RefreshTokenFailedReason};
pub use manager::{AuthManager, shared_api_key_provider};
pub use meta::{AuthMeta, GateInfo};
pub use model::{AuthMode, GrokAuth, lookup_auth};
pub(crate) use model::{TOKEN_TTL, UserInfo, is_expired, token_suffix};
pub use model::{AuthMode, KimiAuth, lookup_auth};
pub(crate) use model::{TOKEN_TTL, is_expired, token_suffix};
pub use storage::{
clear_api_key, read_api_key, read_auth_json, read_token_by_scope, store_api_key,
};
crates/codegen/kigi-shell/src/auth/model.rs
+171 -364
View File
@@ -1,108 +1,75 @@
//! Kimi Code auth data model: the persisted token set + expiry policy.
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use super::is_xai_oauth2_issuer;
/// Fallback TTL for credentials without a server-provided expiry
/// (plain API keys).
pub(crate) const TOKEN_TTL: Duration = Duration::days(30);
const DEFAULT_EARLY_INVALIDATION_SECS: u64 = 300; // 5 minutes
/// Legacy auth.json scope key. Fallback for old devbox auth files.
pub(super) const LEGACY_SCOPE: &str = "https://accounts.x.ai/sign-in";
/// Minimum refresh threshold (PRD F1): refresh when the remaining lifetime
/// drops below `max(300, expires_in × 0.5)` seconds.
const DEFAULT_EARLY_INVALIDATION_SECS: u64 = 300;
/// auth.json scope key for plain API key auth (desktop login, `grok login --api-key`).
pub const API_KEY_SCOPE: &str = "xai::api_key";
/// Fraction of `expires_in` that drives the dynamic refresh threshold.
const REFRESH_THRESHOLD_RATIO: f64 = 0.5;
const BLOCKED_REASON_NO_LOGS: &str = "BLOCKED_REASON_NO_LOGS";
const BLOCKED_REASON_NO_LOGS_MODERATED: &str = "BLOCKED_REASON_NO_LOGS_MODERATED";
/// auth.json scope key for plain API key auth (`kigi login --api-key`, F2).
pub const API_KEY_SCOPE: &str = "kigi::api_key";
/// Token provenance (debugging/auth.json only -- no code branches on this).
/// How this credential was obtained.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum AuthMode {
/// Deprecated. Kept for deserializing old auth.json files.
#[serde(alias = "grok")]
WebLogin,
/// OIDC or OAuth2 interactive login via customer IdP
#[serde(alias = "oidc")]
Oidc,
/// External auth provider binary
External,
/// Plain API key (e.g. from grok-desktop login or `grok login --api-key`)
/// Kimi Code subscription OAuth (device-code flow).
#[serde(rename = "oauth")]
OAuth,
/// Plain API key.
ApiKey,
}
/// Wire value of `principal_type` for team OAuth principals (capitalized by
/// the auth service). Single source for every comparison site.
pub(crate) const TEAM_PRINCIPAL_TYPE: &str = "Team";
/// The Kimi Code credential: the OAuth token set (or a bare API key) plus
/// local bookkeeping. The Kimi token response carries no user info; `user_id`
/// / `email` stay empty until a later feature surfaces account info.
#[derive(Clone, Serialize, Deserialize)]
pub struct GrokAuth {
pub struct KimiAuth {
/// The bearer sent on API calls (`Authorization: Bearer {key}`):
/// the OAuth access token, or the API key in `ApiKey` mode.
pub key: String,
pub auth_mode: AuthMode,
pub create_time: DateTime<Utc>,
pub user_id: String,
pub email: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub first_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub last_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub profile_image_asset_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub principal_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub principal_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub team_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub team_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub team_role: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub organization_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub organization_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub organization_role: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub user_blocked_reason: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub team_blocked_reasons: Vec<String>,
/// Account id — the Kimi token response has none; empty until a later
/// feature surfaces it.
#[serde(default)]
pub coding_data_retention_opt_out: bool,
/// Deprecated. Kept for deserializing existing auth.json files.
pub user_id: String,
/// Account email — the Kimi token response has none; `None` until a
/// later feature surfaces it.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub has_grok_code_access: Option<bool>,
/// Refresh token (OIDC/OAuth2 or external provider).
pub email: Option<String>,
/// OAuth refresh token; `None` for API keys.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub refresh_token: Option<String>,
/// Server-provided expiration (from OIDC `expires_in`).
/// When present, takes precedence over the hardcoded `TOKEN_TTL`.
/// `create_time + expires_in`, computed when the token was minted.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub expires_at: Option<DateTime<Utc>>,
/// Issuer URL that issued this token. For OIDC credentials it drives
/// refresh via discovery; for external-provider credentials it is the
/// provider's `issuer` claim. In both modes an x.ai issuer marks the
/// credential first-party (`is_xai_auth`).
/// Server-reported token lifetime in seconds; drives the dynamic
/// refresh threshold `max(300, expires_in × 0.5)`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub oidc_issuer: Option<String>,
/// OIDC client_id used to obtain this token (needed for refresh).
pub expires_in: Option<i64>,
/// OAuth scope string as returned by the token endpoint.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub oidc_client_id: Option<String>,
pub scope: Option<String>,
/// Token type as returned by the token endpoint (e.g. "bearer").
#[serde(default, skip_serializing_if = "Option::is_none")]
pub token_type: Option<String>,
}
impl std::fmt::Debug for GrokAuth {
impl std::fmt::Debug for KimiAuth {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GrokAuth")
f.debug_struct("KimiAuth")
.field("key", &token_suffix(&self.key))
.field("auth_mode", &self.auth_mode)
.field("user_id", &self.user_id)
.field("expires_at", &self.expires_at)
.field(
"refresh_token",
@@ -112,128 +79,44 @@ impl std::fmt::Debug for GrokAuth {
}
}
impl GrokAuth {
impl KimiAuth {
/// Seconds since this credential was minted. Negative when the local
/// clock stepped back past `create_time` (NTP correction, VM restore, or
/// a sibling machine's clock via an adopted auth.json) — `create_time`
/// is always stamped from the minting machine's local clock.
/// clock stepped back past `create_time` (NTP correction, VM restore).
pub(crate) fn mint_age_seconds(&self) -> i64 {
Utc::now()
.signed_duration_since(self.create_time)
.num_seconds()
}
/// `true` when the token comes from a first-party xAI account —
/// either an OIDC login against https://auth.x.ai (or the local-dev
/// equivalent), or an external auth provider that declared an xAI
/// issuer for its token.
///
/// The issuer is a client-side hint, not a trust assertion: everything
/// it unlocks still authenticates the actual token server-side, and it
/// never influences endpoints.
pub fn is_xai_auth(&self) -> bool {
match self.auth_mode {
AuthMode::Oidc | AuthMode::External => self
.oidc_issuer
.as_deref()
.is_some_and(is_xai_oauth2_issuer),
AuthMode::ApiKey | AuthMode::WebLogin => false,
}
}
/// `true` when this auth can access grok.com managed MCP connectors.
pub fn is_managed_mcp_eligible(&self) -> bool {
self.is_xai_auth() || self.auth_mode == AuthMode::WebLogin
}
/// Whether this credential can access `supported_in_api: false` models.
///
/// Session logins (WebLogin, OIDC — including enterprise issuers) always
/// qualify; external-provider credentials qualify only when first-party
/// (`is_xai_auth`), matching the built-in devbox login they replace.
/// Plain API keys never do.
/// `true` for a refreshable subscription session (vs a bare API key).
pub fn is_session_auth(&self) -> bool {
match self.auth_mode {
AuthMode::WebLogin | AuthMode::Oidc => true,
AuthMode::External => self.is_xai_auth(),
AuthMode::ApiKey => false,
}
}
pub fn is_team_principal(&self) -> bool {
self.principal_type.as_deref() == Some(TEAM_PRINCIPAL_TYPE) && self.team_id.is_some()
}
/// `true` when the team has Zero Data Retention (ZDR) enabled.
pub fn is_zdr_team(&self) -> bool {
self.team_blocked_reasons
.iter()
.any(|r| r == BLOCKED_REASON_NO_LOGS || r == BLOCKED_REASON_NO_LOGS_MODERATED)
}
/// `true` when the team has ZDR or the user opted out of coding data
/// retention. Use this for trace-upload and research-data gates.
/// Product analytics (`telemetry_enabled`) and user-facing sync
/// features should use `is_zdr_team()` directly.
pub fn is_data_collection_disabled(&self) -> bool {
self.is_zdr_team() || self.coding_data_retention_opt_out
}
/// Carry `/user`-derived fields from a previous auth so refresh rebuilds don't drop them.
pub(crate) fn carry_user_profile_from(&mut self, prev: &GrokAuth) {
self.user_id = prev.user_id.clone();
self.email = prev.email.clone();
self.principal_type = prev.principal_type.clone();
self.principal_id = prev.principal_id.clone();
self.team_id = prev.team_id.clone();
self.team_name = prev.team_name.clone();
self.team_role = prev.team_role.clone();
self.organization_id = prev.organization_id.clone();
self.organization_name = prev.organization_name.clone();
self.organization_role = prev.organization_role.clone();
self.user_blocked_reason = prev.user_blocked_reason.clone();
self.team_blocked_reasons = prev.team_blocked_reasons.clone();
self.coding_data_retention_opt_out = prev.coding_data_retention_opt_out;
self.auth_mode == AuthMode::OAuth
}
}
impl Default for GrokAuth {
impl Default for KimiAuth {
fn default() -> Self {
Self {
key: String::new(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
create_time: Utc::now(),
user_id: String::new(),
email: None,
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: None,
team_name: None,
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token: None,
expires_at: None,
oidc_issuer: None,
oidc_client_id: None,
expires_in: None,
scope: None,
token_type: None,
}
}
}
#[cfg(test)]
impl GrokAuth {
/// Returns a `GrokAuth` with sensible defaults for tests. Override fields
/// with struct update syntax:
impl KimiAuth {
/// A `KimiAuth` with sensible defaults for tests. Override fields with
/// struct update syntax:
/// ```ignore
/// GrokAuth { key: "my-key".into(), ..GrokAuth::test_default() }
/// KimiAuth { key: "my-key".into(), ..KimiAuth::test_default() }
/// ```
pub fn test_default() -> Self {
Self {
@@ -244,82 +127,24 @@ impl GrokAuth {
}
}
pub(crate) type AuthStore = BTreeMap<String, GrokAuth>;
pub(crate) type AuthStore = BTreeMap<String, KimiAuth>;
/// User information from the cli-chat-proxy `GET /v1/user` endpoint.
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct UserInfo {
pub(crate) user_id: String,
#[serde(default)]
pub(super) email: Option<String>,
#[serde(default)]
pub(super) first_name: Option<String>,
#[serde(default)]
pub(super) last_name: Option<String>,
#[serde(default)]
pub(super) profile_image_asset_id: Option<String>,
#[serde(default)]
pub(super) principal_type: Option<String>,
#[serde(default)]
pub(super) principal_id: Option<String>,
#[serde(default)]
pub(super) team_id: Option<String>,
#[serde(default)]
pub(super) team_name: Option<String>,
#[serde(default)]
pub(super) team_role: Option<String>,
#[serde(default)]
pub(super) organization_id: Option<String>,
#[serde(default)]
pub(super) organization_name: Option<String>,
#[serde(default)]
pub(super) organization_role: Option<String>,
#[serde(default)]
pub(super) user_blocked_reason: Option<String>,
#[serde(default)]
pub(super) team_blocked_reasons: Option<Vec<String>>,
#[serde(default)]
pub(super) coding_data_retention_opt_out: Option<bool>,
/// Live subscription tier from the backend (only present when
/// `?include=subscription` is passed to `/user`).
#[serde(default)]
pub(crate) subscription_tier: Option<String>,
}
/// Last 12 chars of a token string, safe for diagnostic logging.
/// Uses the tail because JWT access tokens all share the same base64
/// header prefix (`eyJ0eXAiOiJh…`); the tail (signature bytes) is
/// unique per token and makes `key_changed` / `is_stale_snapshot`
/// diagnostics meaningful.
/// Last 12 chars of a token string, safe for diagnostic logging. Uses the
/// tail because token prefixes are shared across a family; the tail is
/// unique per token and makes `key_changed` diagnostics meaningful.
pub(crate) fn token_suffix(t: &str) -> &str {
let len = t.len();
if len > 12 { &t[len - 12..] } else { t }
}
/// Look up auth from the store by scope key.
///
/// Legacy `WebLogin` tokens (from the pre-OIDC `grok login --legacy`
/// flow) are skipped — they are validated via a per-request DB lookup
/// server-side which fails at high volume. Skipping them here forces
/// affected users to re-authenticate via OIDC on next launch.
pub fn lookup_auth(map: &AuthStore, scope: &str) -> Option<GrokAuth> {
let auth = map.get(scope).cloned().or_else(|| {
if scope == LEGACY_SCOPE {
None
} else {
map.get(LEGACY_SCOPE).cloned()
}
})?;
if auth.auth_mode == AuthMode::WebLogin {
tracing::info!("auth: ignoring legacy WebLogin token — re-authentication required");
return None;
}
Some(auth)
pub fn lookup_auth(map: &AuthStore, scope: &str) -> Option<KimiAuth> {
map.get(scope).cloned()
}
/// Early-invalidation buffer. Override with `KIGI_AUTH_EARLY_INVALIDATION_SECS`
/// for testing (e.g. `=5` to shrink the buffer to 5 seconds).
/// Minimum refresh-threshold component. Override with
/// `KIGI_AUTH_EARLY_INVALIDATION_SECS` for testing (e.g. `=5` to shrink the
/// buffer to 5 seconds).
pub(super) fn early_invalidation() -> Duration {
std::env::var("KIGI_AUTH_EARLY_INVALIDATION_SECS")
.ok()
@@ -328,14 +153,30 @@ pub(super) fn early_invalidation() -> Duration {
.unwrap_or_else(|| Duration::seconds(DEFAULT_EARLY_INVALIDATION_SECS as i64))
}
pub(crate) fn is_expired(auth: &GrokAuth) -> bool {
is_expired_with_buffer(auth, early_invalidation())
/// Dynamic refresh threshold (PRD F1): `max(min_threshold, expires_in × 0.5)`
/// where `min_threshold` defaults to 300s. Credentials without a positive
/// `expires_in` use the minimum alone.
pub(crate) fn refresh_threshold(auth: &KimiAuth) -> Duration {
let min = early_invalidation();
match auth.expires_in {
Some(expires_in) if expires_in > 0 => {
let ratio = Duration::seconds((expires_in as f64 * REFRESH_THRESHOLD_RATIO) as i64);
std::cmp::max(min, ratio)
}
_ => min,
}
}
/// Whether the credential is inside its refresh threshold (i.e. should be
/// treated as expiring-soon for refresh scheduling).
pub(crate) fn is_expired(auth: &KimiAuth) -> bool {
is_expired_with_buffer(auth, refresh_threshold(auth))
}
/// Like [`is_expired`] but with an explicit pre-expiry buffer. Pass
/// `Duration::zero()` for actual (hard) expiry — the instant the token would
/// really be rejected on the wire, with no early-invalidation margin.
pub(crate) fn is_expired_with_buffer(auth: &GrokAuth, buffer: Duration) -> bool {
/// really be rejected on the wire.
pub(crate) fn is_expired_with_buffer(auth: &KimiAuth, buffer: Duration) -> bool {
if let Some(expires_at) = auth.expires_at {
Utc::now() >= (expires_at - buffer)
} else {
@@ -348,142 +189,108 @@ pub(crate) fn is_expired_with_buffer(auth: &GrokAuth, buffer: Duration) -> bool
mod tests {
use super::*;
fn make_auth(mode: AuthMode) -> GrokAuth {
GrokAuth {
key: "k".into(),
auth_mode: mode,
create_time: Utc::now(),
user_id: "u".into(),
email: None,
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: None,
team_name: None,
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token: None,
expires_at: None,
oidc_issuer: None,
oidc_client_id: None,
fn auth_with_lifetime(expires_in: i64, remaining_secs: i64) -> KimiAuth {
KimiAuth {
expires_in: Some(expires_in),
expires_at: Some(Utc::now() + Duration::seconds(remaining_secs)),
refresh_token: Some("rt".into()),
..KimiAuth::test_default()
}
}
/// PRD threshold math: `max(300, expires_in × 0.5)`.
#[test]
fn is_xai_auth_matrix() {
use crate::auth::XAI_OAUTH2_ISSUER;
let with_issuer = |mode: AuthMode, issuer: Option<&str>| GrokAuth {
oidc_issuer: issuer.map(str::to_owned),
..make_auth(mode)
fn refresh_threshold_is_max_of_min_and_half_life() {
// Short-lived token: the 300s floor wins (600 × 0.5 = 300 → tie; 400 × 0.5 = 200 < 300).
let short = auth_with_lifetime(400, 400);
assert_eq!(refresh_threshold(&short).num_seconds(), 300);
// Long-lived token: half the lifetime wins (7200 × 0.5 = 3600).
let long = auth_with_lifetime(7200, 7200);
assert_eq!(refresh_threshold(&long).num_seconds(), 3600);
// No expires_in: the floor alone.
let bare = KimiAuth::test_default();
assert_eq!(refresh_threshold(&bare).num_seconds(), 300);
// Non-positive expires_in must not produce a negative threshold.
let broken = KimiAuth {
expires_in: Some(-5),
..KimiAuth::test_default()
};
assert_eq!(refresh_threshold(&broken).num_seconds(), 300);
}
// Only Oidc/External qualify, and only with an x.ai issuer.
assert!(with_issuer(AuthMode::Oidc, Some(XAI_OAUTH2_ISSUER)).is_xai_auth());
assert!(with_issuer(AuthMode::External, Some(XAI_OAUTH2_ISSUER)).is_xai_auth());
assert!(!with_issuer(AuthMode::Oidc, None).is_xai_auth());
assert!(!with_issuer(AuthMode::External, None).is_xai_auth());
assert!(!with_issuer(AuthMode::Oidc, Some("https://idp.acme.example")).is_xai_auth());
assert!(!with_issuer(AuthMode::External, Some("https://idp.acme.example")).is_xai_auth());
/// A token past its dynamic threshold counts as expiring-soon while a
/// token comfortably before it does not.
#[test]
fn is_expired_uses_dynamic_threshold() {
// 7200s lifetime → threshold 3600s. 3000s remaining < 3600 → expiring.
assert!(is_expired(&auth_with_lifetime(7200, 3000)));
// 5000s remaining > 3600 → fresh.
assert!(!is_expired(&auth_with_lifetime(7200, 5000)));
// Hard expiry ignores the buffer entirely.
assert!(!is_expired_with_buffer(
&auth_with_lifetime(7200, 3000),
Duration::zero()
));
assert!(is_expired_with_buffer(
&auth_with_lifetime(7200, -1),
Duration::zero()
));
}
// ApiKey / WebLogin stay false even with an x.ai issuer set.
assert!(!with_issuer(AuthMode::ApiKey, Some(XAI_OAUTH2_ISSUER)).is_xai_auth());
assert!(!with_issuer(AuthMode::WebLogin, Some(XAI_OAUTH2_ISSUER)).is_xai_auth());
/// Credentials without `expires_at` (API keys) age out via the 30-day TTL.
#[test]
fn no_expiry_falls_back_to_token_ttl() {
let fresh = KimiAuth::test_default();
assert!(!is_expired(&fresh));
let old = KimiAuth {
create_time: Utc::now() - Duration::days(31),
..KimiAuth::test_default()
};
assert!(is_expired(&old));
}
#[test]
fn is_session_auth_requires_first_party_for_external() {
use crate::auth::XAI_OAUTH2_ISSUER;
let with_issuer = |mode: AuthMode, issuer: Option<&str>| GrokAuth {
oidc_issuer: issuer.map(str::to_owned),
..make_auth(mode)
fn lookup_auth_finds_scope_entry() {
let mut map = AuthStore::new();
map.insert("oauth/kimi-code".into(), KimiAuth::test_default());
assert!(lookup_auth(&map, "oauth/kimi-code").is_some());
assert!(lookup_auth(&map, "other").is_none());
}
#[test]
fn debug_redacts_tokens() {
let auth = KimiAuth {
key: "super-secret-access-token".into(),
refresh_token: Some("super-secret-refresh-token".into()),
..KimiAuth::test_default()
};
let debug = format!("{auth:?}");
assert!(!debug.contains("super-secret-access-token"));
assert!(!debug.contains("super-secret-refresh-token"));
}
// Session logins qualify regardless of issuer (incl. enterprise OIDC).
assert!(with_issuer(AuthMode::WebLogin, None).is_session_auth());
assert!(with_issuer(AuthMode::Oidc, None).is_session_auth());
assert!(with_issuer(AuthMode::Oidc, Some("https://idp.acme.example")).is_session_auth());
// External qualifies only when first-party (devbox-login parity).
assert!(with_issuer(AuthMode::External, Some(XAI_OAUTH2_ISSUER)).is_session_auth());
assert!(!with_issuer(AuthMode::External, None).is_session_auth());
#[test]
fn serde_roundtrip_preserves_token_set() {
let auth = KimiAuth {
key: "at".into(),
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now()),
expires_in: Some(3600),
scope: Some("kimi-code".into()),
token_type: Some("bearer".into()),
..KimiAuth::test_default()
};
let json = serde_json::to_string(&auth).unwrap();
assert!(
!with_issuer(AuthMode::External, Some("https://idp.acme.example")).is_session_auth()
json.contains("\"oauth\""),
"wire spelling is \"oauth\": {json}"
);
// Plain API keys never do.
assert!(!with_issuer(AuthMode::ApiKey, Some(XAI_OAUTH2_ISSUER)).is_session_auth());
}
#[test]
fn lookup_auth_skips_weblogin_on_primary_scope() {
let mut map = AuthStore::new();
map.insert("scope".into(), make_auth(AuthMode::WebLogin));
assert!(lookup_auth(&map, "scope").is_none());
}
#[test]
fn lookup_auth_skips_weblogin_on_legacy_fallback() {
let mut map = AuthStore::new();
map.insert(LEGACY_SCOPE.into(), make_auth(AuthMode::WebLogin));
assert!(lookup_auth(&map, "other-scope").is_none());
}
#[test]
fn lookup_auth_returns_oidc_token() {
let mut map = AuthStore::new();
map.insert("scope".into(), make_auth(AuthMode::Oidc));
assert!(lookup_auth(&map, "scope").is_some());
}
#[test]
fn lookup_auth_returns_api_key_token() {
let mut map = AuthStore::new();
map.insert("scope".into(), make_auth(AuthMode::ApiKey));
assert!(lookup_auth(&map, "scope").is_some());
}
/// subscriptionTier present → deserializes to Some.
#[test]
fn user_info_subscription_tier_present() {
let json = r#"{
"userId": "u1",
"subscriptionTier": "SuperGrokPro"
}"#;
let info: UserInfo = serde_json::from_str(json).unwrap();
assert_eq!(info.subscription_tier.as_deref(), Some("SuperGrokPro"));
}
/// subscriptionTier absent → deserializes to None (backwards compat).
#[test]
fn user_info_subscription_tier_absent() {
let json = r#"{"userId": "u1"}"#;
let info: UserInfo = serde_json::from_str(json).unwrap();
assert!(info.subscription_tier.is_none());
}
/// subscriptionTier null → deserializes to None.
#[test]
fn user_info_subscription_tier_null() {
let json = r#"{"userId": "u1", "subscriptionTier": null}"#;
let info: UserInfo = serde_json::from_str(json).unwrap();
assert!(info.subscription_tier.is_none());
}
/// subscriptionTier empty string → deserializes to Some("").
/// The paywall poller treats this as "no subscription" (line 230:
/// `Some(tier) if !tier.is_empty()`) and keeps polling.
#[test]
fn user_info_subscription_tier_empty_string() {
let json = r#"{"userId": "u1", "subscriptionTier": ""}"#;
let info: UserInfo = serde_json::from_str(json).unwrap();
assert_eq!(info.subscription_tier.as_deref(), Some(""));
let back: KimiAuth = serde_json::from_str(&json).unwrap();
assert_eq!(back.key, "at");
assert_eq!(back.refresh_token.as_deref(), Some("rt"));
assert_eq!(back.expires_in, Some(3600));
assert_eq!(back.scope.as_deref(), Some("kimi-code"));
assert_eq!(back.token_type.as_deref(), Some("bearer"));
assert_eq!(back.auth_mode, AuthMode::OAuth);
}
}
crates/codegen/kigi-shell/src/auth/oidc/login.rs
-701
View File
@@ -1,701 +0,0 @@
//! Interactive login orchestration: callback HTTP server, browser
//! handoff, stdin paste fallback, race between the two.
//!
//! Cross-references [`super::protocol`] for OIDC mechanics and
//! [`super::super::AuthManager`] for credential persistence.
use std::collections::HashMap;
use std::io::IsTerminal;
use std::sync::Arc;
use axum::{
Router,
extract::{Query, State},
http::{Method, StatusCode},
response::Html,
routing::get,
};
use tokio::net::TcpListener;
use super::super::config::{GrokComConfig, OidcAuthConfig};
use super::super::{AuthManager, GrokAuth};
use super::protocol::{
OidcError, build_authorize_url, build_grok_auth, discover, enforce_login_principal,
exchange_code, extract_user_info, generate_pkce, login_principal_policy,
peek_access_token_principal, peek_access_token_principal_id, validate_state,
};
/// Maximum time to wait for the browser OAuth callback (or manual paste of the code).
/// 10 minutes is long enough for users who step away briefly during login.
const AUTH_CALLBACK_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600);
/// Parse user-pasted input into `(code, state)`.
///
/// Accepts two formats:
/// 1. Full callback URL: `http://127.0.0.1:PORT/callback?code=XXX&state=YYY`
/// 2. Bare authorization code: `abc123`
fn parse_pasted_input(input: &str) -> Result<Callback, OidcError> {
let input = input.trim();
if input.is_empty() {
return Err(OidcError::InvalidPastedInput("empty input".into()));
}
if let Ok(url) = url::Url::parse(input) {
let params: HashMap<String, String> = url.query_pairs().into_owned().collect();
if let Some(code) = params.get("code") {
let state = params.get("state").cloned().unwrap_or_default();
return Ok(Callback {
code: code.clone(),
state,
});
}
if let Some(error) = params.get("error") {
let desc = params.get("error_description").cloned().unwrap_or_default();
return Err(OidcError::CallbackAuthFailed(if desc.is_empty() {
error.clone()
} else {
format!("{error}: {desc}")
}));
}
return Err(OidcError::InvalidPastedInput(
"URL has no 'code' query parameter".into(),
));
}
Ok(Callback {
code: input.to_owned(),
state: String::new(),
})
}
/// Render a styled callback page shown in the browser after the OAuth redirect.
pub(crate) fn callback_page(title: &str, message: &str, is_success: bool) -> String {
let icon = if is_success {
// Grok logo
r#"<svg xmlns="http://www.w3.org/2000/svg" width="48" height="48" fill="none" viewBox="0 0 33 33"><path fill="currentColor" d="m13.237 21.04 11.082-8.19c.543-.4 1.32-.244 1.578.38 1.363 3.288.754 7.241-1.957 9.955-2.71 2.714-6.482 3.31-9.93 1.954l-3.765 1.745c5.401 3.697 11.96 2.782 16.059-1.324 3.251-3.255 4.258-7.692 3.317-11.693l.008.009c-1.365-5.878.336-8.227 3.82-13.031q.123-.17.247-.345l-4.585 4.59v-.014L13.234 21.044M10.95 23.031c-3.877-3.707-3.208-9.446.1-12.755 2.446-2.449 6.454-3.448 9.952-1.979L24.76 6.56c-.677-.49-1.545-1.017-2.54-1.387A12.465 12.465 0 0 0 8.675 7.901c-3.519 3.523-4.625 8.94-2.725 13.561 1.42 3.454-.907 5.898-3.251 8.364-.83.874-1.664 1.749-2.335 2.674l10.583-9.466"/></svg>"#
} else {
// X circle
r#"<svg width="48" height="48" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" style="color:#ef4444"><circle cx="12" cy="12" r="10"/><line x1="15" y1="9" x2="9" y2="15"/><line x1="9" y1="9" x2="15" y2="15"/></svg>"#
};
format!(
r#"<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8"/>
<meta name="viewport" content="width=device-width,initial-scale=1"/>
<meta name="color-scheme" content="light dark"/>
<title>{title}</title>
<style>
*{{margin:0;padding:0;box-sizing:border-box}}
body{{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,Helvetica,Arial,sans-serif;
display:flex;align-items:center;justify-content:center;min-height:100vh;
background:#0a0a0a;color:#e5e5e5}}
.card{{text-align:center;display:flex;flex-direction:column;align-items:center;gap:16px;padding:48px}}
h1{{font-size:18px;font-weight:600}}
p{{font-size:14px;color:#a3a3a3}}
@media(prefers-color-scheme:light){{
body{{background:#fafafa;color:#171717}}
p{{color:#525252}}
}}
</style>
</head>
<body>
<div class="card">
{icon}
<h1>{title}</h1>
<p>{message}</p>
</div>
</body>
</html>"#,
title = title,
icon = icon,
message = message,
)
}
/// Build the axum router for the OIDC loopback callback server.
fn build_callback_router(tx: tokio::sync::mpsc::Sender<CallbackResult>) -> Router {
let cors =
crate::auth::config::accounts_app_cors_layer(Method::GET).allow_private_network(true);
Router::new()
.route("/callback", get(handle_callback))
.layer(cors)
.with_state(tx)
}
async fn handle_callback(
State(tx): State<tokio::sync::mpsc::Sender<CallbackResult>>,
Query(params): Query<HashMap<String, String>>,
) -> (StatusCode, Html<String>) {
let result = parse_callback_params(&params);
let response = callback_response(&result);
if let Err(e) = tx.try_send(result) {
tracing::error!(?e, "OIDC: callback channel send failed; auth will time out");
}
response
}
fn parse_callback_params(params: &HashMap<String, String>) -> CallbackResult {
if let Some(code) = params.get("code") {
let state = params.get("state").cloned().unwrap_or_default();
tracing::debug!(state = %state, "OIDC: received code via loopback callback");
return Ok(Callback {
code: code.clone(),
state,
});
}
let error = params.get("error").cloned().unwrap_or_default();
let desc = params.get("error_description").cloned().unwrap_or_default();
tracing::error!(error = %error, desc = %desc, "OIDC: IdP returned error");
Err(if desc.is_empty() {
error
} else {
format!("{error}: {desc}")
})
}
fn callback_response(result: &CallbackResult) -> (StatusCode, Html<String>) {
let (title, message) = match result {
Ok(_) => (
"Signed in",
"You can close this window and return to Grok Build.",
),
Err(_) => ("Access denied", "Close this window and try again."),
};
(
StatusCode::OK,
Html(callback_page(title, message, result.is_ok())),
)
}
/// Wait until stdin has data or `tx` is closed. Returns `false` if closed.
#[cfg(unix)]
fn wait_for_stdin_or_closed(
stdin: &std::io::Stdin,
tx: &tokio::sync::mpsc::Sender<CallbackResult>,
) -> bool {
use std::os::unix::io::AsRawFd;
let fd = stdin.as_raw_fd();
loop {
if tx.is_closed() {
return false;
}
let ready = unsafe {
let mut fds = std::mem::zeroed::<libc::pollfd>();
fds.fd = fd;
fds.events = libc::POLLIN;
libc::poll(&mut fds, 1, 200)
};
if ready > 0 {
return true;
}
}
}
fn spawn_stdin_reader(tx: tokio::sync::mpsc::Sender<CallbackResult>) {
tokio::task::spawn_blocking(move || {
use std::io::BufRead;
let stdin = std::io::stdin();
let mut buf = String::new();
loop {
#[cfg(unix)]
if !wait_for_stdin_or_closed(&stdin, &tx) {
tracing::debug!("OIDC: stdin reader exiting, channel closed");
return;
}
#[cfg(not(unix))]
if tx.is_closed() {
tracing::debug!("OIDC: stdin reader exiting, channel closed");
return;
}
buf.clear();
let mut handle = stdin.lock();
match handle.read_line(&mut buf) {
Ok(0) => return,
Ok(_) => {}
Err(_) => return,
}
drop(handle);
let trimmed = buf.trim().to_owned();
if trimmed.is_empty() {
continue;
}
match parse_pasted_input(&trimmed) {
Ok(result) => {
tracing::debug!("OIDC: received code via stdin paste");
let _ = tx.blocking_send(Ok(result));
return;
}
Err(OidcError::InvalidPastedInput(msg)) => {
tracing::debug!(input = %msg, "OIDC: invalid stdin paste, retrying");
eprintln!(" Invalid input: {msg}. Try again:");
}
Err(e) => {
tracing::warn!(error = %e, "OIDC: stdin paste returned auth error");
let _ = tx.blocking_send(Err(e.to_string()));
return;
}
}
}
});
}
/// Race loopback callback against manual paste from `code_rx`.
async fn race_callback_and_client_ui(
listener: TcpListener,
code_rx: &mut tokio::sync::mpsc::Receiver<String>,
) -> anyhow::Result<Callback> {
tracing::debug!("OIDC: waiting for auth code (loopback + client paste)");
let (tx, mut rx) = tokio::sync::mpsc::channel::<CallbackResult>(1);
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
let app = build_callback_router(tx.clone());
let server = tokio::spawn(async move {
let _ = axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = shutdown_rx.await;
})
.await;
});
// Bridge client paste input into the callback channel.
let client_tx = tx.clone();
let client_bridge = async {
while let Some(code) = code_rx.recv().await {
match parse_pasted_input(&code) {
Ok(result) => {
tracing::debug!("OIDC: received code via client paste");
let _ = client_tx.send(Ok(result)).await;
return;
}
Err(e) => {
tracing::debug!(error = %e, "OIDC: invalid client paste input");
}
}
}
};
drop(tx);
let result = tokio::select! {
r = tokio::time::timeout(AUTH_CALLBACK_TIMEOUT, rx.recv()) => {
r.map_err(|_| anyhow::Error::new(OidcError::CallbackTimeout))?
.ok_or_else(|| anyhow::Error::new(OidcError::CallbackChannelClosed))?
}
_ = client_bridge => {
rx.recv().await
.ok_or_else(|| anyhow::Error::new(OidcError::CallbackChannelClosed))?
}
};
let _ = shutdown_tx.send(());
let _ = server.await;
result.map_err(|e| anyhow::Error::new(OidcError::CallbackAuthFailed(e)))
}
/// Race loopback callback against stdin paste.
async fn race_callback_and_stdin(
listener: TcpListener,
enable_stdin: bool,
) -> anyhow::Result<Callback> {
tracing::debug!(
enable_stdin = enable_stdin,
"OIDC: waiting for auth code (loopback + stdin)"
);
let (tx, mut rx) = tokio::sync::mpsc::channel::<CallbackResult>(1);
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
let app = build_callback_router(tx.clone());
let server = tokio::spawn(async move {
let _ = axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = shutdown_rx.await;
})
.await;
});
if enable_stdin {
spawn_stdin_reader(tx.clone());
}
drop(tx);
let result = tokio::time::timeout(AUTH_CALLBACK_TIMEOUT, rx.recv())
.await
.map_err(|_| {
// "10 minutes" must match AUTH_CALLBACK_TIMEOUT above
tracing::error!("auth: timed out after 10 minutes waiting for auth code");
anyhow::Error::new(OidcError::CallbackTimeout)
})?
.ok_or_else(|| {
tracing::error!(
"OIDC: callback channel closed, no code received from loopback or stdin"
);
anyhow::Error::new(OidcError::CallbackChannelClosed)
})?;
let _ = shutdown_tx.send(());
let _ = server.await;
result.map_err(|e| anyhow::Error::new(OidcError::CallbackAuthFailed(e)))
}
/// Run the full OIDC login flow: discovery → PKCE → browser → callback → token exchange → persist.
pub async fn run_login_flow(
config: &GrokComConfig,
auth_manager: &Arc<AuthManager>,
channels: Option<super::super::flow::AuthChannels>,
) -> anyhow::Result<(GrokAuth, bool)> {
let oidc = config
.oidc
.as_ref()
.ok_or_else(|| anyhow::Error::new(OidcError::NotConfigured))?;
run_login_flow_with_config(oidc, auth_manager, channels).await
}
/// Run the OIDC login flow with an explicit [`OidcAuthConfig`].
///
/// Also used by the OAuth2 provider path via [`OAuth2ProviderConfig::as_oidc`].
///
/// The flow races two input paths:
/// - **Path A**: A loopback HTTP server on `127.0.0.1` that receives the IdP redirect.
/// - **Path B**: Stdin paste — the user manually pastes the callback URL or bare auth code.
///
/// Path B is essential for remote VMs where the browser runs on a different machine
/// and the `127.0.0.1` redirect cannot reach the CLI process.
/// * `channels` — `Some`: pushes the auth URL to the TUI and receives pasted codes.
/// `None`: prints to stderr / reads stdin (CLI mode).
pub async fn run_login_flow_with_config(
oidc: &OidcAuthConfig,
auth_manager: &Arc<AuthManager>,
channels: Option<super::super::flow::AuthChannels>,
) -> anyhow::Result<(GrokAuth, bool)> {
tracing::info!(issuer = %oidc.issuer, client_id = %oidc.client_id, "OIDC: starting login flow");
// Ensure jsonwebtoken CryptoProvider is installed (required for JWT validation).
jsonwebtoken::crypto::CryptoProvider::install_default(
&jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER,
)
.ok();
let discovery = discover(&oidc.issuer).await?;
let pkce = generate_pkce();
let state = uuid::Uuid::now_v7().to_string();
let nonce = uuid::Uuid::now_v7().to_string();
// In local-dev mode, use a fixed callback port so the redirect_uri is stable
// and can be pre-registered with the local OAuth2 provider. In production the
// OS picks a random available port.
let callback_port: u16 = if super::super::config::use_local_auth() {
56121
} else {
0
};
let listener = TcpListener::bind(("127.0.0.1", callback_port))
.await
.map_err(|e| anyhow::Error::new(OidcError::BindLoopback(e.to_string())))?;
let port = listener.local_addr()?.port();
let redirect_uri = format!("http://127.0.0.1:{}/callback", port);
let oauth2 = auth_manager.grok_com_config().oauth2.as_ref();
let auth_url = build_authorize_url(
oidc,
oauth2,
&discovery,
&redirect_uri,
&pkce,
&state,
&nonce,
);
tracing::debug!(port = port, redirect_uri = %redirect_uri, "OIDC: callback server bound");
let (url_tx, code_rx) = match channels {
Some(ch) => (ch.url_tx, Some(ch.code_rx)),
None => (None, None),
};
let has_client_ui = code_rx.is_some();
if has_client_ui {
// Client provides its own auth UI; just open the browser.
if let Err(e) = webbrowser::open(&auth_url) {
tracing::debug!(error = %e, "OIDC: failed to open browser");
}
} else {
// No client UI — print to stderr.
eprintln!();
let provider_label = if oidc.issuer == super::super::config::XAI_OAUTH2_ISSUER {
"Grok".to_owned()
} else {
oidc.issuer.clone()
};
eprintln!("Signing in with {}...", provider_label);
eprintln!();
if let Err(e) = webbrowser::open(&auth_url) {
tracing::debug!(error = %e, "OIDC: failed to open browser");
}
eprintln!("Open this URL to sign in:");
eprintln!(" {}", auth_url);
}
let use_stdin = !has_client_ui && std::io::stdin().is_terminal();
if use_stdin {
eprintln!();
eprintln!("Paste the URL here if it doesn't connect:");
}
// Push auth URL to the TUI via oneshot.
if let Some(tx) = url_tx {
let _ = tx.send(super::super::flow::AuthUrlInfo {
url: auth_url.clone(),
mode: super::super::flow::AuthUrlMode::Loopback,
});
}
let Callback {
code,
state: received_state,
} = if let Some(mut rx) = code_rx {
// Client UI: race loopback against manual paste via code_rx.
race_callback_and_client_ui(listener, &mut rx).await?
} else {
// No client UI: race loopback against stdin paste.
race_callback_and_stdin(listener, use_stdin).await?
};
// Validate state (skip for bare code paste where state is empty)
if !received_state.is_empty() {
validate_state(&state, &received_state)?;
}
let tokens = exchange_code(
&discovery.token_endpoint,
&code,
&redirect_uri,
&oidc.client_id,
&pkce.code_verifier,
)
.await?;
tracing::info!(
has_refresh = tokens.refresh_token.is_some(),
expires_in = ?tokens.expires_in,
"OIDC: token exchange complete"
);
// Resolve the actual principal chosen on the consent screen.
//
// The shell's config may not have principal_type set (personal login),
// but the user might pick "Team" on the consent screen. The server
// encodes the chosen principal in the access token JWT. If the config
// doesn't specify a principal, peek at the token to discover it.
let token_principal = peek_access_token_principal(&tokens.access_token);
// The authorize URL only pre-selects; verify the token's principal here.
// Match the principal id even if `principal_type` is absent.
let principal_policy = login_principal_policy(auth_manager.grok_com_config());
enforce_login_principal(
principal_policy.as_ref(),
peek_access_token_principal_id(&tokens.access_token).as_deref(),
)?;
let (resolved_principal_type, resolved_principal_id, resolved_team_id) = {
let cfg_pt = oauth2.and_then(|cfg| cfg.principal_type.clone());
let cfg_pid = oauth2.and_then(|cfg| cfg.principal_id.clone());
if cfg_pt.is_some() {
(cfg_pt, cfg_pid, None)
} else if let Some((pt, pid, tid)) = token_principal {
tracing::info!(
principal_type = %pt,
principal_id = %pid,
team_id = ?tid,
"OIDC: resolved principal from access token"
);
(Some(pt), Some(pid), tid)
} else {
(cfg_pt, cfg_pid, None)
}
};
let user_info = extract_user_info(
tokens.id_token.as_deref(),
&discovery,
&oidc.issuer,
&oidc.client_id,
&nonce,
resolved_principal_type.as_deref(),
resolved_principal_id.as_deref(),
resolved_team_id,
)
.await?;
tracing::debug!(user_id = %user_info.user_id, "OIDC: extracted user info");
let mut auth = build_grok_auth(tokens, user_info, &oidc.issuer, &oidc.client_id);
auth_manager.enrich_auth_inline(&mut auth).await;
let auth = auth_manager
.update(auth)
.await
.map_err(|e| anyhow::Error::new(OidcError::SaveAuth(e.to_string())))?;
tracing::info!(user_id = %auth.user_id, "OIDC: login complete, credentials saved");
Ok((auth, true))
}
/// Successful OIDC callback payload.
#[derive(Debug, PartialEq, Eq)]
struct Callback {
code: String,
state: String,
}
/// Result from the OIDC callback: either a [`Callback`] or an IdP error message.
type CallbackResult = Result<Callback, String>;
#[cfg(test)]
mod tests {
use super::super::test_helpers::*;
use super::*;
/// End-to-end test: mock IdP + full login flow with code arriving via loopback.
/// Exercises discovery → PKCE → race_callback_and_stdin → token exchange → user info → persist.
#[tokio::test]
async fn full_login_flow_via_race() {
ensure_crypto_provider();
let (issuer, idp_server) = start_mock_idp().await;
let temp_dir = tempfile::tempdir().unwrap();
// Dead proxy port: inline `/user` enrichment fails fast in tests.
let dead_proxy = {
let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
format!("http://127.0.0.1:{}", l.local_addr().unwrap().port())
};
let auth_manager = Arc::new(
AuthManager::new(temp_dir.path(), GrokComConfig::default())
.with_proxy_base_url(&dead_proxy),
);
let oidc_cfg = OidcAuthConfig {
issuer: issuer.clone(),
client_id: TEST_CLIENT_ID.into(),
scopes: vec!["openid".into(), "email".into()],
audience: None,
};
let discovery = discover(&oidc_cfg.issuer).await.unwrap();
let pkce = generate_pkce();
let state = "test-state".to_string();
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let redirect_uri = format!("http://127.0.0.1:{port}/callback");
let _auth_url = build_authorize_url(
&oidc_cfg,
None,
&discovery,
&redirect_uri,
&pkce,
&state,
TEST_NONCE,
);
// Simulate browser callback via race_callback_and_stdin
let Callback {
code,
state: received_state,
} = tokio::join!(race_callback_and_stdin(listener, false), async {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
reqwest::get(format!(
"http://127.0.0.1:{port}/callback?code=mock-auth-code&state={state}"
))
.await
.unwrap();
})
.0
.unwrap();
assert_eq!(code, "mock-auth-code");
assert_eq!(received_state, state);
let tokens = exchange_code(
&discovery.token_endpoint,
&code,
&redirect_uri,
&oidc_cfg.client_id,
&pkce.code_verifier,
)
.await
.unwrap();
assert_eq!(tokens.access_token, "mock-access-token");
let user_info = extract_user_info(
tokens.id_token.as_deref(),
&discovery,
&oidc_cfg.issuer,
&oidc_cfg.client_id,
TEST_NONCE,
None,
None,
None,
)
.await
.unwrap();
let auth = build_grok_auth(tokens, user_info, &oidc_cfg.issuer, &oidc_cfg.client_id);
let auth = auth_manager.update(auth).await.unwrap();
assert_eq!(auth.key, "mock-access-token");
assert_eq!(auth.refresh_token.as_deref(), Some("mock-refresh-token"));
assert_eq!(auth.user_id, "user-42");
assert_eq!(auth.email.as_deref(), Some("test@corp.com"));
assert!(auth.principal_type.is_none());
assert!(auth.principal_id.is_none());
assert!(auth.expires_at.is_some());
assert_eq!(auth.oidc_issuer.as_deref(), Some(issuer.as_str()));
let auth_json = std::fs::read_to_string(temp_dir.path().join("auth.json")).unwrap();
assert!(auth_json.contains("mock-access-token"));
assert!(auth_json.contains("user-42"));
idp_server.abort();
}
/// Parser matrix: full callback URL, bare code, error URL, empty.
/// Each case is one bug class:
/// - full URL: regression in URL extraction
/// - bare code: paste-friendly fallback
/// - error URL: surfaces IdP error to user
/// - empty: input validation
#[test]
fn parse_pasted_input_matrix() {
// (input, expected: Ok((code, state)) | Err substring)
let ok_cases: &[(&str, &str, &str)] = &[
(
"http://127.0.0.1:54321/callback?code=abc123&state=xyz789",
"abc123",
"xyz789",
),
("abc123def456", "abc123def456", ""),
];
for (input, code, state) in ok_cases {
let cb =
parse_pasted_input(input).unwrap_or_else(|e| panic!("parse {input:?} failed: {e}"));
assert_eq!(cb.code, *code, "code for {input:?}");
assert_eq!(cb.state, *state, "state for {input:?}");
}
let err_cases: &[(&str, &str)] = &[
(
"http://127.0.0.1:54321/callback?error=access_denied&error_description=User+denied",
"access_denied",
),
("", ""),
(" ", ""),
];
for (input, expected_substr) in err_cases {
let err = parse_pasted_input(input).unwrap_err();
if !expected_substr.is_empty() {
assert!(
err.to_string().contains(expected_substr),
"input {input:?} -> unexpected err: {err}"
);
}
}
}
}
crates/codegen/kigi-shell/src/auth/oidc/mod.rs
-14
View File
@@ -1,14 +0,0 @@
//! OIDC authentication: protocol, login, and refresh submodules.
mod login;
pub(crate) mod protocol;
pub(crate) mod refresh;
#[cfg(test)]
mod test_helpers;
pub use login::{run_login_flow, run_login_flow_with_config};
pub(crate) use protocol::{
enforce_login_principal, is_configured, login_principal_policy, peek_access_token_principal,
peek_access_token_principal_id, with_alpha_test_key,
};
pub(crate) use refresh::{OidcRefreshResult, oidc_token_exchange};
crates/codegen/kigi-shell/src/auth/oidc/protocol.rs
-1283
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/oidc/refresh.rs
-249
View File
@@ -1,249 +0,0 @@
//! Pure-data OIDC refresh. Talks to the IdP and returns
//! [`OidcRefreshResult`] without touching [`AuthManager`].
use super::super::GrokAuth;
use super::protocol::{OidcError, OidcUserInfo, build_grok_auth, discover, refresh_tokens};
use crate::auth::error::RefreshTokenFailedReason;
/// Outcome of a pure OIDC token refresh (no AuthManager mutations).
pub(crate) enum OidcRefreshResult {
/// Fresh token obtained. Caller must persist.
Success(Box<GrokAuth>),
/// Terminal error from the IdP, already classified into a reason.
TerminalError { reason: RefreshTokenFailedReason },
/// Non-terminal failure (discovery failed, network error, etc.)
Failed,
}
/// Classify an OAuth2 `error` code as a terminal refresh failure. `None` means
/// non-terminal (retryable). Single source of truth for which codes are fatal;
/// the retry gate (`protocol::is_transient_refresh_error`) defers to this too.
pub(super) fn classify_terminal(error_code: &str) -> Option<RefreshTokenFailedReason> {
match error_code {
"invalid_grant" => Some(RefreshTokenFailedReason::RefreshTokenRejected),
"invalid_client" => Some(RefreshTokenFailedReason::ClientRejected),
_ => None,
}
}
/// `oauth2-provider` refresh-token rotation-grace window (ms). Only a clock
/// divergence past this bound is flagged as a suspected suspend-straddle, since
/// a longer suspend can turn a lost refresh response into a revoked RT.
const ROTATION_GRACE_MS: u64 = 60_000;
/// Exchange a refresh_token for fresh tokens at the IdP. Pure data return, no
/// `AuthManager` mutations; the caller (`OidcRefresher`) routes the result
/// through `refresh_chain`.
pub(crate) async fn oidc_token_exchange(auth: &GrokAuth) -> OidcRefreshResult {
let has_rt = auth.refresh_token.is_some();
let has_issuer = auth.oidc_issuer.is_some();
let has_client_id = auth.oidc_client_id.is_some();
tracing::debug!(
has_rt,
has_issuer,
has_client_id,
"oidc try_refresh_pure enter"
);
if !has_rt || !has_issuer || !has_client_id {
kigi_log::unified_log::warn(
"oidc try_refresh skipped: missing fields",
None,
Some(serde_json::json!({
"has_refresh_token": has_rt,
"has_issuer": has_issuer,
"has_client_id": has_client_id,
"auth_mode": format!("{:?}", auth.auth_mode),
})),
);
}
let Some(refresh_tok) = auth.refresh_token.as_ref() else {
return OidcRefreshResult::Failed;
};
let Some(issuer) = auth.oidc_issuer.as_ref() else {
return OidcRefreshResult::Failed;
};
let Some(client_id) = auth.oidc_client_id.as_ref() else {
return OidcRefreshResult::Failed;
};
crate::unified_log::info(
"oidc try_refresh_pure enter",
None,
Some(serde_json::json!({ "issuer": issuer, "client_id": client_id })),
);
// Suspend probe: the monotonic clock pauses while the machine is asleep
// but the wall clock does not, so a large divergence around the IdP call
// means the process was suspended mid-refresh — the exact condition that
// can revoke the refresh token (response lost across sleep).
let started_mono = std::time::Instant::now();
let started_wall = chrono::Utc::now();
let timing = || {
let mono_ms = started_mono.elapsed().as_millis() as u64;
let wall_ms = (chrono::Utc::now() - started_wall)
.num_milliseconds()
.max(0) as u64;
let suspended_ms = wall_ms.saturating_sub(mono_ms);
(
mono_ms,
wall_ms,
suspended_ms,
suspended_ms > ROTATION_GRACE_MS,
)
};
let discovery = match discover(issuer).await {
Ok(d) => d,
Err(e) => {
let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing();
crate::unified_log::error(
"oidc try_refresh_pure discovery failed",
None,
Some(serde_json::json!({
"error": format!("{e:#}"),
"mono_ms": mono_ms,
"wall_ms": wall_ms,
"suspended_ms": suspended_ms,
"suspected_suspend": suspected_suspend,
})),
);
if suspected_suspend {
emit_suspend_spanned("discovery_failed", suspended_ms);
}
return OidcRefreshResult::Failed;
}
};
let tokens = match refresh_tokens(
&discovery.token_endpoint,
refresh_tok,
client_id,
auth.principal_type.as_deref(),
auth.principal_id.as_deref(),
)
.await
{
Ok(t) => t,
Err(e) => {
if let Some(OidcError::TokenRefreshHttp { body, .. }) = e.downcast_ref::<OidcError>()
&& let Some(error_code) = serde_json::from_str::<serde_json::Value>(body)
.ok()
.and_then(|v| v.get("error")?.as_str().map(str::to_owned))
&& let Some(reason) = classify_terminal(&error_code)
{
let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing();
let cred_age_secs = auth.mint_age_seconds();
crate::unified_log::error(
"oidc try_refresh_pure terminal error",
None,
Some(serde_json::json!({
"error_code": error_code,
"client_id": client_id,
"tried_rt_prefix": auth.refresh_token.as_deref().map(crate::auth::token_suffix),
"error_description": serde_json::from_str::<serde_json::Value>(body)
.ok()
.and_then(|v| v.get("error_description").cloned()),
"mono_ms": mono_ms,
"wall_ms": wall_ms,
"suspended_ms": suspended_ms,
"suspected_suspend": suspected_suspend,
"cred_age_secs": cred_age_secs,
})),
);
if suspected_suspend {
emit_suspend_spanned(&error_code, suspended_ms);
}
return OidcRefreshResult::TerminalError { reason };
}
let http_status = e.downcast_ref::<OidcError>().and_then(|oe| match oe {
OidcError::TokenRefreshHttp { status, .. } => Some(*status),
_ => None,
});
let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing();
crate::unified_log::error(
"oidc try_refresh_pure token exchange failed",
None,
Some(serde_json::json!({
"error": e.to_string(),
"client_id": client_id,
"http_status": http_status,
"mono_ms": mono_ms,
"wall_ms": wall_ms,
"suspended_ms": suspended_ms,
"suspected_suspend": suspected_suspend,
})),
);
tracing::warn!(
error = %e,
http_status = ?http_status,
client_id = %client_id,
issuer = %issuer,
"OIDC: token refresh failed"
);
if suspected_suspend {
emit_suspend_spanned("transient_failed", suspended_ms);
}
return OidcRefreshResult::Failed;
}
};
// Reuse identity from original login; new id_token from refresh is intentionally skipped.
let user_info = OidcUserInfo {
user_id: auth.user_id.clone(),
email: auth.email.clone(),
first_name: auth.first_name.clone(),
last_name: auth.last_name.clone(),
profile_image_asset_id: auth.profile_image_asset_id.clone(),
principal_type: auth.principal_type.clone(),
principal_id: auth.principal_id.clone(),
team_id: auth.team_id.clone(),
team_name: auth.team_name.clone(),
team_role: auth.team_role.clone(),
organization_id: auth.organization_id.clone(),
organization_name: auth.organization_name.clone(),
organization_role: auth.organization_role.clone(),
user_blocked_reason: auth.user_blocked_reason.clone(),
team_blocked_reasons: auth.team_blocked_reasons.clone(),
coding_data_retention_opt_out: auth.coding_data_retention_opt_out,
};
let mut new_auth = build_grok_auth(tokens, user_info, issuer, client_id);
let idp_rotated = new_auth.refresh_token.is_some();
// Keep old refresh token if IdP didn't rotate it
if new_auth.refresh_token.is_none() {
new_auth.refresh_token = auth.refresh_token.clone();
}
tracing::debug!(
idp_rotated,
key_prefix = crate::auth::token_suffix(&new_auth.key),
"oidc try_refresh_pure token obtained"
);
let (mono_ms, wall_ms, suspended_ms, suspected_suspend) = timing();
crate::unified_log::info(
"oidc try_refresh_pure succeeded",
None,
Some(serde_json::json!({
"expires_at": new_auth.expires_at.map(|e| e.to_rfc3339()),
"mono_ms": mono_ms,
"wall_ms": wall_ms,
"suspended_ms": suspended_ms,
"suspected_suspend": suspected_suspend,
})),
);
if suspected_suspend {
emit_suspend_spanned("ok", suspended_ms);
}
OidcRefreshResult::Success(Box::new(new_auth))
}
/// Alertable event: an OIDC refresh's network call spanned a suspend (wall
/// clock ran far ahead of the monotonic clock) — the precondition for a
/// lost-response refresh-token revocation.
fn emit_suspend_spanned(outcome: &str, suspended_ms: u64) {
crate::unified_log::warn(
"auth.refresh.suspend_spanned",
None,
Some(serde_json::json!({
"outcome": outcome,
"suspended_ms": suspended_ms,
})),
);
}
crates/codegen/kigi-shell/src/auth/oidc/test_helpers.rs
-130
View File
@@ -1,130 +0,0 @@
//! Shared test helpers for `oidc::protocol::tests` and `oidc::login::tests`.
//! Both test modules need a mock IdP server (`start_mock_idp`), JWT
//! signing primitives (`generate_test_rsa_key`, `mock_idp_token`), and
//! the same constants. Extracted here so neither test mod has to
//! re-implement them.
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use super::protocol::{Discovery, discover};
pub(super) const TEST_KID: &str = "test-kid";
pub(super) const TEST_NONCE: &str = "test-nonce-value";
pub(super) const TEST_CLIENT_ID: &str = "test-client-id";
pub(super) fn ensure_crypto_provider() {
let _ = rustls::crypto::ring::default_provider().install_default();
let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default();
}
pub(super) fn generate_test_rsa_key() -> (String, String, String) {
use rsa::pkcs8::EncodePrivateKey;
use rsa::traits::PublicKeyParts;
let private_key = rsa::RsaPrivateKey::new(&mut rsa::rand_core::OsRng, 2048).unwrap();
let pem = private_key
.to_pkcs8_pem(rsa::pkcs8::LineEnding::LF)
.unwrap()
.to_string();
let jwk_n = URL_SAFE_NO_PAD.encode(private_key.n().to_bytes_be());
let jwk_e = URL_SAFE_NO_PAD.encode(private_key.e().to_bytes_be());
(pem, jwk_n, jwk_e)
}
pub(super) async fn mock_idp_token() -> (String, String, Discovery, tokio::task::JoinHandle<()>) {
let (issuer, handle) = start_mock_idp().await;
let discovery = discover(&issuer).await.unwrap();
let resp: serde_json::Value = crate::http::shared_client()
.post(&discovery.token_endpoint)
.form(&[("grant_type", "authorization_code")])
.send()
.await
.unwrap()
.json()
.await
.unwrap();
let id_token = resp["id_token"]
.as_str()
.expect("mock missing id_token")
.to_string();
(issuer, id_token, discovery, handle)
}
pub(super) async fn start_mock_idp() -> (String, tokio::task::JoinHandle<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let issuer = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port());
let issuer_for_discovery = issuer.clone();
let (rsa_pem, jwk_n, jwk_e) = generate_test_rsa_key();
#[derive(serde::Serialize)]
struct Claims {
sub: &'static str,
email: &'static str,
iss: String,
aud: &'static str,
nonce: &'static str,
exp: usize,
}
let id_token = {
let mut hdr = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256);
hdr.kid = Some(TEST_KID.to_owned());
jsonwebtoken::encode(
&hdr,
&Claims {
sub: "user-42",
email: "test@corp.com",
iss: issuer.clone(),
aud: TEST_CLIENT_ID,
nonce: TEST_NONCE,
exp: (chrono::Utc::now() + chrono::Duration::hours(1)).timestamp() as usize,
},
&jsonwebtoken::EncodingKey::from_rsa_pem(rsa_pem.as_bytes()).unwrap(),
)
.unwrap()
};
let app = axum::Router::new()
.route(
"/.well-known/openid-configuration",
axum::routing::get(move || {
let iss = issuer_for_discovery.clone();
async move {
axum::Json(serde_json::json!({
"authorization_endpoint": format!("{iss}/authorize"),
"token_endpoint": format!("{iss}/token"),
"jwks_uri": format!("{iss}/jwks"),
"id_token_signing_alg_values_supported": ["RS256"],
}))
}
}),
)
.route(
"/jwks",
axum::routing::get(move || {
let n = jwk_n.clone();
let e = jwk_e.clone();
async move {
axum::Json(serde_json::json!({
"keys": [{
"kty": "RSA", "alg": "RS256", "kid": TEST_KID,
"n": n, "e": e,
}]
}))
}
}),
)
.route(
"/token",
axum::routing::post(move || {
let tok = id_token.clone();
async move {
axum::Json(serde_json::json!({
"access_token": "mock-access-token",
"refresh_token": "mock-refresh-token",
"id_token": tok,
"expires_in": 3600,
}))
}
}),
);
let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(issuer, handle)
}
crates/codegen/kigi-shell/src/auth/recovery.rs
+123 -311
View File
@@ -3,71 +3,54 @@
//! When the server rejects a token, `UnauthorizedRecovery` walks through
//! a sequence of recovery steps before giving up:
//!
//! 1. **ReloadFromDisk** — re-read `auth.json` under a file lock; if the
//! on-disk token differs from the rejected one, accept it (another
//! process may have refreshed).
//! 2. **RefreshFromAuthority** — run the appropriate refresh chain
//! (OIDC token refresh, external binary, etc.) based on `TokenType`,
//! unless the live token was minted moments ago (fresh-mint guard).
//! 3. **DevboxRecovery** — on devboxes, purge `auth.json` and mint fresh
//! OIDC credentials.
//! 4. **Done** — all recovery strategies exhausted.
//! 1. **ReloadFromDisk** — re-read the persisted credential under a file
//! lock; if it differs from the rejected one, accept it (another process
//! may have refreshed).
//! 2. **RefreshFromAuthority** — run the refresh chain against the Kimi
//! OAuth host, unless the live token was minted moments ago (fresh-mint
//! guard).
//! 3. **Done** — all recovery strategies exhausted.
use std::sync::Arc;
use crate::auth::error::{AuthError, RefreshTokenError, RefreshTokenFailedReason};
use crate::auth::manager::AuthManager;
use crate::auth::model::GrokAuth;
use crate::auth::model::KimiAuth;
use crate::auth::token_type::TokenType;
/// Whether a terminal `AuthError` forces a manual re-login (`None` cases
/// Whether a terminal `AuthError` forces a manual re-login (`false` cases
/// self-heal or are transient). Lives here (not on `AuthError`) so the error
/// model stays free of recovery policy.
pub(crate) fn forces_manual_reauth(err: &AuthError) -> bool {
match err {
AuthError::Refresh(RefreshTokenError::Permanent(e)) => match e.reason {
RefreshTokenFailedReason::RefreshTokenRejected => true,
// Self-healing via the TTL, not a manual re-auth.
RefreshTokenFailedReason::ClientRejected | RefreshTokenFailedReason::Other => false,
// Self-healing via the tombstone cooldown, not a manual re-auth.
RefreshTokenFailedReason::Other => false,
},
AuthError::ServerRejectedNoRecovery
| AuthError::RecoveryExhausted
| AuthError::TokenExpiredNoRefresh
| AuthError::PinnedTeamMismatch { .. } => true,
// API-key lockouts are out of scope: an admin disabling API-key auth
// means rotate the key, not `/login`.
AuthError::ApiKeyAuthDisabled
| AuthError::Refresh(RefreshTokenError::Transient(_))
| AuthError::NotLoggedIn => false,
| AuthError::TokenExpiredNoRefresh => true,
AuthError::Refresh(RefreshTokenError::Transient(_)) | AuthError::NotLoggedIn => false,
}
}
/// Whether the relay should stop reconnecting on this recovery error. Its own
/// predicate rather than reusing `forces_manual_reauth`: the relay must give up
/// on any terminal auth failure, including `ApiKeyAuthDisabled` (a kill-switched
/// API key), which deliberately doesn't force a manual re-login.
pub(crate) fn relay_should_cancel(err: &AuthError) -> bool {
forces_manual_reauth(err) || matches!(err, AuthError::ApiKeyAuthDisabled)
}
/// Fresh-mint guard window (±) for `ServerRejected` refreshes
/// ([`UnauthorizedRecovery::fresh_mint_guard`]). 120s outlasts in-flight
/// requests sent with a previous key plus validation lag (observed stale
/// 401s land ~20s after mint), while `current()`'s 300s early-invalidation
/// buffer keeps any guard-returned token wire-valid. A genuinely-dead fresh
/// token waits at most this long to re-mint; the symmetric bound caps that
/// delay when the clock stepped back.
/// 401s land ~20s after mint), while the refresh-threshold buffer keeps any
/// guard-returned token wire-valid. A genuinely-dead fresh token waits at
/// most this long to re-mint; the symmetric bound caps that delay when the
/// clock stepped back.
const FRESH_MINT_GUARD_SECS: i64 = 120;
/// Which recovery step to attempt next.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RecoveryStep {
/// Re-read auth.json from disk (file-locked).
/// Re-read the persisted credential (file-locked).
ReloadFromDisk,
/// Refresh via the authority (OIDC, external binary, etc.).
/// Refresh via the Kimi OAuth host.
RefreshFromAuthority,
/// On devboxes: purge auth.json and mint fresh OIDC credentials.
DevboxRecovery,
/// All strategies exhausted.
Done,
}
@@ -79,8 +62,7 @@ pub struct UnauthorizedRecovery {
rejected_token: String,
/// Current step in the recovery sequence.
step: RecoveryStep,
/// Error from `RefreshFromAuthority`, propagated as fallback when
/// devbox recovery doesn't apply.
/// Error from `RefreshFromAuthority`, propagated on exhaustion.
authority_error: Option<AuthError>,
/// Whether the last authority failure was transient. Kept past the
/// `authority_error` handoff so exhaustion preserves the
@@ -90,7 +72,7 @@ pub struct UnauthorizedRecovery {
impl UnauthorizedRecovery {
/// `rejected` is the credential the server rejected: its key drives recovery.
pub(crate) fn new(auth_manager: Arc<AuthManager>, rejected: Option<GrokAuth>) -> Self {
pub(crate) fn new(auth_manager: Arc<AuthManager>, rejected: Option<KimiAuth>) -> Self {
let rejected_token = rejected.as_ref().map(|a| a.key.clone()).unwrap_or_default();
Self {
auth_manager,
@@ -102,14 +84,14 @@ impl UnauthorizedRecovery {
}
/// Attempt the next recovery step. Walks
/// `ReloadFromDisk -> RefreshFromAuthority -> DevboxRecovery -> Done`.
/// `ReloadFromDisk -> RefreshFromAuthority -> Done`.
/// `token_type` span field is recorded lazily via
/// `Span::is_disabled()` to avoid the lock when tracing is off.
#[tracing::instrument(
skip(self),
fields(step = ?self.step, token_type = tracing::field::Empty),
)]
pub async fn next(&mut self) -> Result<GrokAuth, AuthError> {
pub async fn next(&mut self) -> Result<KimiAuth, AuthError> {
let span = tracing::Span::current();
if !span.is_disabled() {
// Only acquire the inner-lock when tracing actually
@@ -122,23 +104,10 @@ impl UnauthorizedRecovery {
tracing::field::debug(self.auth_manager.token_type()),
);
}
self.resolve_next().await
self.next_step_loop().await
}
/// Walk the recovery steps and apply the team-pin policy gate.
async fn resolve_next(&mut self) -> Result<GrokAuth, AuthError> {
// Team-pin gate: 401 recovery must not resurrect a wrong-team session
// (disk adoption / refresh / devbox mint) for the relay to reconnect
// with. Clear + reject on mismatch.
let auth = self.next_step_loop().await?;
if let Some(e) = self.auth_manager.cached_token_policy_error(&auth) {
self.auth_manager.reject_and_clear(&e);
return Err(e);
}
Ok(auth)
}
async fn next_step_loop(&mut self) -> Result<GrokAuth, AuthError> {
async fn next_step_loop(&mut self) -> Result<KimiAuth, AuthError> {
loop {
match self.step {
RecoveryStep::ReloadFromDisk => {
@@ -148,30 +117,20 @@ impl UnauthorizedRecovery {
}
}
RecoveryStep::RefreshFromAuthority => {
self.step = RecoveryStep::DevboxRecovery;
self.step = RecoveryStep::Done;
match self.try_refresh_from_authority().await {
Ok(auth) => return Ok(auth),
Err(e) => {
self.authority_was_transient =
matches!(e, AuthError::Refresh(RefreshTokenError::Transient(_)));
self.authority_error = Some(e);
return Err(self
.authority_error
.take()
.unwrap_or(AuthError::RecoveryExhausted));
}
}
}
RecoveryStep::DevboxRecovery => {
self.step = RecoveryStep::Done;
// preferred_method=api_key forbids automatic OIDC mint.
if !self.auth_manager.grok_com_config().blocks_automatic_oidc()
&& self.auth_manager.is_devbox_environment()
&& let Ok(auth) = self.auth_manager.try_devbox_recovery().await
{
return Ok(auth);
}
return Err(self
.authority_error
.take()
.unwrap_or(AuthError::RecoveryExhausted));
}
RecoveryStep::Done => {
// Exhaustion after a *transient* authority failure stays
// transient: `RecoveryExhausted` here would count a network
@@ -187,9 +146,9 @@ impl UnauthorizedRecovery {
}
}
/// Re-read `auth.json` from disk. Accept the token only if it differs
/// Re-read the persisted credential. Accept the token only if it differs
/// from the one that was rejected.
async fn try_reload_from_disk(&self) -> Option<GrokAuth> {
async fn try_reload_from_disk(&self) -> Option<KimiAuth> {
let _lock = self
.auth_manager
.try_lock_auth_file_async(crate::auth::manager::AUTH_LOCK_TIMEOUT)
@@ -203,13 +162,13 @@ impl UnauthorizedRecovery {
// same-as-rejected / no entry): a silent arm hides which path
// a recovery loop is taking. Debug level — the disk-state
// *transition* is logged once by `read_disk_auth` itself.
kigi_log::unified_log::debug("auth recovery: no disk entry", None, None);
kigi_log::unified_log::debug("auth recovery: no persisted entry", None, None);
return None;
};
if crate::auth::is_expired(&disk_auth) {
tracing::debug!("auth recovery: disk token is expired, skipping");
tracing::debug!("auth recovery: persisted token is expired, skipping");
kigi_log::unified_log::debug(
"auth recovery: disk token expired",
"auth recovery: persisted token expired",
None,
Some(serde_json::json!({
"disk_key_prefix": crate::auth::token_suffix(&disk_auth.key),
@@ -219,9 +178,9 @@ impl UnauthorizedRecovery {
return None;
}
if self.is_different_token(&disk_auth) {
tracing::info!("auth recovery: disk has a different token, accepting");
tracing::info!("auth recovery: persisted store has a different token, accepting");
kigi_log::unified_log::info(
"auth recovery: adopted disk token",
"auth recovery: adopted persisted token",
None,
Some(serde_json::json!({
"adopted_key_prefix": crate::auth::token_suffix(&disk_auth.key),
@@ -231,8 +190,12 @@ impl UnauthorizedRecovery {
self.auth_manager.hot_swap(disk_auth.clone());
Some(disk_auth)
} else {
tracing::debug!("auth recovery: disk token is same as rejected, skipping");
kigi_log::unified_log::debug("auth recovery: disk token same as rejected", None, None);
tracing::debug!("auth recovery: persisted token is same as rejected, skipping");
kigi_log::unified_log::debug(
"auth recovery: persisted token same as rejected",
None,
None,
);
None
}
}
@@ -242,14 +205,12 @@ impl UnauthorizedRecovery {
/// clock that stepped far back) falls through to a normal refresh.
///
/// A 401 moments after a successful mint is a stale rejection (sent with
/// the previous key and mis-attributed — see `is_stale_snapshot`) or
/// validation lag on the new key — re-minting fixes neither, and a crash
/// between the IdP grant and persisting the response orphans the
/// replacement RT (forced re-login). Consumers retry with the returned
/// token; a genuinely-bad one refreshes once the window passes. Lives
/// here, not in `refresh_chain`, so paywall claims re-mints that call
/// `refresh_chain(ServerRejected)` directly are unaffected.
fn fresh_mint_guard(&self) -> Option<GrokAuth> {
/// the previous key) or validation lag on the new key — re-minting fixes
/// neither, and a crash between the token grant and persisting the
/// response orphans the replacement RT (forced re-login). Consumers retry
/// with the returned token; a genuinely-bad one refreshes once the window
/// passes.
fn fresh_mint_guard(&self) -> Option<KimiAuth> {
let auth = self.auth_manager.current()?;
let mint_age_seconds = auth.mint_age_seconds();
if !(-FRESH_MINT_GUARD_SECS..FRESH_MINT_GUARD_SECS).contains(&mint_age_seconds) {
@@ -272,14 +233,14 @@ impl UnauthorizedRecovery {
Some(auth)
}
/// Dispatch to the correct refresh chain based on the current `TokenType`.
/// Dispatch to the refresh chain based on the current `TokenType`.
///
/// Per-variant outcome:
///
/// - **OidcSession / ExternalBinary**: full refresh chain via the
/// authority, unless the live token is inside the fresh-mint guard
/// window ([`Self::fresh_mint_guard`]).
/// - **LegacySession / ApiKey**: no refresh authority for these
/// - **OAuthSession**: full refresh chain via the OAuth host, unless the
/// live token is inside the fresh-mint guard window
/// ([`Self::fresh_mint_guard`]).
/// - **SessionNoRefresh / ApiKey**: no refresh authority for these
/// types. We've already tried `ReloadFromDisk` (the previous
/// recovery step), so the server's 401 stands. Surface
/// [`AuthError::ServerRejectedNoRecovery`] -- *not*
@@ -289,10 +250,10 @@ impl UnauthorizedRecovery {
/// reading the variant can distinguish "ran past local TTL" from
/// "server actively rejected".
/// - **None**: no credentials at all.
async fn try_refresh_from_authority(&self) -> Result<GrokAuth, AuthError> {
async fn try_refresh_from_authority(&self) -> Result<KimiAuth, AuthError> {
let tt = self.auth_manager.token_type();
match tt {
TokenType::OidcSession | TokenType::ExternalBinary => {
TokenType::OAuthSession => {
if let Some(auth) = self.fresh_mint_guard() {
return Ok(auth);
}
@@ -325,7 +286,7 @@ impl UnauthorizedRecovery {
}
result
}
TokenType::LegacySession | TokenType::ApiKey => {
TokenType::SessionNoRefresh | TokenType::ApiKey => {
kigi_log::unified_log::warn(
"auth recovery: no refresh authority for token type",
None,
@@ -338,7 +299,7 @@ impl UnauthorizedRecovery {
}
/// Check if a candidate token is different from the rejected one.
fn is_different_token(&self, candidate: &GrokAuth) -> bool {
fn is_different_token(&self, candidate: &KimiAuth) -> bool {
candidate.key != self.rejected_token
}
}
@@ -348,29 +309,28 @@ mod tests {
//! State-machine matrix tests for `UnauthorizedRecovery`.
//!
//! Coverage targets:
//! - All 5 `TokenType` variants x dispatch in `try_refresh_from_authority`.
//! - All 4 `TokenType` variants x dispatch in `try_refresh_from_authority`.
//! - `try_reload_from_disk`: same/different/no token on disk.
//! - `next()` exhaustion (Done -> RecoveryExhausted).
//! - Fresh-mint guard: ±window bounds, ExternalBinary, verdict grace,
//! policy-hidden fall-through (fail closed).
//! - Fresh-mint guard: ±window bounds, tombstone grace.
//!
//! These tests use the same in-process `AuthManager` that production
//! does and inject a counting refresher so we can observe whether the
//! authority was consulted.
use super::*;
use crate::auth::config::GrokComConfig;
use crate::auth::error::{RefreshTokenError, RefreshTokenFailedReason};
use crate::auth::model::{AuthMode, GrokAuth};
use crate::auth::config::KimiCodeConfig;
use crate::auth::error::RefreshTokenError;
use crate::auth::model::{AuthMode, KimiAuth};
use crate::auth::refresh::{RefreshOutcome, TokenRefresher};
use crate::auth::storage::{read_auth_json, write_auth_json};
use chrono::{Duration, Utc};
use std::sync::atomic::{AtomicU32, Ordering};
/// The rejected wire bearer these tests seed into the manager.
fn rejected_cred() -> Option<GrokAuth> {
Some(GrokAuth {
fn rejected_cred() -> Option<KimiAuth> {
Some(KimiAuth {
key: "rejected-tok".into(),
..GrokAuth::test_default()
..KimiAuth::test_default()
})
}
@@ -382,17 +342,17 @@ mod tests {
impl TokenRefresher for OkRefresher {
async fn refresh(&self, _reason: crate::auth::manager::RefreshReason) -> RefreshOutcome {
self.calls.fetch_add(1, Ordering::SeqCst);
RefreshOutcome::Success(Box::new(GrokAuth {
RefreshOutcome::Success(Box::new(KimiAuth {
key: "fresh-from-authority".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
}))
}
}
/// Refresher fake: returns PermanentFailure (invalid_grant).
/// Refresher fake: returns PermanentFailure (rejected refresh token).
struct FailRefresher {
calls: Arc<AtomicU32>,
}
@@ -406,19 +366,19 @@ mod tests {
fn mgr() -> (tempfile::TempDir, Arc<AuthManager>) {
let dir = tempfile::tempdir().unwrap();
let m = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let m = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
(dir, m)
}
fn seed(mgr: &AuthManager, mode: AuthMode, refresh_token: Option<&str>) {
let auth = GrokAuth {
let auth = KimiAuth {
key: "rejected-tok".into(),
auth_mode: mode,
refresh_token: refresh_token.map(str::to_string),
// Past expiry so `current()` returns None and the refresh
// chain actually has to do work.
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
};
mgr.hot_swap(auth);
}
@@ -426,9 +386,9 @@ mod tests {
// -- TokenType dispatch matrix ----------------------------------------
#[tokio::test]
async fn dispatch_oidc_session_uses_refresh_chain() {
async fn dispatch_oauth_session_uses_refresh_chain() {
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, Some("rt"));
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
@@ -441,39 +401,25 @@ mod tests {
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn dispatch_external_binary_uses_refresh_chain() {
let (_d, m) = mgr();
seed(&m, AuthMode::External, None);
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
}));
let mut rec = m.unauthorized_recovery(rejected_cred());
let auth = rec.next().await.expect("external-binary recovery succeeds");
assert_eq!(auth.key, "fresh-from-authority");
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
// -- Fresh-mint guard --------------------------------------------------
/// Seed a *valid* (unexpired) in-memory token whose `create_time` lies
/// `mint_age` in the past (negative = clock stepped back since mint).
fn seed_valid(mgr: &AuthManager, mode: AuthMode, mint_age: Duration) {
mgr.hot_swap(GrokAuth {
fn seed_valid(mgr: &AuthManager, mint_age: Duration) {
mgr.hot_swap(KimiAuth {
key: "rejected-tok".into(),
auth_mode: mode,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
create_time: Utc::now() - mint_age,
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
expires_in: Some(3600),
..KimiAuth::test_default()
});
}
/// Run one recovery against a counting refresher; return the outcome and
/// how many times the authority was consulted.
async fn recover_with_ok_refresher(m: &Arc<AuthManager>) -> (Result<GrokAuth, AuthError>, u32) {
async fn recover_with_ok_refresher(m: &Arc<AuthManager>) -> (Result<KimiAuth, AuthError>, u32) {
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
@@ -484,9 +430,9 @@ mod tests {
}
#[tokio::test]
async fn fresh_mint_guard_skips_idp_for_freshly_minted_token() {
async fn fresh_mint_guard_skips_wire_for_freshly_minted_token() {
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::seconds(10));
seed_valid(&m, Duration::seconds(10));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("guard returns the live token").key,
@@ -495,23 +441,11 @@ mod tests {
assert_eq!(calls, 0, "a 10s-old token must not be re-minted");
}
#[tokio::test]
async fn fresh_mint_guard_applies_to_external_binary_tokens() {
let (_d, m) = mgr();
seed_valid(&m, AuthMode::External, Duration::seconds(10));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("guard returns the live token").key,
"rejected-tok"
);
assert_eq!(calls, 0);
}
#[tokio::test]
async fn fresh_mint_guard_treats_small_negative_age_as_fresh() {
// Clock stepped back slightly since mint (NTP nudge).
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::seconds(-60));
seed_valid(&m, Duration::seconds(-60));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("guard returns the live token").key,
@@ -525,19 +459,19 @@ mod tests {
// A large backwards clock step must not wedge recovery for the whole
// step: outside the ±window the guard stands down.
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::hours(-1));
seed_valid(&m, Duration::hours(-1));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("recovery should succeed").key,
"fresh-from-authority"
);
assert_eq!(calls, 1, "far-negative mint age must reach the IdP");
assert_eq!(calls, 1, "far-negative mint age must reach the wire");
}
#[tokio::test]
async fn fresh_mint_guard_lets_old_token_refresh() {
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::minutes(10));
seed_valid(&m, Duration::minutes(10));
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result.expect("recovery should succeed").key,
@@ -545,25 +479,25 @@ mod tests {
);
assert_eq!(
calls, 1,
"outside the guard window ServerRejected must reach the IdP"
"outside the guard window ServerRejected must reach the wire"
);
}
#[tokio::test]
async fn fresh_mint_guard_wins_over_cached_permanent_failure() {
// A fresh *valid* token is served even when a permanent-failure
// verdict is cached for it — mirrors `auth()`'s wire-valid grace arm;
// the verdict re-applies once the guard window passes.
async fn fresh_mint_guard_wins_over_cached_tombstone() {
// A fresh *valid* token is served even when a tombstone is cached
// for its refresh token — mirrors `auth()`'s wire-valid grace arm;
// the tombstone re-applies once the guard window passes.
let (_d, m) = mgr();
seed_valid(&m, AuthMode::Oidc, Duration::seconds(10));
seed_valid(&m, Duration::seconds(10));
m.record_permanent_failure(
"rejected-tok".into(),
"rt".into(),
RefreshTokenFailedReason::RefreshTokenRejected.into(),
);
let (result, calls) = recover_with_ok_refresher(&m).await;
assert_eq!(
result
.expect("guard precedes the verdict short-circuit")
.expect("guard precedes the tombstone short-circuit")
.key,
"rejected-tok"
);
@@ -571,65 +505,10 @@ mod tests {
}
#[tokio::test]
async fn fresh_mint_guard_never_returns_policy_hidden_token() {
// Wrong-team fresh token: `current()` hides it (vet_cached), so the
// guard must fall through to a normal refresh — fail closed.
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig {
force_login_team_uuid: Some(crate::auth::config::ForceLoginTeam::Single(
"team-good".into(),
)),
..GrokComConfig::default()
};
let m = Arc::new(AuthManager::new(dir.path(), cfg));
m.hot_swap(GrokAuth {
key: team_jwt("team-wrong"),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt".into()),
create_time: Utc::now(),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
});
let calls = Arc::new(AtomicU32::new(0));
m.set_refresher(Arc::new(OkRefresher {
calls: calls.clone(),
}));
let mut rec = m.unauthorized_recovery(rejected_cred());
let result = rec.next().await;
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"hidden token must not satisfy the guard"
);
if let Ok(auth) = result {
assert_ne!(
auth.key,
team_jwt("team-wrong"),
"wrong-team token must never be returned"
);
}
}
#[tokio::test]
async fn dispatch_legacy_session_returns_server_rejected_no_recovery() {
async fn dispatch_session_without_refresh_token_returns_server_rejected_no_recovery() {
// OAuth without refresh_token classifies as SessionNoRefresh.
let (_d, m) = mgr();
// WebLogin (no refresh_token) -> LegacySession.
seed(&m, AuthMode::WebLogin, None);
let mut rec = m.unauthorized_recovery(rejected_cred());
let err = rec.next().await.unwrap_err();
assert!(
matches!(err, AuthError::ServerRejectedNoRecovery),
"LegacySession recovery should surface ServerRejectedNoRecovery, got {err:?}",
);
}
#[tokio::test]
async fn dispatch_oidc_without_refresh_token_returns_server_rejected_no_recovery() {
// Oidc without refresh_token classifies as LegacySession.
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, None);
seed(&m, AuthMode::OAuth, None);
let mut rec = m.unauthorized_recovery(rejected_cred());
let err = rec.next().await.unwrap_err();
@@ -668,16 +547,17 @@ mod tests {
#[tokio::test]
async fn reload_from_disk_picks_up_different_token() {
let (dir, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, Some("rt"));
// Sibling process wrote a different valid token to disk.
let scope = m.grok_com_config().auth_scope();
let fresh = GrokAuth {
let scope = m.kimi_code_config().auth_scope();
let fresh = KimiAuth {
key: "fresh-from-disk".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
expires_in: Some(3600),
..KimiAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, fresh);
@@ -694,16 +574,17 @@ mod tests {
#[tokio::test]
async fn reload_from_disk_skips_same_token_then_proceeds_to_authority() {
let (dir, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, Some("rt"));
// Disk has the SAME token that was rejected -- skip, fall through.
let scope = m.grok_com_config().auth_scope();
let same = GrokAuth {
let scope = m.kimi_code_config().auth_scope();
let same = KimiAuth {
key: "rejected-tok".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
expires_in: Some(3600),
..KimiAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, same);
@@ -732,15 +613,11 @@ mod tests {
#[tokio::test]
async fn next_after_done_returns_recovery_exhausted() {
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, Some("rt"));
m.set_refresher(Arc::new(OkRefresher {
calls: Arc::new(AtomicU32::new(0)),
}));
// Pin non-devbox so DevboxRecovery can't adopt the seeded token (CI runs
// in K8s pods where is_devbox_environment() is true).
m.set_devbox_env_for_test(false);
let mut rec = m.unauthorized_recovery(rejected_cred());
let _ = rec.next().await.unwrap();
let err = loop {
@@ -773,9 +650,8 @@ mod tests {
}
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, Some("rt"));
m.set_refresher(Arc::new(TransientFailRefresher));
m.set_devbox_env_for_test(false);
let mut rec = m.unauthorized_recovery(rejected_cred());
// First next(): the authority's transient error propagates as-is.
@@ -799,21 +675,17 @@ mod tests {
!forces_manual_reauth(&err),
"a transient exhaustion must not force a manual re-login",
);
assert!(
!relay_should_cancel(&err),
"the relay must reconnect (not cancel) on a transient exhaustion",
);
}
// -- Permanent failure short-circuit (cross-check) ------------
// -- Tombstone short-circuit (cross-check) ------------
#[tokio::test]
async fn refresh_authority_short_circuits_on_cached_permanent_failure() {
async fn refresh_authority_short_circuits_on_cached_tombstone() {
let (_d, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
// Pre-record a permanent failure scoped to the seeded credential.
seed(&m, AuthMode::OAuth, Some("rt"));
// Pre-record a tombstone scoped to the seeded refresh token.
m.record_permanent_failure(
"rejected-tok".into(),
"rt".into(),
RefreshTokenFailedReason::RefreshTokenRejected.into(),
);
@@ -831,7 +703,7 @@ mod tests {
assert_eq!(
calls.load(Ordering::SeqCst),
0,
"refresher must not be invoked when permanent_failure is cached",
"refresher must not be invoked while the tombstone cooldown is live",
);
}
@@ -843,15 +715,15 @@ mod tests {
#[tokio::test]
async fn reload_from_disk_rejects_expired_different_token() {
let (dir, m) = mgr();
seed(&m, AuthMode::Oidc, Some("rt"));
seed(&m, AuthMode::OAuth, Some("rt"));
let scope = m.grok_com_config().auth_scope();
let expired_different = GrokAuth {
let scope = m.kimi_code_config().auth_scope();
let expired_different = KimiAuth {
key: "different-but-expired".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt-new".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
};
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(scope, expired_different);
@@ -870,64 +742,4 @@ mod tests {
);
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
// -- force_login_team_uuid pin enforced on the 401-recovery path -------
fn ensure_crypto_provider() {
let _ = jsonwebtoken::crypto::rust_crypto::DEFAULT_PROVIDER.install_default();
}
fn team_jwt(principal_id: &str) -> String {
ensure_crypto_provider();
jsonwebtoken::encode(
&jsonwebtoken::Header::new(jsonwebtoken::Algorithm::HS256),
&serde_json::json!({
"sub": "user-1",
"principal_type": "Team",
"principal_id": principal_id,
"exp": 9999999999u64,
}),
&jsonwebtoken::EncodingKey::from_secret(b"test-secret"),
)
.unwrap()
}
/// A sibling writes a wrong-team token to disk; 401 recovery (relay path)
/// must reject + clear it at `next()`, not hand it back as a bearer.
#[tokio::test]
async fn recovery_rejects_wrong_team_adopted_disk_token() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig {
force_login_team_uuid: Some(crate::auth::config::ForceLoginTeam::Single(
"team-good".into(),
)),
..GrokComConfig::default()
};
let scope = cfg.auth_scope();
let m = Arc::new(AuthManager::new(dir.path(), cfg));
// In-memory: the rejected (expired) session that triggered recovery.
seed(&m, AuthMode::Oidc, Some("rt"));
// Disk: a different, non-expired, *wrong-team* token a sibling wrote.
let mut store = read_auth_json(&dir.path().join("auth.json")).unwrap_or_default();
store.insert(
scope,
GrokAuth {
key: team_jwt("team-wrong"),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-sibling".into()),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
},
);
write_auth_json(&dir.path().join("auth.json"), &store).unwrap();
let mut rec = m.unauthorized_recovery(rejected_cred());
let err = rec.next().await.unwrap_err();
assert!(
matches!(err, AuthError::PinnedTeamMismatch { .. }),
"recovery must reject a wrong-team disk token, got {err:?}"
);
}
}
crates/codegen/kigi-shell/src/auth/refresh/auth_backend_contract_tests.rs
-351
View File
@@ -1,351 +0,0 @@
//! End-to-end auth-backend contract tests: a mock IdP whose `/token` response
//! is forced per case, asserting the refresh outcome, the storm cap, and the
//! terminal-error classification on the live recovery path.
use super::*;
use crate::auth::error::RefreshTokenFailedReason;
use crate::auth::{GrokAuth, GrokComConfig};
use chrono::{Duration, Utc};
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
/// Mock IdP: OIDC discovery + a `/token` endpoint returning a fixed
/// `(status, body)` and counting every hit, plus the `/user` endpoint
/// `AuthManager::update` calls after a successful refresh. `delay_ms` widens
/// the in-lock window so concurrent callers queue on `refresh_lock`.
async fn start_idp(
token_status: u16,
token_body: String,
hits: Arc<AtomicU32>,
delay_ms: u64,
) -> (String, tokio::task::JoinHandle<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port());
let disco = base.clone();
let app = axum::Router::new()
.route(
"/.well-known/openid-configuration",
axum::routing::get(move || {
let b = disco.clone();
async move {
axum::Json(serde_json::json!({
"authorization_endpoint": format!("{b}/authorize"),
"token_endpoint": format!("{b}/token"),
}))
}
}),
)
.route(
"/token",
axum::routing::post(move || {
let hits = hits.clone();
let body = token_body.clone();
async move {
hits.fetch_add(1, Ordering::SeqCst);
if delay_ms > 0 {
tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
}
(
axum::http::StatusCode::from_u16(token_status).unwrap(),
body,
)
}
}),
)
.route(
"/user",
axum::routing::get(|| async {
axum::Json(serde_json::json!({ "userId": "user-42", "email": "u@corp.com" }))
}),
);
let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(base, handle)
}
fn expired_oidc(base_url: &str) -> GrokAuth {
GrokAuth {
key: "expired-at".into(),
create_time: Utc::now() - Duration::hours(2),
user_id: "user-42".into(),
auth_mode: crate::auth::model::AuthMode::Oidc,
refresh_token: Some("rt-under-test".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
oidc_issuer: Some(base_url.to_owned()),
oidc_client_id: Some("client-under-test".into()),
..GrokAuth::test_default()
}
}
#[derive(Debug)]
enum Expect {
Success,
Permanent(RefreshTokenFailedReason),
Transient,
}
/// The IdP token-endpoint contract: each response shape maps to one outcome.
/// `invalid_grant`/`invalid_client` are the only permanent verdicts; status
/// blips and unrecognized codes stay transient (never permanent-lock).
#[tokio::test]
async fn auth_backend_contract_token_responses_map_to_outcomes() {
use RefreshTokenFailedReason::{ClientRejected, RefreshTokenRejected};
let cases: &[(&str, u16, &str, Expect)] = &[
(
"success",
200,
r#"{"access_token":"fresh","refresh_token":"fresh-rt","expires_in":3600}"#,
Expect::Success,
),
(
"invalid_grant",
400,
r#"{"error":"invalid_grant"}"#,
Expect::Permanent(RefreshTokenRejected),
),
(
"invalid_client",
401,
r#"{"error":"invalid_client"}"#,
Expect::Permanent(ClientRejected),
),
("server_error_5xx", 503, "{}", Expect::Transient),
("rate_limited_429", 429, "{}", Expect::Transient),
(
"temporarily_unavailable",
400,
r#"{"error":"temporarily_unavailable"}"#,
Expect::Transient,
),
("bare_4xx_no_body", 400, "", Expect::Transient),
("malformed_body", 400, "not json", Expect::Transient),
// Proxy/WAF-mangled bodies must degrade to retry, never a false permanent
// lock: a nested error object or a non-string `error` is not a recognized
// top-level code, so it stays transient.
(
"nested_error_object",
400,
r#"{"error":{"code":"invalid_grant"}}"#,
Expect::Transient,
),
(
"non_string_error",
400,
r#"{"error":123}"#,
Expect::Transient,
),
];
for (name, status, body, expect) in cases {
let hits = Arc::new(AtomicU32::new(0));
let (base_url, server) = start_idp(*status, body.to_string(), hits.clone(), 0).await;
let dir = tempfile::tempdir().unwrap();
let auth_manager = Arc::new(
AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url),
);
auth_manager.hot_swap(expired_oidc(&base_url));
let refresher = OidcRefresher::new(auth_manager.clone());
let result = refresher.refresh(RefreshReason::ServerRejected).await;
match (expect, &result) {
(Expect::Success, RefreshOutcome::Success(_)) => {}
(Expect::Permanent(want), RefreshOutcome::PermanentFailure { error, .. }) => {
assert_eq!(error.reason, *want, "{name}: wrong permanent reason");
}
(Expect::Transient, RefreshOutcome::TransientFailure { .. }) => {}
(exp, got) => panic!("{name}: expected {exp:?}, got {got:?}"),
}
server.abort();
}
}
/// A burst of concurrent 401s on the same revoked refresh token must hit the
/// IdP exactly once. The callers serialize on `refresh_lock`; the leader records
/// the verdict before releasing, so the in-lock re-check (`refresh_chain` step
/// 1b) short-circuits every follower. Delete step 1b and the count climbs to N.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn auth_backend_contract_concurrent_401s_hit_idp_once() {
let hits = Arc::new(AtomicU32::new(0));
// 100ms /token delay so every caller passes the pre-lock check and queues
// on refresh_lock before the leader records the verdict, exercising step 1b.
let (base_url, server) = start_idp(
400,
r#"{"error":"invalid_grant"}"#.to_string(),
hits.clone(),
100,
)
.await;
let dir = tempfile::tempdir().unwrap();
let auth_manager = Arc::new(
AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url),
);
auth_manager.hot_swap(expired_oidc(&base_url));
auth_manager.set_refresher(Arc::new(OidcRefresher::new(auth_manager.clone())));
let mut tasks = Vec::new();
for _ in 0..6 {
let auth_manager = auth_manager.clone();
tasks.push(tokio::spawn(async move { auth_manager.auth().await }));
}
for t in tasks {
let outcome = t.await.unwrap();
assert!(
matches!(
outcome,
Err(crate::auth::AuthError::Refresh(
crate::auth::RefreshTokenError::Permanent(_)
))
),
"every concurrent caller must fail permanently on a revoked refresh token, got {outcome:?}",
);
}
assert_eq!(
hits.load(Ordering::SeqCst),
1,
"concurrent 401s on one dead credential must hit the IdP exactly once",
);
server.abort();
}
/// The classification loop through the live recovery state machine: a dead
/// refresh token terminates recovery with an error that forces a manual
/// re-login; a refreshable token auto-refreshes.
#[tokio::test]
async fn auth_backend_contract_dead_token_forces_manual_reauth() {
// A dead refresh token terminates recovery with a forced-relogin error.
let hits = Arc::new(AtomicU32::new(0));
let (url, server) = start_idp(400, r#"{"error":"invalid_grant"}"#.to_string(), hits, 0).await;
let dir = tempfile::tempdir().unwrap();
let auth_manager =
Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&url));
auth_manager.hot_swap(expired_oidc(&url));
auth_manager.set_refresher(Arc::new(OidcRefresher::new(auth_manager.clone())));
let err = auth_manager
.unauthorized_recovery(auth_manager.current_or_expired())
.next()
.await
.expect_err("a dead refresh token must fail recovery");
assert!(
crate::auth::recovery::forces_manual_reauth(&err),
"a dead refresh token must be a forced-relogin error, got {err:?}",
);
server.abort();
// Refreshable token: recovery auto-refreshes.
let ok_hits = Arc::new(AtomicU32::new(0));
let (ok_url, ok_server) = start_idp(
200,
r#"{"access_token":"fresh","refresh_token":"fresh-rt","expires_in":3600}"#.to_string(),
ok_hits,
0,
)
.await;
let ok_dir = tempfile::tempdir().unwrap();
let ok_manager = Arc::new(
AuthManager::new(ok_dir.path(), GrokComConfig::default()).with_proxy_base_url(&ok_url),
);
ok_manager.hot_swap(expired_oidc(&ok_url));
ok_manager.set_refresher(Arc::new(OidcRefresher::new(ok_manager.clone())));
let refreshed = ok_manager
.unauthorized_recovery(ok_manager.current_or_expired())
.next()
.await
.expect("a refreshable token must auto-refresh");
assert_eq!(
refreshed.key, "fresh",
"recovery must return the fresh token"
);
ok_server.abort();
}
/// Consecutive transient failures self-heal up to a bound, then escalate to a
/// non-sticky `Other` permanent failure (which ages out via the TTL). A
/// regression here would turn recoverable blips into a permanent `/login`.
#[tokio::test]
async fn auth_backend_contract_transient_failures_escalate_to_non_sticky_permanent() {
let hits = Arc::new(AtomicU32::new(0));
// Persistent 503: every refresh attempt is transient.
let (base_url, server) = start_idp(503, "{}".to_string(), hits, 0).await;
let dir = tempfile::tempdir().unwrap();
let auth_manager = Arc::new(
AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&base_url),
);
auth_manager.hot_swap(expired_oidc(&base_url));
// One refresher instance: it owns the consecutive-failure counter.
let refresher = OidcRefresher::new(auth_manager.clone());
let mut outcomes = Vec::new();
for _ in 0..3 {
outcomes.push(refresher.refresh(RefreshReason::ServerRejected).await);
}
assert!(
matches!(outcomes[0], RefreshOutcome::TransientFailure { .. }),
"first blip is transient, not a lockout: {:?}",
outcomes[0],
);
match &outcomes[2] {
RefreshOutcome::PermanentFailure { error, .. } => {
assert_eq!(
error.reason,
RefreshTokenFailedReason::Other,
"escalation must use the generic Other reason",
);
assert!(
!error.reason.is_sticky(),
"an escalated transient must age out, not strand the user forever",
);
}
other => panic!("repeated transients must escalate to a permanent Other, got {other:?}"),
}
server.abort();
}
/// Two `AuthManager`s sharing one auth.json stand in for two CLI processes: the
/// auth.json flock must serialize their refreshes so the shared refresh token is
/// spent at the IdP exactly once. The loser adopts the rotated token from disk
/// instead of racing a second exchange (which the IdP could revoke as reuse).
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn auth_backend_contract_two_instances_share_one_idp_call() {
let hits = Arc::new(AtomicU32::new(0));
let (url, server) = start_idp(
200,
r#"{"access_token":"fresh","refresh_token":"fresh-rt","expires_in":3600}"#.to_string(),
hits.clone(),
100,
)
.await;
let dir = tempfile::tempdir().unwrap();
// Distinct managers, same on-disk auth.json (separate flock OFDs => they
// genuinely contend, like two processes).
let new_instance = || {
let m = Arc::new(
AuthManager::new(dir.path(), GrokComConfig::default()).with_proxy_base_url(&url),
);
m.hot_swap(expired_oidc(&url));
m.set_refresher(Arc::new(OidcRefresher::new(m.clone())));
m
};
let a = new_instance();
let b = new_instance();
let (ra, rb) = tokio::join!(a.auth(), b.auth());
assert_eq!(ra.expect("instance A must obtain a token").key, "fresh");
assert_eq!(rb.expect("instance B must obtain a token").key, "fresh");
assert_eq!(
hits.load(Ordering::SeqCst),
1,
"two instances sharing auth.json must spend the refresh token at the IdP only once",
);
server.abort();
}
crates/codegen/kigi-shell/src/auth/refresh/external_refresher.rs
-176
View File
@@ -1,176 +0,0 @@
use std::sync::Arc;
use crate::auth::error::RefreshTokenFailedReason;
use crate::auth::manager::RefreshReason;
use super::{ExternalCommandRunner, RefreshOutcome, TokenRefresher};
/// Refreshes by re-running the operator's external auth binary via
/// `spawn_blocking`. Pure data return -- mutation lives in
/// `refresh_chain` (honors the [`TokenRefresher`] no-mutation contract).
pub(crate) struct ExternalBinaryRefresher {
runner: Arc<dyn ExternalCommandRunner>,
command: String,
timeout: std::time::Duration,
}
impl ExternalBinaryRefresher {
pub(crate) fn new(runner: Arc<dyn ExternalCommandRunner>, command: String) -> Self {
Self {
runner,
command,
timeout: EXTERNAL_REFRESH_TIMEOUT,
}
}
/// Override the binary timeout (tests use a short one to exercise the
/// timeout arm without a real 30s wait).
#[cfg(test)]
pub(crate) fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
self.timeout = timeout;
self
}
/// A failed binary run is a single-strike `Other` permanent failure; the
/// `PERMANENT_FAILURE_TTL` lets a flaky binary self-heal without `/login`.
/// No consecutive-blip tolerance like OIDC: a local binary failure is a
/// stronger signal than a network refresh blip.
fn record_failure(&self, message: String) -> RefreshOutcome {
tracing::warn!(%message, "auth: external binary refresh failed -> permanent");
// No token key in the binary flow; the caller scopes the verdict.
RefreshOutcome::permanent(RefreshTokenFailedReason::Other, None)
}
}
/// Timeout for the external auth binary. If the binary hangs, the
/// `spawn_blocking` thread is leaked (it cannot be interrupted), but this is
/// acceptable: the thread holds no locks and mutates no shared state.
const EXTERNAL_REFRESH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
#[async_trait::async_trait]
impl TokenRefresher for ExternalBinaryRefresher {
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome {
tracing::debug!(?reason, "auth: external binary refresh starting");
let runner = self.runner.clone();
let cmd = self.command.clone();
let timeout_ms = self.timeout.as_millis() as u64;
match tokio::time::timeout(
self.timeout,
tokio::task::spawn_blocking(move || runner.run_external_command(&cmd)),
)
.await
{
Err(_elapsed) => {
tracing::warn!(
timeout_ms,
"auth: external binary refresh timed out (thread leaked)"
);
crate::unified_log::warn(
"auth.refresh.external_timeout",
None,
Some(serde_json::json!({ "timeout_ms": timeout_ms })),
);
self.record_failure(format!("external binary timed out after {timeout_ms}ms"))
}
Ok(Ok(Some(auth))) => {
crate::unified_log::info("auth: external binary refresh succeeded", None, None);
RefreshOutcome::success(auth)
}
Ok(Ok(None)) => {
crate::unified_log::warn(
"auth: external binary refresh returned no token",
None,
None,
);
self.record_failure("external binary returned no token".into())
}
Ok(Err(e)) => {
tracing::warn!(error = %e, "auth: external binary refresh task failed");
self.record_failure(format!("external binary task failed: {e}"))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::GrokAuth;
/// Minimal runner whose external command returns a fixed result.
struct FakeRunner {
external_result: Option<GrokAuth>,
}
impl ExternalCommandRunner for FakeRunner {
fn run_external_command(&self, _command: &str) -> Option<GrokAuth> {
self.external_result.clone()
}
}
/// A failed binary run is a single-strike `Other` permanent failure that is
/// NON-sticky: it must age out via the TTL, never lock an external-binary
/// user out forever. (Flipping this to a sticky reason would be a silent
/// lockout regression.)
#[tokio::test]
async fn external_binary_failure_is_single_strike_non_sticky_permanent() {
let refresher = ExternalBinaryRefresher::new(
Arc::new(FakeRunner {
external_result: None,
}),
"auth-binary".into(),
);
match refresher.refresh(RefreshReason::ServerRejected).await {
RefreshOutcome::PermanentFailure { error, .. } => {
assert_eq!(error.reason, RefreshTokenFailedReason::Other);
assert!(
!error.reason.is_sticky(),
"external-binary failure must age out, not strand the user forever",
);
}
other => panic!("a failed binary run must be a permanent Other failure, got {other:?}"),
}
}
/// A binary that outlives the (test-shortened) timeout hits the `Elapsed`
/// arm and maps to the same non-sticky `Other` permanent failure.
#[tokio::test]
async fn external_binary_timeout_is_non_sticky_permanent() {
struct SlowRunner;
impl ExternalCommandRunner for SlowRunner {
fn run_external_command(&self, _command: &str) -> Option<GrokAuth> {
std::thread::sleep(std::time::Duration::from_millis(50));
Some(GrokAuth::test_default())
}
}
let refresher = ExternalBinaryRefresher::new(Arc::new(SlowRunner), "auth-binary".into())
.with_timeout(std::time::Duration::from_millis(5));
match refresher.refresh(RefreshReason::ServerRejected).await {
RefreshOutcome::PermanentFailure { error, .. } => {
assert_eq!(error.reason, RefreshTokenFailedReason::Other);
assert!(
!error.reason.is_sticky(),
"timeout must age out, not strand"
);
}
other => panic!("a timed-out binary must be a permanent Other failure, got {other:?}"),
}
}
#[tokio::test]
async fn external_binary_success_returns_fresh_token() {
let token = GrokAuth {
key: "ext-fresh".into(),
..GrokAuth::test_default()
};
let refresher = ExternalBinaryRefresher::new(
Arc::new(FakeRunner {
external_result: Some(token),
}),
"auth-binary".into(),
);
match refresher.refresh(RefreshReason::ServerRejected).await {
RefreshOutcome::Success(auth) => assert_eq!(auth.key, "ext-fresh"),
other => panic!("a successful binary run must return Success, got {other:?}"),
}
}
}
crates/codegen/kigi-shell/src/auth/refresh/kimi_refresher.rs
+356
View File
@@ -0,0 +1,356 @@
//! Kimi Code token refresher: drives `POST /api/oauth/token` with
//! `grant_type=refresh_token` through the [`TokenRefresher`] seam.
//!
//! Ports kimi-cli `OAuthManager._refresh_tokens`' sibling-safety behavior:
//! the persisted credential is re-read before the wire call (adopt a
//! rotation instead of refreshing), and after a 401/403 the persisted
//! credential is re-read once more (with a 1s grace) so a concurrent
//! process's freshly rotated token is adopted instead of tombstoning it.
use std::sync::Arc;
use crate::auth::error::RefreshTokenFailedReason;
use crate::auth::kimi_oauth::{self, RefreshError};
use crate::auth::manager::RefreshReason;
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
/// (kimi-cli parity: `await asyncio.sleep(1)`).
const POST_UNAUTHORIZED_GRACE: std::time::Duration = std::time::Duration::from_secs(1);
pub(crate) struct KimiRefresher {
auth: Arc<dyn AuthSnapshot>,
/// OAuth host; `kigi_env::oauth_host()` in production, injectable for
/// wiremock tests.
host: String,
}
impl KimiRefresher {
pub(crate) fn new(auth: Arc<dyn AuthSnapshot>, host: String) -> Self {
Self { auth, host }
}
/// Post-401 sibling check (kimi-cli parity): wait a beat, re-read the
/// persisted credential, and adopt it when its refresh token differs
/// from the one the server just rejected.
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!({
"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 KimiRefresher {
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome {
tracing::info!(?reason, "auth: kimi refresh attempt starting");
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 between the
// refresh_chain disk check (under lock) and here. 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!({
"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");
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");
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"
);
match kimi_oauth::refresh_token(&self.host, &refresh_token).await {
Ok(new_auth) => {
kigi_log::unified_log::info(
"auth.refresh.token_rotated",
None,
Some(serde_json::json!({
"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");
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!({
"status": status,
"description": description,
"rt_prefix": crate::auth::token_suffix(&refresh_token),
})),
);
RefreshOutcome::permanent(
RefreshTokenFailedReason::RefreshTokenRejected,
Some(refresh_token),
)
}
// kimi-cli parity: non-401 failures never tombstone; the next
// 60s tick (or pre-request check) retries.
Err(
e @ (RefreshError::Exhausted { .. }
| RefreshError::Fatal { .. }
| RefreshError::Local(_)),
) => {
tracing::warn!(error = %e, "auth: refresh attempt failed (transient)");
kigi_log::unified_log::warn(
"auth.refresh.transient_wire_failure",
None,
Some(serde_json::json!({ "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 parking_lot::Mutex;
use wiremock::matchers::{body_string_contains, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
/// Scriptable snapshot: `disk` can be swapped mid-test to simulate a
/// sibling process rotating the persisted credential.
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 valid_session(key: &str, rt: &str) -> KimiAuth {
KimiAuth {
key: key.into(),
refresh_token: Some(rt.into()),
expires_at: Some(Utc::now() + Duration::hours(2)),
expires_in: Some(7200),
..KimiAuth::test_default()
}
}
fn token_json(access: &str, refresh: &str) -> serde_json::Value {
serde_json::json!({
"access_token": access,
"refresh_token": refresh,
"expires_in": 3600,
"scope": "kimi-code",
"token_type": "bearer",
})
}
#[tokio::test]
async fn refresh_success_returns_rotated_token() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.and(body_string_contains("refresh_token=rt-old"))
.respond_with(ResponseTemplate::new(200).set_body_json(token_json("at-new", "rt-new")))
.expect(1)
.mount(&server)
.await;
let stale = expired_session("at-old", "rt-old");
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
let refresher = KimiRefresher::new(snap, server.uri());
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
let RefreshOutcome::Success(new_auth) = outcome else {
panic!("expected success, got {outcome:?}");
};
assert_eq!(new_auth.key, "at-new");
assert_eq!(new_auth.refresh_token.as_deref(), Some("rt-new"));
}
#[tokio::test]
async fn adopts_valid_sibling_token_without_wire_call() {
// Disk has a fresh token with a different key: adopt, no HTTP.
let server = MockServer::start().await;
// No mock mounted: any request would 404 and fail the refresh.
let snap = FakeSnapshot::new(
Some(expired_session("at-old", "rt-old")),
Some(valid_session("at-sibling", "rt-sibling")),
);
let refresher = KimiRefresher::new(snap, server.uri());
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
let RefreshOutcome::Success(adopted) = outcome else {
panic!("expected sibling adoption, got {outcome:?}");
};
assert_eq!(adopted.key, "at-sibling");
assert!(
server.received_requests().await.unwrap().is_empty(),
"sibling adoption must not consume a refresh token on the wire"
);
}
#[tokio::test]
async fn unauthorized_tombstones_the_tried_refresh_token() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(
ResponseTemplate::new(401)
.set_body_json(serde_json::json!({ "error_description": "revoked" })),
)
.expect(1)
.mount(&server)
.await;
let stale = expired_session("at-old", "rt-dead");
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
let refresher = KimiRefresher::new(snap, server.uri());
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("rt-dead"));
}
#[tokio::test]
async fn unauthorized_adopts_sibling_rotation_instead_of_tombstoning() {
// 401 lands, but by the time we re-check, a sibling has persisted a
// rotated credential — adopt it (the mutual-logout race guard).
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(401))
.expect(1)
.mount(&server)
.await;
let stale = expired_session("at-old", "rt-dead");
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
let refresher = KimiRefresher::new(snap.clone(), server.uri());
// Swap the persisted credential while the wire call is in flight.
let rotator = {
let snap = snap.clone();
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(300)).await;
*snap.disk.lock() = Some(valid_session("at-rotated", "rt-rotated"));
})
};
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
rotator.await.unwrap();
let RefreshOutcome::Success(adopted) = outcome else {
panic!("expected rotation adoption, got {outcome:?}");
};
assert_eq!(adopted.key, "at-rotated");
}
#[tokio::test]
async fn wire_exhaustion_is_transient_not_tombstoned() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/oauth/token"))
.respond_with(ResponseTemplate::new(503))
.expect(3)
.mount(&server)
.await;
let stale = expired_session("at-old", "rt-old");
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
let refresher = KimiRefresher::new(snap, server.uri());
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
assert!(
matches!(outcome, RefreshOutcome::TransientFailure { .. }),
"5xx exhaustion must stay transient: {outcome:?}"
);
}
#[tokio::test]
async fn no_credential_is_transient() {
let server = MockServer::start().await;
let snap = FakeSnapshot::new(None, None);
let refresher = KimiRefresher::new(snap, server.uri());
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
assert!(matches!(outcome, RefreshOutcome::TransientFailure { .. }));
}
}
crates/codegen/kigi-shell/src/auth/refresh/mod.rs
+39 -117
View File
@@ -1,42 +1,38 @@
mod external_refresher;
mod oidc_refresher;
mod kimi_refresher;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use crate::auth::manager::AuthManager;
pub(crate) use crate::auth::manager::RefreshReason;
use crate::auth::model::GrokAuth;
use crate::auth::model::KimiAuth;
use external_refresher::ExternalBinaryRefresher;
pub(crate) use oidc_refresher::OidcRefresher;
pub(crate) use kimi_refresher::KimiRefresher;
/// Read-only view of `AuthManager` for refreshers. Enforces the
/// no-mutation contract on *credential* state at the type level: refreshers
/// hold `Arc<dyn AuthSnapshot>` and physically cannot call `update()`,
/// `clear()`, `hot_swap()`, or `refresh_chain()`.
pub(crate) trait AuthSnapshot: Send + Sync {
/// Read the current in-memory bearer outside the early-invalidation buffer.
fn current(&self) -> Option<GrokAuth>;
/// Read the current in-memory bearer outside the refresh threshold.
fn current(&self) -> Option<KimiAuth>;
/// Read the expired in-memory bearer (for its `refresh_token`).
fn expired_auth(&self) -> Option<GrokAuth>;
/// Re-read auth.json from disk for the configured scope. Read-only w.r.t.
/// credentials, but may advance disk-observation state and emit transition
/// telemetry (not credential mutation).
fn read_disk_auth(&self) -> Option<GrokAuth>;
fn expired_auth(&self) -> Option<KimiAuth>;
/// Re-read the persisted credential (keyring → file) for the configured
/// scope. Read-only w.r.t. credentials, but may advance disk-observation
/// state and emit transition telemetry (not credential mutation).
fn read_disk_auth(&self) -> Option<KimiAuth>;
/// Whether the in-memory bearer is expired.
fn is_expired(&self) -> bool;
}
impl AuthSnapshot for AuthManager {
fn current(&self) -> Option<GrokAuth> {
fn current(&self) -> Option<KimiAuth> {
self.current()
}
fn expired_auth(&self) -> Option<GrokAuth> {
fn expired_auth(&self) -> Option<KimiAuth> {
self.expired_auth()
}
fn read_disk_auth(&self) -> Option<GrokAuth> {
fn read_disk_auth(&self) -> Option<KimiAuth> {
self.read_disk_auth()
}
fn is_expired(&self) -> bool {
@@ -44,31 +40,18 @@ impl AuthSnapshot for AuthManager {
}
}
/// Capability to run the operator's external auth binary. Split out of
/// [`AuthSnapshot`] so OIDC refreshers (read-only) physically cannot reach it
/// (interface segregation); only [`ExternalBinaryRefresher`] depends on it.
pub(crate) trait ExternalCommandRunner: Send + Sync {
/// Run the external auth binary and return the parsed output.
fn run_external_command(&self, command: &str) -> Option<GrokAuth>;
}
impl ExternalCommandRunner for AuthManager {
fn run_external_command(&self, command: &str) -> Option<GrokAuth> {
self.run_external_refresh_command(command)
}
}
/// The credential a refresh would send to the IdP: disk refresh-token first,
/// then the expired in-mem bearer, then current (only on `ServerRejected`).
/// Single source of truth shared by [`OidcRefresher::refresh`] (the attempt) and
/// `AuthManager::attempted_verdict_key` (the verdict scope), so the two can't
/// drift. The caller supplies the disk read: the verdict path passes a
/// side-effect-free read, the refresher the observing one.
/// The credential a refresh would send to the OAuth host: persisted
/// refresh-token first, then the expired in-mem bearer, then current (only on
/// `ServerRejected`). Single source of truth shared by
/// [`KimiRefresher::refresh`] (the attempt) and
/// `AuthManager::attempted_tombstone_key` (the tombstone scope), so the two
/// can't drift. The caller supplies the persisted read: the tombstone path
/// passes a side-effect-free read, the refresher the observing one.
pub(crate) fn resolve_refresh_credential(
snap: &dyn AuthSnapshot,
disk_auth: Option<GrokAuth>,
disk_auth: Option<KimiAuth>,
reason: RefreshReason,
) -> Option<GrokAuth> {
) -> Option<KimiAuth> {
disk_auth
.filter(|a| a.refresh_token.is_some())
.or_else(|| snap.expired_auth())
@@ -84,18 +67,16 @@ pub(crate) fn resolve_refresh_credential(
#[must_use = "RefreshOutcome encodes a state transition; route it through refresh_chain"]
pub(crate) enum RefreshOutcome {
/// Authority returned a fresh token. Caller persists via `update()`.
Success(Box<GrokAuth>),
/// Terminal failure (e.g. invalid_grant), or a transient escalated to
/// `Other` after repeated blips. Caller records a verdict scoped to the
/// rejected credential and retains it (`RefreshTokenRejected` is sticky,
/// the rest age out past the TTL).
Success(Box<KimiAuth>),
/// Terminal failure (401/403 from the OAuth host). Caller records a
/// tombstone scoped to the rejected refresh token; the 300s cooldown (or
/// a rotated persisted refresh token) clears it.
PermanentFailure {
error: crate::auth::error::RefreshTokenFailedError,
/// Key of the credential the refresher actually sent to the IdP, so
/// `refresh_chain` scopes the verdict to it. `None` when the authority
/// has no token key (external binary flow); the caller falls back to
/// its own resolution.
tried_key: Option<String>,
/// The refresh-token value the refresher actually sent, so
/// `refresh_chain` scopes the tombstone to it. `None` when the
/// attempt never reached the wire.
rejected_refresh_token: Option<String>,
},
/// Transient / unknown failure. Caller may retry later. Message-only: the
/// underlying cause is logged structurally at the refresher, then flattened
@@ -105,19 +86,19 @@ pub(crate) enum RefreshOutcome {
impl RefreshOutcome {
/// A fresh credential from the authority (hides the `Box`).
pub(crate) fn success(auth: GrokAuth) -> Self {
pub(crate) fn success(auth: KimiAuth) -> Self {
Self::Success(Box::new(auth))
}
/// Terminal failure for an already-classified reason against the credential
/// `tried_key` (the one actually sent to the IdP).
/// Terminal failure for an already-classified reason against the
/// refresh token actually sent to the OAuth host.
pub(crate) fn permanent(
reason: crate::auth::error::RefreshTokenFailedReason,
tried_key: Option<String>,
rejected_refresh_token: Option<String>,
) -> Self {
Self::PermanentFailure {
error: reason.into(),
tried_key,
rejected_refresh_token,
}
}
@@ -139,67 +120,8 @@ pub(crate) trait TokenRefresher: Send + Sync {
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome;
}
pub(crate) fn build_refresher(
auth_manager: Arc<AuthManager>,
auth_provider_command: Option<String>,
) -> Arc<dyn TokenRefresher> {
match auth_provider_command {
Some(cmd) => {
let runner: Arc<dyn ExternalCommandRunner> = auth_manager;
Arc::new(ExternalBinaryRefresher::new(runner, cmd))
}
None => {
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
Arc::new(OidcRefresher::new(snapshot))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::{AuthMode, GrokAuth, GrokComConfig};
use chrono::{Duration, Utc};
/// auth_token_ttl makes is_token_expired use create_time + ttl for
/// External tokens without expires_at, instead of the 30-day fallback.
#[test]
fn token_ttl_expires_external_token_by_create_time() {
let dir = tempfile::tempdir().unwrap();
let cfg = GrokComConfig {
auth_token_ttl: Some(3600), // 1 hour
..GrokComConfig::default()
};
let mgr = AuthManager::new(dir.path(), cfg);
// Token created 2 hours ago, no expires_at. With auth_token_ttl=3600,
// is_token_expired should return true (age 2h > ttl 1h).
let old_token = GrokAuth {
key: "old-external-token".into(),
auth_mode: AuthMode::External,
create_time: Utc::now() - Duration::hours(2),
expires_at: None,
..GrokAuth::test_default()
};
mgr.hot_swap(old_token);
assert!(
mgr.current().is_none(),
"expired external token via auth_token_ttl"
);
assert!(mgr.is_expired());
// Fresh token created just now — should be valid.
let new_token = GrokAuth {
key: "new-external-token".into(),
auth_mode: AuthMode::External,
create_time: Utc::now(),
expires_at: None,
..GrokAuth::test_default()
};
mgr.hot_swap(new_token);
assert!(
mgr.current().is_some(),
"fresh external token should be valid"
);
}
/// Build the production refresher against `kigi_env::oauth_host()`.
pub(crate) fn build_refresher(auth_manager: Arc<AuthManager>) -> Arc<dyn TokenRefresher> {
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host()))
}
crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher.rs
-260
View File
@@ -1,260 +0,0 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use crate::auth::error::RefreshTokenFailedReason;
use crate::auth::manager::RefreshReason;
use crate::auth::oidc::OidcRefreshResult;
use super::{AuthSnapshot, RefreshOutcome, TokenRefresher};
#[cfg(test)]
use crate::auth::manager::AuthManager;
/// Escalate to `PermanentFailure` after this many consecutive transient
/// failures (then `PERMANENT_FAILURE_TTL` allows recovery). OIDC tolerates more
/// blips than `ExternalBinaryRefresher` (1) since network refreshes flake more
/// than a local binary.
const MAX_CONSECUTIVE_TRANSIENT_FAILURES: u32 = 3;
/// Consecutive transient-failure budget, scoped to the credential it accrued
/// against. Held under one lock so the credential check, reset, and increment
/// are a single atomic step.
#[derive(Default)]
struct TransientBudget {
/// Credential the count belongs to. A different credential (e.g. after
/// re-login on this long-lived refresher) re-arms the budget so a fresh,
/// valid token never inherits a dead one's escalation.
key: Option<String>,
count: u32,
}
pub(crate) struct OidcRefresher {
auth: Arc<dyn AuthSnapshot>,
transient_budget: parking_lot::Mutex<TransientBudget>,
}
impl OidcRefresher {
pub(crate) fn new(auth: Arc<dyn AuthSnapshot>) -> Self {
Self {
auth,
transient_budget: parking_lot::Mutex::new(TransientBudget::default()),
}
}
/// Clear the transient-blip budget on refresh progress (a fresh token or an
/// adopted sibling token), so later blips start from a full budget.
fn note_refresh_progress(&self) {
*self.transient_budget.lock() = TransientBudget::default();
}
fn record_transient_failure(
&self,
message: String,
tried_key: Option<String>,
) -> RefreshOutcome {
let escalate = {
let mut budget = self.transient_budget.lock();
// Re-arm when the credential changes so a fresh token never inherits
// a prior credential's accrued blips.
if budget.key != tried_key {
budget.key = tried_key.clone();
budget.count = 0;
}
budget.count += 1;
let escalate = budget.count >= MAX_CONSECUTIVE_TRANSIENT_FAILURES;
// On escalation reset the count so the next TTL window gets the full
// budget (the verdict gates refresh() meanwhile). The key is left in
// place; a same-key retry resumes from zero, a new key re-arms.
if escalate {
budget.count = 0;
}
escalate
};
if escalate {
tracing::warn!(%message, "auth: escalating consecutive transient failures to permanent");
RefreshOutcome::permanent(RefreshTokenFailedReason::Other, tried_key)
} else {
RefreshOutcome::transient(message)
}
}
/// One-shot retry with disk's RT after `invalid_grant`.
///
/// If disk already has a valid (unexpired) AT with a different key,
/// adopt it directly, without consuming the disk's RT in another IdP
/// call. This prevents cascading `invalid_grant` when a sibling
/// already refreshed and wrote a valid token.
async fn retry_with_fresh_disk_token(
&self,
tried: &crate::auth::GrokAuth,
) -> Option<RefreshOutcome> {
let disk_now = self.auth.read_disk_auth()?;
// If disk has a valid AT that differs from what we tried,
// a sibling already refreshed. Adopt directly — no IdP call.
if !crate::auth::is_expired(&disk_now) && disk_now.key != tried.key {
crate::unified_log::info(
"oidc refresh: disk has valid AT, adopting instead of consuming RT",
None,
Some(serde_json::json!({
"disk_key_prefix": crate::auth::token_suffix(&disk_now.key),
"tried_key_prefix": crate::auth::token_suffix(&tried.key),
})),
);
self.note_refresh_progress();
return Some(RefreshOutcome::success(disk_now));
}
if disk_now.refresh_token.is_none()
|| disk_now.refresh_token.as_deref() == tried.refresh_token.as_deref()
{
return None;
}
crate::unified_log::info(
"oidc refresh retrying with disk token",
None,
Some(serde_json::json!({
"tried_rt_prefix": tried
.refresh_token
.as_deref()
.map(crate::auth::token_suffix),
"disk_rt_prefix": disk_now
.refresh_token
.as_deref()
.map(crate::auth::token_suffix),
})),
);
match crate::auth::oidc::oidc_token_exchange(&disk_now).await {
OidcRefreshResult::Success(new_auth) => {
self.note_refresh_progress();
Some(RefreshOutcome::Success(new_auth))
}
OidcRefreshResult::TerminalError { reason } => {
crate::unified_log::warn(
"oidc refresh disk retry exhausted",
None,
Some(serde_json::json!({ "reason": format!("{reason:?}") })),
);
Some(RefreshOutcome::permanent(
reason,
Some(disk_now.key.clone()),
))
}
OidcRefreshResult::Failed => {
Some(RefreshOutcome::transient("OIDC disk-retry refresh failed"))
}
}
}
}
#[async_trait::async_trait]
impl TokenRefresher for OidcRefresher {
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome {
crate::unified_log::debug(
"oidc refresh enter",
None,
Some(serde_json::json!({
"reason": format!("{reason:?}"),
"has_current": self.auth.current().is_some(),
"is_expired": self.auth.is_expired(),
})),
);
let disk_auth = self.auth.read_disk_auth();
// Short-circuit: if disk has a valid unexpired AT that differs
// from in-memory, a sibling refreshed between refresh_chain
// step 2 (disk check under lock) and here. Adopt it directly,
// no IdP call needed.
if let Some(ref d) = disk_auth
&& !crate::auth::is_expired(d)
&& self.auth.current().map(|a| a.key).as_deref() != Some(&d.key)
{
crate::unified_log::info(
"oidc refresh: sibling refreshed, adopting valid disk AT",
None,
Some(serde_json::json!({
"disk_key_prefix": crate::auth::token_suffix(&d.key),
})),
);
self.note_refresh_progress();
return RefreshOutcome::success(d.clone());
}
let auth = super::resolve_refresh_credential(self.auth.as_ref(), disk_auth, reason);
let Some(auth) = auth else {
crate::unified_log::warn(
"oidc refresh no token available",
None,
Some(serde_json::json!({ "reason": format!("{reason:?}") })),
);
return RefreshOutcome::transient("no token with refresh_token available");
};
crate::unified_log::info(
"oidc refresh attempting idp",
None,
Some(serde_json::json!({
"has_rt": auth.refresh_token.is_some(),
"issuer": auth.oidc_issuer,
"client_id": auth.oidc_client_id,
"expires_at": auth.expires_at.map(|e| e.to_rfc3339()),
})),
);
match crate::auth::oidc::oidc_token_exchange(&auth).await {
OidcRefreshResult::Success(new_auth) => {
self.note_refresh_progress();
RefreshOutcome::Success(new_auth)
}
OidcRefreshResult::TerminalError { reason } => {
// Sibling-rotation race: disk may hold a
// fresher RT than the one we tried. One-shot retry.
if reason == RefreshTokenFailedReason::RefreshTokenRejected
&& let Some(retry_outcome) = self.retry_with_fresh_disk_token(&auth).await
{
return retry_outcome;
}
RefreshOutcome::permanent(reason, Some(auth.key.clone()))
}
OidcRefreshResult::Failed => {
tracing::warn!(
refresh_reason = ?reason,
user_id = %auth.user_id,
has_refresh_token = auth.refresh_token.is_some(),
issuer = ?auth.oidc_issuer,
client_id = ?auth.oidc_client_id,
expires_at = ?auth.expires_at,
"auth: OIDC token refresh failed"
);
crate::unified_log::error(
"oidc refresh failed",
None,
Some(serde_json::json!({
"has_refresh_token": auth.refresh_token.is_some(),
"auth_mode": format!("{:?}", auth.auth_mode),
"issuer": auth.oidc_issuer,
"client_id": auth.oidc_client_id,
"expires_at": auth.expires_at.map(|e| e.to_rfc3339()),
})),
);
self.record_transient_failure(
"OIDC token refresh failed".into(),
Some(auth.key.clone()),
)
}
}
}
}
#[cfg(test)]
#[path = "oidc_refresher_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "auth_backend_contract_tests.rs"]
mod auth_backend_contract_tests;
crates/codegen/kigi-shell/src/auth/refresh/oidc_refresher_tests.rs
-1311
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/auth/storage.rs
+257 -11
View File
@@ -2,7 +2,184 @@ use std::fs::File;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use super::model::{API_KEY_SCOPE, AuthMode, AuthStore, GrokAuth, lookup_auth};
use super::model::{API_KEY_SCOPE, AuthMode, AuthStore, KimiAuth, lookup_auth};
// ── System-keyring storage for the Kimi Code OAuth session ─────────────
//
// PRD F1: the OAuth token set lives in the system keyring (service `kigi`,
// entry `oauth/kimi-code`); when the keyring is unavailable we fall back to
// the file mechanism below (`auth.json`, owner-only, atomic writes). The
// official Kimi client's keyring entries (service `kimi-code`) and `~/.kimi`
// files are never touched.
/// Keyring service name — deliberately distinct from the official client's
/// `kimi-code` service.
#[cfg(any(target_os = "macos", windows))]
pub(crate) const KEYRING_SERVICE: &str = "kigi";
/// Outcome of a keyring read for the session scope.
#[derive(Debug)]
pub(crate) enum KeyringRead {
/// Backend reachable and the entry exists.
Found(Box<KimiAuth>),
/// Backend reachable, no entry stored.
Missing,
/// Keyring disabled, unsupported on this platform, or the backend
/// errored — callers fall back to the file store.
Unavailable,
}
/// Whether keyring storage participates for the session credential.
///
/// Disabled when:
/// - the platform has no supported backend (non-macOS/Windows builds),
/// - `KIGI_DISABLE_KEYRING` is set to a truthy value,
/// - a non-default credential location is in use (`KIGI_SHARE_DIR` /
/// `KIGI_AUTH_PATH`): the keyring entry belongs to the default user
/// install; alternate profiles (and tests) stay file-scoped, and
/// - in unit-test builds, unless a test explicitly opted into the mock
/// keyring via [`enable_mock_keyring_for_test`].
pub(crate) fn keyring_enabled() -> bool {
#[cfg(test)]
{
// Thread-local so keyring-specific tests (which opt in via
// `enable_mock_keyring_for_test`) can't leak the toggle into
// concurrently running persistence tests on other threads.
TEST_KEYRING_ENABLED.with(|flag| flag.get())
}
#[cfg(not(test))]
{
#[cfg(not(any(target_os = "macos", windows)))]
{
false
}
#[cfg(any(target_os = "macos", windows))]
{
let disabled = std::env::var("KIGI_DISABLE_KEYRING")
.is_ok_and(|v| !matches!(v.trim(), "" | "0" | "false" | "off" | "no"));
!disabled
&& std::env::var_os("KIGI_SHARE_DIR").is_none()
&& std::env::var_os("KIGI_AUTH_PATH").is_none()
}
}
}
#[cfg(test)]
thread_local! {
static TEST_KEYRING_ENABLED: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
/// Route keyring calls through an in-memory mock store for this process,
/// enable [`keyring_enabled`] on this thread, and clear any entry left by an
/// earlier test. Tests using this must serialize on the `kigi_keyring` key:
/// the mock entry is shared process-wide.
#[cfg(all(test, any(target_os = "macos", windows)))]
pub(crate) fn enable_mock_keyring_for_test() {
keyring::set_default_credential_builder(keyring::mock::default_credential_builder());
TEST_KEYRING_ENABLED.with(|flag| flag.set(true));
if let Err(e) = keyring_delete_session() {
panic!("mock keyring cleanup failed: {e}");
}
}
/// Disable the test keyring again (paired with
/// [`enable_mock_keyring_for_test`] in an RAII guard or test teardown).
#[cfg(test)]
pub(crate) fn disable_mock_keyring_for_test() {
TEST_KEYRING_ENABLED.with(|flag| flag.set(false));
}
/// The process-wide keyring entry handle. Cached so every reader/writer talks
/// to the same credential object: real backends read live state from the OS
/// store on each call, and the test mock keeps its state on the entry itself.
#[cfg(any(target_os = "macos", windows))]
fn keyring_entry() -> Result<&'static keyring::Entry, keyring::Error> {
static ENTRY: std::sync::OnceLock<Result<keyring::Entry, keyring::Error>> =
std::sync::OnceLock::new();
match ENTRY.get_or_init(|| {
keyring::Entry::new(KEYRING_SERVICE, crate::auth::config::KIMI_CODE_OAUTH_SCOPE)
}) {
Ok(entry) => Ok(entry),
// `keyring::Error` is not `Clone`; surface a stable equivalent.
Err(e) => {
tracing::warn!(error = %e, "auth: keyring entry construction failed");
Err(keyring::Error::Invalid(
"keyring entry".into(),
e.to_string(),
))
}
}
}
/// Read the session credential from the system keyring.
#[cfg(any(target_os = "macos", windows))]
pub(crate) fn keyring_read_session() -> KeyringRead {
if !keyring_enabled() {
return KeyringRead::Unavailable;
}
let entry = match keyring_entry() {
Ok(entry) => entry,
Err(e) => {
tracing::warn!(error = %e, "auth: keyring entry unavailable, falling back to file");
return KeyringRead::Unavailable;
}
};
match entry.get_password() {
Ok(raw) => match serde_json::from_str::<KimiAuth>(&raw) {
Ok(auth) => KeyringRead::Found(Box::new(auth)),
Err(e) => {
tracing::warn!(error = %e, "auth: keyring entry is not valid JSON, ignoring");
KeyringRead::Missing
}
},
Err(keyring::Error::NoEntry) => KeyringRead::Missing,
Err(e) => {
tracing::warn!(error = %e, "auth: keyring read failed, falling back to file");
KeyringRead::Unavailable
}
}
}
#[cfg(not(any(target_os = "macos", windows)))]
pub(crate) fn keyring_read_session() -> KeyringRead {
KeyringRead::Unavailable
}
/// Write the session credential to the system keyring.
#[cfg(any(target_os = "macos", windows))]
pub(crate) fn keyring_write_session(auth: &KimiAuth) -> anyhow::Result<()> {
anyhow::ensure!(keyring_enabled(), "keyring storage disabled");
let payload = serde_json::to_string(auth)?;
keyring_entry()?.set_password(&payload)?;
tracing::info!("auth: session credential written to system keyring");
Ok(())
}
#[cfg(not(any(target_os = "macos", windows)))]
pub(crate) fn keyring_write_session(_auth: &KimiAuth) -> anyhow::Result<()> {
anyhow::bail!("keyring storage is not supported on this platform")
}
/// Delete the session credential from the system keyring (Ok when absent).
#[cfg(any(target_os = "macos", windows))]
pub(crate) fn keyring_delete_session() -> anyhow::Result<()> {
if !keyring_enabled() {
return Ok(());
}
match keyring_entry()?.delete_credential() {
Ok(()) => {
tracing::info!("auth: session credential removed from system keyring");
Ok(())
}
Err(keyring::Error::NoEntry) => Ok(()),
Err(e) => Err(e.into()),
}
}
#[cfg(not(any(target_os = "macos", windows)))]
pub(crate) fn keyring_delete_session() -> anyhow::Result<()> {
Ok(())
}
/// RAII guard for an exclusive advisory lock on `auth.json.lock`.
/// The lock is released when the inner `File` is dropped (closing the FD).
@@ -325,25 +502,23 @@ fn restore_prior_bytes(auth_file: &Path, bytes: &[u8]) -> std::io::Result<()> {
}
/// Read a single auth token from `auth.json` by scope key.
/// Falls back to the legacy `https://accounts.x.ai/sign-in` scope key
/// when the requested scope is not found (devbox auth.json migration).
pub fn read_token_by_scope(kigi_home: &Path, scope: &str) -> anyhow::Result<String> {
let path = kigi_home.join("auth.json");
let store =
read_auth_json(&path).map_err(|_| anyhow::anyhow!("Not logged in. Run `grok login`."))?;
read_auth_json(&path).map_err(|_| anyhow::anyhow!("Not logged in. Run `kigi login`."))?;
lookup_auth(&store, scope).map(|a| a.key).ok_or_else(|| {
anyhow::anyhow!("Your auth token is invalid. Run `grok login` to re-authenticate.")
anyhow::anyhow!("Your auth token is invalid. Run `kigi login` to re-authenticate.")
})
}
/// Read the API key from the `xai::api_key` scope in auth.json.
/// Read the API key from the `kigi::api_key` scope in auth.json.
pub fn read_api_key(kigi_home: &Path) -> Option<String> {
let path = kigi_home.join("auth.json");
let map = read_auth_json(&path).ok()?;
map.get(API_KEY_SCOPE).map(|a| a.key.clone())
}
/// Store a plain API key in auth.json under the `xai::api_key` scope.
/// Store a plain API key in auth.json under the `kigi::api_key` scope.
///
/// Uses the corrupt-recovery reader so a malformed auth.json (e.g. from a
/// previous crash) can be healed when the user sets an API key.
@@ -352,7 +527,7 @@ pub fn store_api_key(kigi_home: &Path, api_key: &str) -> std::io::Result<()> {
let mut map = read_auth_json_or_empty_recovering_corrupt(&path)?;
map.insert(
API_KEY_SCOPE.to_owned(),
GrokAuth {
KimiAuth {
key: api_key.to_owned(),
auth_mode: AuthMode::ApiKey,
..Default::default()
@@ -361,7 +536,7 @@ pub fn store_api_key(kigi_home: &Path, api_key: &str) -> std::io::Result<()> {
write_auth_json(&path, &map)
}
/// Remove the `xai::api_key` scope from auth.json.
/// Remove the `kigi::api_key` scope from auth.json.
pub fn clear_api_key(kigi_home: &Path) -> std::io::Result<()> {
let path = kigi_home.join("auth.json");
if let Ok(mut map) = read_auth_json(&path) {
@@ -383,7 +558,7 @@ mod write_fallback_tests {
let mut map = AuthStore::new();
map.insert(
API_KEY_SCOPE.to_owned(),
GrokAuth {
KimiAuth {
key: "secret-key".to_owned(),
auth_mode: AuthMode::ApiKey,
..Default::default()
@@ -481,7 +656,7 @@ mod write_fallback_tests {
let mut replacement = AuthStore::new();
replacement.insert(
API_KEY_SCOPE.to_owned(),
GrokAuth {
KimiAuth {
key: "replacement-key".to_owned(),
auth_mode: AuthMode::ApiKey,
..Default::default()
@@ -510,3 +685,74 @@ mod write_fallback_tests {
assert_eq!(mode & 0o777, 0o600, "restored file must stay 0o600");
}
}
#[cfg(all(test, any(target_os = "macos", windows)))]
mod keyring_tests {
use super::*;
use chrono::Utc;
/// RAII teardown so a panicking test doesn't leave the process-global
/// test-keyring toggle enabled for later tests.
struct MockKeyringGuard;
impl MockKeyringGuard {
fn enable() -> Self {
enable_mock_keyring_for_test();
Self
}
}
impl Drop for MockKeyringGuard {
fn drop(&mut self) {
disable_mock_keyring_for_test();
}
}
fn session_auth(key: &str, rt: &str) -> KimiAuth {
KimiAuth {
key: key.into(),
refresh_token: Some(rt.into()),
expires_at: Some(Utc::now() + chrono::Duration::seconds(3600)),
expires_in: Some(3600),
scope: Some("kimi-code".into()),
token_type: Some("bearer".into()),
..KimiAuth::test_default()
}
}
#[test]
#[serial_test::serial(kigi_keyring)]
fn keyring_session_roundtrip() {
let _guard = MockKeyringGuard::enable();
// Fresh mock store: nothing there yet.
assert!(matches!(keyring_read_session(), KeyringRead::Missing));
keyring_write_session(&session_auth("at-1", "rt-1")).unwrap();
let KeyringRead::Found(read) = keyring_read_session() else {
panic!("expected Found after write");
};
assert_eq!(read.key, "at-1");
assert_eq!(read.refresh_token.as_deref(), Some("rt-1"));
assert_eq!(read.expires_in, Some(3600));
// Overwrite rotates in place.
keyring_write_session(&session_auth("at-2", "rt-2")).unwrap();
let KeyringRead::Found(read) = keyring_read_session() else {
panic!("expected Found after rotate");
};
assert_eq!(read.key, "at-2");
// Delete is idempotent.
keyring_delete_session().unwrap();
assert!(matches!(keyring_read_session(), KeyringRead::Missing));
keyring_delete_session().unwrap();
}
#[test]
#[serial_test::serial(kigi_keyring)]
fn keyring_disabled_reads_unavailable() {
disable_mock_keyring_for_test();
assert!(matches!(keyring_read_session(), KeyringRead::Unavailable));
assert!(keyring_write_session(&session_auth("a", "r")).is_err());
// Delete when disabled is a no-op success (logout stays best-effort).
keyring_delete_session().unwrap();
}
}
crates/codegen/kigi-shell/src/auth/token_type.rs
+36 -20
View File
@@ -1,17 +1,13 @@
use crate::auth::model::{AuthMode, GrokAuth};
use crate::auth::model::{AuthMode, KimiAuth};
/// What kind of bearer is loaded right now. Dispatch key for
/// `auth()`, `unauthorized_recovery()`, and proactive refresh.
///
/// Not a session classifier — use `is_session_based_method` for that.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum TokenType {
/// OIDC/OAuth2 session with a refresh_token available.
OidcSession,
/// Legacy web-login session or OIDC without a refresh_token.
LegacySession,
/// External auth binary provides tokens.
ExternalBinary,
/// Kimi Code OAuth session with a refresh_token available.
OAuthSession,
/// OAuth session without a refresh_token (cannot be silently renewed).
SessionNoRefresh,
/// Plain API key (no refresh possible).
ApiKey,
/// No credentials loaded.
@@ -20,36 +16,56 @@ pub(crate) enum TokenType {
impl TokenType {
/// Classify the loaded credential (pure; no manager state).
pub(crate) fn from_auth(auth: Option<&GrokAuth>) -> Self {
pub(crate) fn from_auth(auth: Option<&KimiAuth>) -> Self {
match auth {
None => Self::None,
// Oidc without a refresh_token degrades to the unrefreshable LegacySession shape.
Some(a) => match a.auth_mode {
AuthMode::Oidc if a.refresh_token.is_some() => Self::OidcSession,
AuthMode::Oidc | AuthMode::WebLogin => Self::LegacySession,
AuthMode::External => Self::ExternalBinary,
AuthMode::OAuth if a.refresh_token.is_some() => Self::OAuthSession,
AuthMode::OAuth => Self::SessionNoRefresh,
AuthMode::ApiKey => Self::ApiKey,
},
}
}
/// `true` for types that can be silently refreshed (OIDC, external binary).
/// `true` for types that can be silently refreshed.
pub(crate) fn is_refreshable(self) -> bool {
matches!(self, Self::OidcSession | Self::ExternalBinary)
matches!(self, Self::OAuthSession)
}
}
#[cfg(test)]
mod tests {
//! Per-variant matrix for `is_refreshable`.
//! Per-variant matrix for `is_refreshable` and classification.
use super::*;
#[test]
fn is_refreshable_matrix() {
assert!(TokenType::OidcSession.is_refreshable());
assert!(TokenType::ExternalBinary.is_refreshable());
assert!(!TokenType::LegacySession.is_refreshable());
assert!(TokenType::OAuthSession.is_refreshable());
assert!(!TokenType::SessionNoRefresh.is_refreshable());
assert!(!TokenType::ApiKey.is_refreshable());
assert!(!TokenType::None.is_refreshable());
}
#[test]
fn from_auth_classifies_by_mode_and_refresh_token() {
assert_eq!(TokenType::from_auth(None), TokenType::None);
let with_rt = KimiAuth {
refresh_token: Some("rt".into()),
..KimiAuth::test_default()
};
assert_eq!(
TokenType::from_auth(Some(&with_rt)),
TokenType::OAuthSession
);
let no_rt = KimiAuth::test_default();
assert_eq!(
TokenType::from_auth(Some(&no_rt)),
TokenType::SessionNoRefresh
);
let api = KimiAuth {
auth_mode: AuthMode::ApiKey,
..KimiAuth::test_default()
};
assert_eq!(TokenType::from_auth(Some(&api)), TokenType::ApiKey);
}
}
crates/codegen/kigi-shell/src/cli_models.rs
+16 -59
View File
@@ -22,28 +22,24 @@ impl AuthStatus {
/// Banner status: env key → session → BYOK → deployment → none.
///
/// Differs from sampling (`resolve_credentials`: BYOK → session → env) so a
/// logged-in user sees the login host. BYOK uses
/// [`crate::agent::auth_method::should_advertise_xai_api_key`] so
/// `disable_api_key_auth` is honored.
/// logged-in user sees the login host.
pub fn resolve(agent_config: &AgentConfig) -> Self {
if crate::agent::auth_method::has_xai_api_key_env() {
return Self::ApiKey;
}
if agent_config.create_auth_manager().current().is_some() {
let origin = &agent_config.grok_com_config.grok_ws_origin;
let origin = kigi_env::oauth_host();
let host = origin
.strip_prefix("https://")
.or_else(|| origin.strip_prefix("http://"))
.unwrap_or(origin);
.unwrap_or(&origin);
return Self::LoggedIn(host.to_owned());
}
let models = crate::agent::config::resolve_model_list(agent_config, None);
if crate::agent::auth_method::should_advertise_xai_api_key(
agent_config.grok_com_config.api_key_auth_disabled(),
models.values(),
) && let Some(name) = models
.iter()
.find_map(|(name, entry)| entry.has_own_credentials().then(|| name.clone()))
if crate::agent::auth_method::should_advertise_xai_api_key(models.values())
&& let Some(name) = models
.iter()
.find_map(|(name, entry)| entry.has_own_credentials().then(|| name.clone()))
{
return Self::ModelCredentials(name);
}
@@ -94,7 +90,7 @@ mod tests {
use super::*;
use crate::agent::auth_method::{LEGACY_XAI_API_KEY_ENV_VAR, XAI_API_KEY_ENV_VAR};
use crate::agent::config::Config;
use crate::auth::{AuthMode, GrokAuth};
use crate::auth::{AuthMode, KimiAuth};
use kigi_test_support::EnvGuard;
use serial_test::serial;
@@ -155,17 +151,17 @@ mod tests {
#[serial]
fn resolve_oauth_session() {
let (_dir, _g) = isolate_auth_sources();
let token = GrokAuth {
let token = KimiAuth {
key: "session-token".into(),
auth_mode: AuthMode::WebLogin,
..GrokAuth::test_default()
auth_mode: AuthMode::OAuth,
..KimiAuth::test_default()
};
let json = serde_json::to_string(&token).unwrap();
let _auth = EnvGuard::set("KIGI_AUTH", &json);
assert_eq!(
AuthStatus::resolve(&Config::default()),
AuthStatus::LoggedIn("grok.com".to_owned())
AuthStatus::LoggedIn("auth.kimi.com".to_owned())
);
}
@@ -247,10 +243,10 @@ mod tests {
#[serial]
fn resolve_priority_session_over_byok_and_deployment() {
let (_dir, _g) = isolate_auth_sources();
let token = GrokAuth {
let token = KimiAuth {
key: "session-token".into(),
auth_mode: AuthMode::WebLogin,
..GrokAuth::test_default()
auth_mode: AuthMode::OAuth,
..KimiAuth::test_default()
};
let json = serde_json::to_string(&token).unwrap();
let _auth = EnvGuard::set("KIGI_AUTH", &json);
@@ -259,7 +255,7 @@ mod tests {
let cfg = config_from_toml(&byok_and_deployment_toml(dm));
assert_eq!(
AuthStatus::resolve(&cfg),
AuthStatus::LoggedIn("grok.com".to_owned())
AuthStatus::LoggedIn("auth.kimi.com".to_owned())
);
}
@@ -275,45 +271,6 @@ mod tests {
);
}
#[test]
#[serial]
fn resolve_disable_api_key_auth_suppresses_byok_banner() {
let (_dir, _g) = isolate_auth_sources();
let dm = crate::models::default_model();
let cfg = config_from_toml(&format!(
r#"
[grok_com_config]
disable_api_key_auth = true
[model."{dm}"]
model = "{dm}"
api_key = "sk-byok"
"#
));
assert_eq!(AuthStatus::resolve(&cfg), AuthStatus::NotAuthenticated);
}
#[test]
#[serial]
fn resolve_disable_api_key_auth_falls_through_to_deployment() {
let (_dir, _g) = isolate_auth_sources();
let dm = crate::models::default_model();
let cfg = config_from_toml(&format!(
r#"
[grok_com_config]
disable_api_key_auth = true
[endpoints]
deployment_key = "deploy-key"
[model."{dm}"]
model = "{dm}"
api_key = "sk-byok"
"#
));
assert_eq!(AuthStatus::resolve(&cfg), AuthStatus::DeploymentKey);
}
#[test]
#[serial]
fn resolve_model_credentials_uses_first_catalog_key() {
crates/codegen/kigi-shell/src/config/reloader.rs
+7 -7
View File
@@ -7,7 +7,7 @@ use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info};
use crate::auth::{GrokAuth, read_auth_json};
use crate::auth::{KimiAuth, read_auth_json};
use super::watcher::ConfigChangeEvent;
@@ -15,7 +15,7 @@ use super::watcher::ConfigChangeEvent;
#[derive(Debug)]
pub enum ConfigUpdate {
/// New auth credentials from disk.
Auth(Box<GrokAuth>),
Auth(Box<KimiAuth>),
/// Auth scope was removed (user logged out).
AuthCleared,
/// A **broadcast** MCP reload — applies to every active session
@@ -535,14 +535,14 @@ fn extract_ui_fields(config: &toml::Value) -> (Option<String>, bool, Option<Stri
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::GrokAuth;
use crate::auth::KimiAuth;
use std::collections::BTreeMap;
fn make_auth(key: &str) -> GrokAuth {
GrokAuth {
fn make_auth(key: &str) -> KimiAuth {
KimiAuth {
key: key.to_string(),
email: Some("test@test.com".to_string()),
..GrokAuth::test_default()
..KimiAuth::test_default()
}
}
@@ -604,7 +604,7 @@ mod tests {
reloader.reload_auth().unwrap();
let update = rx.try_recv().expect("should send Auth update");
assert!(
matches!(update, ConfigUpdate::Auth(a) if a.key == "new-key"), // a is Box<GrokAuth>, Deref coercion
matches!(update, ConfigUpdate::Auth(a) if a.key == "new-key"), // a is Box<KimiAuth>, Deref coercion
"should contain new key"
);
}
crates/codegen/kigi-shell/src/extensions/auth.rs
+9 -64
View File
@@ -21,7 +21,6 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
"x.ai/auth/get_url" => handle_get_url(agent).await,
"x.ai/auth/logout" => handle_logout(agent, args).await,
"x.ai/auth/info" => handle_info(agent),
"x.ai/auth/check_subscription" => handle_check_subscription(agent).await,
_ => Err(acp::Error::method_not_found()),
}
}
@@ -111,7 +110,7 @@ async fn handle_get_url(agent: &MvpAgent) -> ExtResult {
to_raw_response(&serde_json::json!({
"auth_url": auth_url,
// `external_provider` kept for older clients; `mode` is authoritative.
"external_provider": mode.is_some_and(|m| m.is_external_provider()),
"external_provider": false,
"mode": mode.map(|m| m.as_wire_str()),
}))
}
@@ -141,43 +140,16 @@ async fn handle_logout(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
}))
}
/// Single-shot subscription re-check (retry button on paywall screen).
///
/// Calls `retry_subscription_check()`, then returns the updated auth
/// response with gate info so the pager can refresh the gate state.
async fn handle_check_subscription(agent: &MvpAgent) -> ExtResult {
agent.retry_subscription_check().await;
let response = agent.auth_response_with_meta();
to_raw_response(&serde_json::json!({
"authenticated": response.meta.is_some(),
"meta": response.meta,
}))
}
/// Returns current auth method ID, user profile fields, and team/principal
/// metadata.
/// Returns current auth method ID and the account fields the Kimi flow
/// exposes (email/user id are empty until a later feature surfaces them).
fn handle_info(agent: &MvpAgent) -> ExtResult {
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct AuthInfoResponse {
method_id: Option<String>,
email: Option<String>,
first_name: Option<String>,
last_name: Option<String>,
/// `grok-asset://` URL resolved by the Electron protocol handler,
/// or a full `http(s)://` URL passed through unchanged.
profile_image_url: Option<String>,
team_id: Option<String>,
team_name: Option<String>,
team_role: Option<String>,
organization_id: Option<String>,
organization_name: Option<String>,
organization_role: Option<String>,
principal_type: Option<String>,
principal_id: Option<String>,
user_blocked_reason: Option<String>,
team_blocked_reasons: Vec<String>,
coding_data_retention_opt_out: bool,
user_id: Option<String>,
auth_mode: Option<String>,
}
let method_id = agent
@@ -186,40 +158,13 @@ fn handle_info(agent: &MvpAgent) -> ExtResult {
.as_ref()
.map(|m| m.0.to_string());
let auth = agent.auth_manager.current();
let raw_asset_id = auth.as_ref().and_then(|a| a.profile_image_asset_id.clone());
// Return a grok-asset:// URL that the Electron renderer resolves at
// display time via a custom protocol handler. The handler proxies
// through cli-chat-proxy's /asset endpoint; Electron's HTTP cache
// handles reuse. No disk-cache or network call needed here.
let profile_image_url = match raw_asset_id.as_deref().filter(|k| !k.is_empty()) {
Some(key) if key.starts_with("http://") || key.starts_with("https://") => {
Some(key.to_owned())
}
Some(key) => Some(format!("grok-asset:///{key}")),
None => None,
};
to_raw_response(&AuthInfoResponse {
method_id,
email: auth.as_ref().and_then(|a| a.email.clone()),
first_name: auth.as_ref().and_then(|a| a.first_name.clone()),
last_name: auth.as_ref().and_then(|a| a.last_name.clone()),
profile_image_url,
team_id: auth.as_ref().and_then(|a| a.team_id.clone()),
team_name: auth.as_ref().and_then(|a| a.team_name.clone()),
team_role: auth.as_ref().and_then(|a| a.team_role.clone()),
organization_id: auth.as_ref().and_then(|a| a.organization_id.clone()),
organization_name: auth.as_ref().and_then(|a| a.organization_name.clone()),
organization_role: auth.as_ref().and_then(|a| a.organization_role.clone()),
principal_type: auth.as_ref().and_then(|a| a.principal_type.clone()),
principal_id: auth.as_ref().and_then(|a| a.principal_id.clone()),
user_blocked_reason: auth.as_ref().and_then(|a| a.user_blocked_reason.clone()),
team_blocked_reasons: auth
user_id: auth
.as_ref()
.map(|a| a.team_blocked_reasons.clone())
.unwrap_or_default(),
coding_data_retention_opt_out: auth
.as_ref()
.is_some_and(|a| a.coding_data_retention_opt_out),
.map(|a| a.user_id.clone())
.filter(|id| !id.is_empty()),
auth_mode: auth.as_ref().map(|a| format!("{:?}", a.auth_mode)),
})
}
crates/codegen/kigi-shell/src/extensions/auth_gate.rs
+4 -4
View File
@@ -1,17 +1,17 @@
use agent_client_protocol as acp;
use crate::auth::{AuthManager, GrokAuth};
use crate::auth::{AuthManager, KimiAuth};
/// Require xAI auth from a sync context, accepting tokens in the client-side buffer window.
/// Require a Kimi Code session from a sync context, accepting tokens in the client-side buffer window.
pub(crate) fn require_xai_auth(
auth_manager: &AuthManager,
missing_message: &'static str,
non_xai_message: &'static str,
) -> Result<GrokAuth, acp::Error> {
) -> Result<KimiAuth, acp::Error> {
let auth = auth_manager
.current_or_expired()
.ok_or_else(|| acp::Error::auth_required().data(missing_message))?;
if !auth.is_xai_auth() {
if !auth.is_session_auth() {
return Err(acp::Error::auth_required().data(non_xai_message));
}
Ok(auth)
crates/codegen/kigi-shell/src/extensions/billing.rs
-8
View File
@@ -213,10 +213,6 @@ async fn handle_get_billing(agent: &MvpAgent) -> ExtResult {
let credits_resp = crate::http::shared_client()
.get(&credits_url)
.header("Authorization", format!("Bearer {}", auth.key))
.header(
"X-XAI-Token-Auth",
crate::auth::GrokComConfig::default().token_header,
)
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
@@ -304,10 +300,6 @@ async fn handle_get_auto_topup_rule(agent: &MvpAgent) -> ExtResult {
let response = crate::http::shared_client()
.get(&url)
.header("Authorization", format!("Bearer {}", auth.key))
.header(
"X-XAI-Token-Auth",
crate::auth::GrokComConfig::default().token_header,
)
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crates/codegen/kigi-shell/src/extensions/bundle.rs
+7 -21
View File
@@ -484,37 +484,23 @@ mod tests {
.insert("review".to_string(), "# Review skill\n".to_string());
bundle
}
fn test_auth() -> crate::auth::GrokAuth {
crate::auth::GrokAuth {
fn test_auth() -> crate::auth::KimiAuth {
crate::auth::KimiAuth {
key: "token".to_string(),
auth_mode: crate::auth::AuthMode::Oidc,
auth_mode: crate::auth::AuthMode::OAuth,
create_time: chrono::Utc::now(),
user_id: "user-1".to_string(),
email: Some("test@example.com".to_string()),
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: None,
team_name: None,
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token: None,
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
oidc_issuer: None,
oidc_client_id: None,
expires_in: Some(3600),
scope: None,
token_type: None,
}
}
fn test_auth_manager() -> Arc<crate::auth::AuthManager> {
let dir = tempfile::tempdir().unwrap();
let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::GrokComConfig::default());
let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::KimiCodeConfig::default());
mgr.hot_swap(test_auth());
std::mem::forget(dir);
Arc::new(mgr)
crates/codegen/kigi-shell/src/extensions/mod.rs
-1
View File
@@ -17,7 +17,6 @@ pub mod memory;
pub mod notification;
pub mod plugins;
pub mod pr;
pub mod privacy;
pub mod prompt_history;
pub mod prompt_meta;
pub mod recap;
crates/codegen/kigi-shell/src/extensions/privacy.rs
-91
View File
@@ -1,91 +0,0 @@
//! `x.ai/privacy/setCodingDataRetention` extension handler.
//!
//! PUTs the new opt-out flag to cli-chat-proxy and updates local auth state
//! to match. The local update is fire-and-forget (best-effort cache refresh).
use agent_client_protocol as acp;
use serde::Deserialize;
use super::{ExtResult, parse_params, to_raw_response};
use crate::agent::MvpAgent;
#[tracing::instrument(skip_all, fields(method = %args.method))]
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
match args.method.as_ref() {
"x.ai/privacy/setCodingDataRetention" => handle_set(agent, args).await,
_ => Err(acp::Error::method_not_found()),
}
}
async fn handle_set(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct Params {
coding_data_retention_opt_out: bool,
}
let params: Params = parse_params(args)?;
let auth = agent.auth_manager.auth().await.map_err(|e| {
tracing::warn!(error = %e, "privacy: auth resolution failed");
acp::Error::auth_required()
.data("Authentication required. Run `grok login` to re-authenticate.")
})?;
let proxy_url = agent.cfg.borrow().endpoints.proxy_url();
let url = format!("{proxy_url}/privacy/coding-data-retention");
let token_header = agent.auth_manager.grok_com_config().token_header.clone();
let body = serde_json::json!({
"codingDataRetentionOptOut": params.coding_data_retention_opt_out,
});
let provider: std::sync::Arc<dyn kigi_auth::AuthCredentialProvider> = std::sync::Arc::new(
crate::auth::credential_provider::ShellAuthCredentialProvider::new(
agent.auth_manager.clone(),
None,
None,
),
);
let client = crate::http::with_auth_retry(crate::http::shared_client(), provider);
let resp = client
.put(&url)
.header("X-XAI-Token-Auth", &token_header)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
)
.json(&body)
.send()
.await
.map_err(|e| acp::Error::internal_error().data(format!("HTTP request failed: {e}")))?;
if !resp.status().is_success() {
let status = resp.status().as_u16();
let body = resp.text().await.unwrap_or_default();
tracing::warn!(status, "setCodingDataRetention request failed");
let friendly = serde_json::from_str::<serde_json::Value>(&body)
.ok()
.and_then(|v| {
v.get("error")
.or_else(|| v.get("message"))
.and_then(|e| e.as_str().map(String::from))
})
.unwrap_or_else(|| format!("server returned HTTP {status}"));
return Err(acp::Error::internal_error().data(friendly));
}
// Update local auth state to reflect the change.
// Use save_without_enrichment to avoid a race: update() spawns a
// background GET /user enrichment that may read stale ACL state
// and overwrite the opt-out flag back to its previous value.
let mut updated = auth.clone();
updated.coding_data_retention_opt_out = params.coding_data_retention_opt_out;
let _ = agent.auth_manager.save_without_enrichment(updated).await;
to_raw_response(&serde_json::json!({
"codingDataRetentionOptOut": params.coding_data_retention_opt_out,
}))
}
crates/codegen/kigi-shell/src/extensions/session_admin.rs
+3 -15
View File
@@ -111,10 +111,7 @@ async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR
// Send a SessionSummaryGenerated notification so the TUI updates its title
notify_session_title(agent, session_id, &req.title).await;
if agent.is_writeback_storage()
&& let Some(auth) = agent.current_auth()
&& !auth.is_zdr_team()
{
if agent.is_writeback_storage() && agent.current_auth().is_some() {
use crate::remote::client::BackendClient;
use crate::session::export::ExportedMetadata;
@@ -133,15 +130,7 @@ async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR
// Hook 2: update session replica with summary (fire-and-forget)
if let Some(client) = agent.session_registry_client() {
let sid = req.session_id.to_string();
let title = if agent
.auth_manager
.current_or_expired()
.is_some_and(|a| a.is_zdr_team())
{
None
} else {
Some(req.title.clone())
};
let title = Some(req.title.clone());
tokio::spawn(async move {
let update = crate::agent::session_registry_client::UpdateRequest {
summary: title,
@@ -248,8 +237,7 @@ async fn handle_session_delete(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR
// For writeback storage (non-ZDR): remote delete is authoritative for
// the cloud history and runs first; on failure no local bits are
// touched so the pager does not remove the row or toast success.
let needs_remote =
agent.is_writeback_storage() && agent.current_auth().is_some_and(|a| !a.is_zdr_team());
let needs_remote = agent.is_writeback_storage() && agent.current_auth().is_some();
// Shared delete: remote-first, then local disk + FTS eviction.
// Mirrored by the `grok sessions delete <id>` CLI path.
crates/codegen/kigi-shell/src/extensions/share.rs
+9 -17
View File
@@ -45,13 +45,6 @@ async fn handle_share_session(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRe
);
}
// Only block for ZDR teams (hard data-retention policy), not for
// coding-data-retention opt-out — sharing is user-initiated.
if auth.is_zdr_team() {
return Err(acp::Error::invalid_params()
.data("Session sharing is disabled for your team's data retention policy"));
}
// Find session info by searching through summaries
let summaries = list_summaries(None).await.map_err(|e| {
acp::Error::internal_error().data(format!("Failed to list sessions: {}", e))
@@ -94,7 +87,7 @@ async fn handle_share_session(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRe
fn require_xai_auth_for_share(
auth_manager: &crate::auth::AuthManager,
) -> Result<crate::auth::GrokAuth, acp::Error> {
) -> Result<crate::auth::KimiAuth, acp::Error> {
super::auth_gate::require_xai_auth(
auth_manager,
"Authentication required to share session",
@@ -105,8 +98,8 @@ fn require_xai_auth_for_share(
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::GrokComConfig;
use crate::auth::{AuthMode, GrokAuth};
use crate::auth::KimiCodeConfig;
use crate::auth::{AuthMode, KimiAuth};
use chrono::{Duration, Utc};
use std::sync::Arc;
use tempfile::tempdir;
@@ -117,7 +110,7 @@ mod tests {
let dir = tempdir().expect("tempdir for share auth test");
let mgr = Arc::new(crate::auth::AuthManager::new(
dir.path(),
GrokComConfig::default(),
KimiCodeConfig::default(),
));
let expires_at = Utc::now() + ttl;
@@ -126,9 +119,8 @@ mod tests {
// Only OIDC tokens against https://auth.x.ai (or the local-dev equivalent)
// return true from is_xai_auth(). This is required for the share tests to
// exercise the happy path through require_xai_auth_for_share.
let auth = GrokAuth {
auth_mode: AuthMode::Oidc,
oidc_issuer: Some("https://auth.x.ai".to_string()),
let auth = KimiAuth {
auth_mode: AuthMode::OAuth,
key: "test-key".into(),
expires_at: Some(expires_at),
create_time: Utc::now() - Duration::hours(1),
@@ -168,7 +160,7 @@ mod tests {
let dir = tempdir().expect("tempdir");
let mgr = Arc::new(crate::auth::AuthManager::new(
dir.path(),
GrokComConfig::default(),
KimiCodeConfig::default(),
));
assert!(require_xai_auth_for_share(&mgr).is_err());
}
@@ -178,12 +170,12 @@ mod tests {
let dir = tempdir().expect("tempdir");
let mgr = Arc::new(crate::auth::AuthManager::new(
dir.path(),
GrokComConfig::default(),
KimiCodeConfig::default(),
));
// API key is the simplest non-xAI credential (External and enterprise OIDC
// are also rejected the same way).
let non_xai = GrokAuth {
let non_xai = KimiAuth {
auth_mode: AuthMode::ApiKey,
key: "xai-test-key".into(),
create_time: Utc::now(),
crates/codegen/kigi-shell/src/inspect/mod.rs
-62
View File
@@ -17,7 +17,6 @@ use std::path::{Path, PathBuf};
use serde::Serialize;
use crate::auth::ForceLoginTeam;
use kigi_tools::types::config_source::ConfigSource;
use kigi_tools::util::truncate::estimate_tokens;
@@ -63,7 +62,6 @@ pub struct InspectReport {
pub project_trusted: bool,
pub project_instructions: Vec<InstructionFile>,
pub permissions: PermissionsReport,
pub login_policy: LoginPolicyReport,
pub hooks: Vec<HookEntry>,
pub skills: Vec<SkillEntry>,
pub agents: Vec<AgentEntry>,
@@ -139,20 +137,6 @@ pub struct SkippedRule {
pub reason: String,
}
/// Enterprise login-hardening policy resolved from `[grok_com_config]`
/// (TOML + env). Surfaced so admins can verify the deployment loaded it.
/// The team pin is admin policy, not a secret, so it is shown verbatim.
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LoginPolicyReport {
/// Raw `disable_api_key_auth` knob (env `KIGI_DISABLE_API_KEY_AUTH`).
pub disable_api_key_auth: Option<bool>,
/// Configured team pin: single string, list, or null when unset.
pub force_login_team_uuid: Option<ForceLoginTeam>,
/// Resolved verdict — true when either knob forces first-party login.
pub api_key_auth_disabled: bool,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct HookEntry {
@@ -395,7 +379,6 @@ async fn build_report(cwd: &Path) -> InspectReport {
project_trusted,
project_instructions: instructions,
permissions,
login_policy: login_policy_report(parsed_config.as_ref()),
hooks,
skills,
agents,
@@ -615,20 +598,6 @@ async fn list_permissions(cwd: &Path) -> PermissionsReport {
}
}
/// Resolves the enterprise login-hardening knobs from the merged config
/// (`[grok_com_config]`, the `[auth]` alias, and env overrides) so admins can
/// confirm the deployment's auth policy actually loaded.
fn login_policy_report(config: Option<&crate::agent::config::Config>) -> LoginPolicyReport {
let grok_com_config = config
.map(|c| c.grok_com_config.clone())
.unwrap_or_default();
LoginPolicyReport {
api_key_auth_disabled: grok_com_config.api_key_auth_disabled(),
disable_api_key_auth: grok_com_config.disable_api_key_auth,
force_login_team_uuid: grok_com_config.force_login_team_uuid,
}
}
/// Discovers hooks with every vendor enabled so compatibility can be annotated later.
fn list_hooks(
git_root: Option<&Path>,
@@ -1174,19 +1143,6 @@ fn print_columns<T>(
}
}
/// Render the team pin for the human view: single value, comma-joined list,
/// or an explicit empty-list marker (which fails closed at login).
fn format_force_login_team(team: &Option<ForceLoginTeam>) -> String {
match team {
None => "(none)".to_string(),
Some(ForceLoginTeam::Single(s)) => s.clone(),
Some(ForceLoginTeam::AnyOf(list)) if list.is_empty() => {
"(empty -- fail closed)".to_string()
}
Some(ForceLoginTeam::AnyOf(list)) => list.join(", "),
}
}
/// Human label for an enforced setting. Uses product vocabulary, not the
/// internal field names (no `ui.yolo` / `--yolo` / `permission_mode`).
fn enforced_label(p: &EnforcedPolicy) -> String {
@@ -1342,24 +1298,6 @@ fn print_human(r: &InspectReport) {
}
}
println!();
println!(" Login Policy");
println!(
" {TREE} disable_api_key_auth: {}",
match r.login_policy.disable_api_key_auth {
Some(v) => v.to_string(),
None => "(unset)".to_string(),
}
);
println!(
" {TREE} force_login_team_uuid: {}",
format_force_login_team(&r.login_policy.force_login_team_uuid)
);
println!(
" {TREE} api_key_auth_disabled: {}",
r.login_policy.api_key_auth_disabled
);
print_columns(
"Skills",
&r.skills,
crates/codegen/kigi-shell/src/leader/server.rs
+6 -16
View File
@@ -210,24 +210,14 @@ impl AuthProvider for LeaderAuthProvider {
AuthCredential::bearer(token)
}
/// Owner identity from the leader's `AuthManager`, surfaced on the auth
/// provider instead of a separate auth.json
/// read. Mirrors the in-process path (`mvp_agent`): prefer `GrokAuth.team_id`
/// (what shell telemetry/snapshot use) mapped onto a `"Team"` principal so
/// team attribution is derived; otherwise pass principal fields through.
/// `None` when no credential is available (identity resolution never blocks).
/// provider instead of a separate auth.json read. The Kimi credential
/// carries no principal metadata; only the (possibly empty) user id.
fn identity(&self) -> Option<AuthIdentity> {
let a = self.auth_manager.current_or_expired()?;
Some(match a.team_id.filter(|t| !t.is_empty()) {
Some(team) => AuthIdentity {
user_id: a.user_id,
principal_type: Some("Team".to_string()),
principal_id: Some(team),
},
None => AuthIdentity {
user_id: a.user_id,
principal_type: a.principal_type,
principal_id: a.principal_id,
},
Some(AuthIdentity {
user_id: a.user_id,
principal_type: None,
principal_id: None,
})
}
}
crates/codegen/kigi-shell/src/lib.rs
-1
View File
@@ -34,7 +34,6 @@ pub mod session;
pub mod terminal;
#[cfg(test)]
pub(crate) mod test_support;
pub mod tier;
pub mod tools;
pub mod trace_classifier;
pub mod util;
crates/codegen/kigi-shell/src/managed_config.rs
+35 -165
View File
@@ -3,7 +3,7 @@
mod response;
use crate::auth::GrokAuth;
use crate::auth::KimiAuth;
pub use response::ManagedConfigError;
use response::{ApplyOutcome, ManagedConfigResponse, ManagedConfigSource, verify_signed_envelope};
@@ -76,35 +76,10 @@ fn remove_managed_path(path: &std::path::Path) -> std::io::Result<bool> {
}
}
/// A team principal is eligible to fetch only if non-expired (an expired token
/// would just 401).
fn eligible_team_principal(auth: GrokAuth) -> Option<GrokAuth> {
(auth.is_team_principal() && !crate::auth::is_expired(&auth)).then_some(auth)
}
/// The eligible team principal in `auth.json`, or `None`. Single-team: managed
/// config is a grok.com feature with one grok.com auth.
fn read_active_team_auth() -> Option<GrokAuth> {
let home = crate::util::kigi_home::kigi_home();
let store = crate::auth::read_auth_json(&home.join("auth.json")).ok()?;
let team = store.values().find(|a| a.is_team_principal())?.clone();
eligible_team_principal(team)
}
/// Team principals were an xAI concept; the Kimi Code auth model has none,
/// so no team credential can ever serve managed config.
pub(crate) fn has_active_team_auth() -> bool {
read_active_team_auth().is_some()
}
/// Whether any team principal is signed in, **ignoring expiry** (a cold-start
/// expired token is not a logout). `Err` = `auth.json` unreadable: callers must
/// NOT treat that as a logout — it would wipe enforced policy on a read blip.
fn team_principal_signed_in() -> std::io::Result<bool> {
let home = crate::util::kigi_home::kigi_home();
match crate::auth::read_auth_json(&home.join("auth.json")) {
Ok(store) => Ok(store.values().any(|a| a.is_team_principal())),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
Err(e) => Err(e),
}
false
}
/// Clear the synced files when no principal could own them: no deployment key
@@ -115,14 +90,6 @@ pub fn clear_orphan() {
if resolve_deployment_key().is_some() {
return;
}
match team_principal_signed_in() {
Ok(true) => return,
Ok(false) => {}
Err(e) => {
tracing::warn!(error = %e, "auth.json unreadable; keeping managed config until it recovers");
return;
}
}
let home = crate::util::kigi_home::kigi_home();
let Some(_lock) = try_lock_managed_config(&home) else {
return; // another process is syncing; retry next call
@@ -427,7 +394,7 @@ pub fn is_fetch_enabled() -> bool {
/// Fetch managed config + requirements and write to `~/.kigi/`, trying the
/// deployment key first, then a signed-in team. `Ok(false)` when neither applies.
pub async fn sync() -> Result<bool, ManagedConfigError> {
Ok(sync_with_budget(SyncBudget::Standard, None).await?.wrote)
Ok(sync_with_budget(SyncBudget::Standard).await?.wrote)
}
struct SyncOutcome {
@@ -479,25 +446,16 @@ impl SyncOutcome {
/// Runs a sync under `budget`'s deadline, returning `None` when the deadline
/// elapses first.
async fn sync_bounded(
budget: SyncBudget,
team_override: Option<GrokAuth>,
) -> Option<Result<SyncOutcome, ManagedConfigError>> {
let sync = sync_with_budget(budget, team_override);
async fn sync_bounded(budget: SyncBudget) -> Option<Result<SyncOutcome, ManagedConfigError>> {
let sync = sync_with_budget(budget);
match budget.deadline() {
Some(deadline) => tokio::time::timeout(deadline, sync).await.ok(),
None => Some(sync.await),
}
}
/// `team_override` pins a specific team principal (the just-authenticated one,
/// post-login) instead of re-deriving the team from `auth.json`; `None` uses
/// [`read_active_team_auth`] (the current eligible team).
async fn sync_with_budget(
budget: SyncBudget,
team_override: Option<GrokAuth>,
) -> Result<SyncOutcome, ManagedConfigError> {
let outcome = sync_inner(budget, team_override).await?;
async fn sync_with_budget(budget: SyncBudget) -> Result<SyncOutcome, ManagedConfigError> {
let outcome = sync_inner(budget).await?;
// Mark only when a principal was consulted AND the fetch wasn't signature-rejected —
// a rejected fetch persisted nothing, so marking would claim an unwritten body. Lock
// contention still marks (the holder persists the same config).
@@ -519,73 +477,31 @@ enum FetchedConfig {
key: String,
body: ManagedConfigResponse,
},
Team {
auth: Box<GrokAuth>,
body: ManagedConfigResponse,
},
/// No deployment key configured and no eligible team signed in.
/// No deployment key configured.
NoPrincipal,
}
/// Fetches the configuration for the current principal without touching disk:
/// the deployment key first, then a signed-in team. The installing sync and the
/// read-only `grok setup --json` both build on this.
async fn fetch_for_principal(
budget: SyncBudget,
team_override: Option<GrokAuth>,
) -> Result<FetchedConfig, ManagedConfigError> {
async fn fetch_for_principal(budget: SyncBudget) -> Result<FetchedConfig, ManagedConfigError> {
let max_attempts = budget.max_attempts();
// Resolve from the merged config (managed_config_url > cli_chat_proxy_base_url,
// including the enterprise single-endpoint derivation) so endpoint overrides
// are honored and the bearer isn't sent to the public default.
// Resolve from the merged config (managed_config_url override) so endpoint
// overrides are honored and the bearer isn't sent to the public default.
let url =
crate::agent::config::EndpointsConfig::from_effective_config().resolve_managed_config_url();
let team_auth = team_override.or_else(read_active_team_auth);
if let Some(dk) = resolve_deployment_key() {
let source = ManagedConfigSource::DeploymentKey;
match fetch_managed_config(&url, &dk, source, max_attempts).await {
// A rejected dk (stale env/config) must not starve a valid team
// sign-in: fall through. Network/5xx do NOT — same unreachable
// server, double the latency for nothing.
Err(ManagedConfigError::DeploymentKeyRejected) if team_auth.is_some() => {
tracing::warn!("deployment key rejected; falling back to the team session token");
}
Err(e) => return Err(e),
// Fall through to the team only when the dk has no config row: an apply
// converges disk to the served set, and the empty dk body must not delete
// the team's files. Gate on row existence, not content (which can serve empty).
Ok(body) if !body.config_exists() && team_auth.is_some() => {
tracing::debug!("deployment key has no config; trying the team principal");
}
Ok(body) => return Ok(FetchedConfig::DeploymentKey { key: dk, body }),
}
}
// The proxy resolves the team from the principal and returns its config.
if let Some(auth) = team_auth {
let body = fetch_managed_config(
&url,
&auth.key,
ManagedConfigSource::TeamOauth,
max_attempts,
)
.await?;
return Ok(FetchedConfig::Team {
auth: Box::new(auth),
body,
});
let body = fetch_managed_config(&url, &dk, source, max_attempts).await?;
return Ok(FetchedConfig::DeploymentKey { key: dk, body });
}
Ok(FetchedConfig::NoPrincipal)
}
async fn sync_inner(
budget: SyncBudget,
team_override: Option<GrokAuth>,
) -> Result<SyncOutcome, ManagedConfigError> {
match fetch_for_principal(budget, team_override).await? {
async fn sync_inner(budget: SyncBudget) -> Result<SyncOutcome, ManagedConfigError> {
match fetch_for_principal(budget).await? {
FetchedConfig::DeploymentKey { key, body } => {
let source = ManagedConfigSource::DeploymentKey;
let fingerprint = deployment_key_fingerprint(&key);
@@ -610,18 +526,6 @@ async fn sync_inner(
&outcome,
))
}
FetchedConfig::Team { auth, body } => {
let source = ManagedConfigSource::TeamOauth;
let outcome = apply_fetched(&body, source, auth.team_id.as_deref(), None)?;
// Team identity is bound via principal (team id), not a key fingerprint.
Ok(SyncOutcome::from_fetch(
&body,
source,
auth.team_id.clone(),
None,
&outcome,
))
}
FetchedConfig::NoPrincipal => Ok(SyncOutcome {
wrote: false,
served: false,
@@ -720,7 +624,9 @@ fn evict_prior_managed_config(home: &std::path::Path) {
fn credential_present(source: ManagedConfigSource) -> bool {
match source {
ManagedConfigSource::DeploymentKey => resolve_deployment_key().is_some(),
ManagedConfigSource::TeamOauth => team_principal_signed_in().unwrap_or(true),
// Team principals no longer exist; a cached team-sourced config has
// no live credential behind it.
ManagedConfigSource::TeamOauth => false,
}
}
@@ -744,20 +650,15 @@ pub enum ManagedConfigSync {
/// waiting for the background tick. `authenticated` pins the just-logged-in
/// principal (`None` = on-disk team). Latency-bounded by [`SyncBudget::Login`];
/// failures are logged, not propagated (the background loop retries).
pub async fn post_login_sync(authenticated: Option<GrokAuth>) -> ManagedConfigSync {
pub async fn post_login_sync(_authenticated: Option<KimiAuth>) -> ManagedConfigSync {
clear_orphan();
if !is_fetch_enabled() {
return ManagedConfigSync::Skipped;
}
// The just-authenticated team, else the on-disk one — reused for the gate
// and the sync (one auth.json read). With no team, only sync if due anyway.
let team = authenticated
.and_then(eligible_team_principal)
.or_else(read_active_team_auth);
if team.is_none() && !crate::config::is_managed_config_stale_for(&current_serving_identity()) {
if !crate::config::is_managed_config_stale_for(&current_serving_identity()) {
return ManagedConfigSync::Skipped;
}
match sync_bounded(SyncBudget::Login, team).await {
match sync_bounded(SyncBudget::Login).await {
// Nothing was persisted for a rejected envelope — that's a failure to
// report, not "no change" (the gate may refuse the next session).
Some(Ok(SyncOutcome {
@@ -791,14 +692,14 @@ pub async fn post_login_sync(authenticated: Option<GrokAuth>) -> ManagedConfigSy
/// Whether a credential exists that `grok setup` could install config for.
pub fn has_principal() -> bool {
resolve_deployment_key().is_some() || read_active_team_auth().is_some()
resolve_deployment_key().is_some()
}
/// Whether a managed identity owns this machine, IGNORING token expiry (unlike [`has_principal`]) so an
/// expired/backdated `auth.json` can't disarm the gate. Unreadable → present (fail-safe; the gate ANDs this
/// with [`crate::config::managed_policy_compromised_for`], which a personal user never satisfies).
fn managed_principal_present() -> bool {
resolve_deployment_key().is_some() || team_principal_signed_in().unwrap_or(true)
resolve_deployment_key().is_some()
}
/// The serving identity for an optional team id: a configured deployment key always
@@ -820,22 +721,12 @@ fn serving_identity_from(team_id: Option<String>) -> crate::config::ServingIdent
/// The identity to check the cache against for whoever serves now: a configured deployment key wins
/// (else the active team, else none).
pub fn current_serving_identity() -> crate::config::ServingIdentity {
serving_identity_from(read_active_team_auth().and_then(|a| a.team_id))
serving_identity_from(None)
}
/// The client's team_id, IGNORING token expiry (the binding must survive the cold-start
/// expired window). Must NOT special-case a configured deployment key — that would
/// disable envelope binding for a real team user. Used at fetch time to bind the envelope.
/// Team principals no longer exist in the Kimi Code auth model.
pub fn active_team_id_any_expiry() -> Option<String> {
let home = crate::util::kigi_home::kigi_home();
let store = crate::auth::read_auth_json(&home.join("auth.json")).ok()?;
store
.values()
.find(|a| a.is_team_principal())
.and_then(|a| a.team_id.clone())
// A blank team_id (malformed auth.json) is unknown, not a distinct identity: it must not
// feed the gate's identity checks, the tenant-switch purge, or the envelope binding.
.filter(|id| !id.trim().is_empty())
None
}
/// Like [`current_serving_identity`] but IGNORING token expiry, for the enforcement gate:
@@ -854,36 +745,16 @@ pub async fn ensure_managed_policy_present(
if !is_fetch_enabled() {
return;
}
// Cheap disk-only gates before any network token refresh, so the boot path doesn't pay
// an `auth()` in the common cases. A personal user (no deploy key, and no team in
// `auth.json` even ignoring expiry) skips entirely; a usable identity whose cache isn't
// hard-stale also skips. Only an expired-but-refreshable team token (identity reads
// `None` before the refresh) or a hard-stale cache falls through to `auth()` below.
// `auth.json` unreadable (`Err`) is NOT treated as "no principal" — that would skip
// enforcement on a transient read blip.
if resolve_deployment_key().is_none() && matches!(team_principal_signed_in(), Ok(false)) {
return;
}
let identity = current_serving_identity();
if !matches!(identity, crate::config::ServingIdentity::None)
&& !crate::config::is_managed_config_hard_stale_for(&identity)
{
return;
}
// Refresh before the heal so an expired-but-refreshable team token isn't dropped by
// the expiry filter. Bounded; deploy-key machines have no OAuth (auth() → None).
let team = tokio::time::timeout(SESSION_START_AUTH_DEADLINE, auth_manager.auth())
.await
.ok()
.and_then(Result::ok)
.filter(GrokAuth::is_team_principal);
if !has_principal() {
// Cheap disk-only gates: only a configured deployment key can own managed
// policy now (team principals no longer exist).
let _ = auth_manager;
if resolve_deployment_key().is_none() {
return;
}
if !crate::config::is_managed_config_hard_stale_for(&current_serving_identity()) {
return;
}
match sync_bounded(SyncBudget::SessionStart, team).await {
match sync_bounded(SyncBudget::SessionStart).await {
Some(Ok(_)) => {}
Some(Err(e)) => tracing::warn!("session-start managed policy refresh failed: {e}"),
None => tracing::warn!("session-start managed policy refresh timed out"),
@@ -987,9 +858,8 @@ pub struct SetupReport {
/// Fetches the report behind `grok setup --json` without writing anything:
/// no artifacts, no signature sidecar, no sync marker.
pub async fn fetch_setup_report() -> Result<SetupReport, ManagedConfigError> {
let (source, body) = match fetch_for_principal(SyncBudget::Standard, None).await? {
let (source, body) = match fetch_for_principal(SyncBudget::Standard).await? {
FetchedConfig::DeploymentKey { body, .. } => (Some("deploymentKey"), body),
FetchedConfig::Team { body, .. } => (Some("teamOauth"), body),
FetchedConfig::NoPrincipal => (None, ManagedConfigResponse::default()),
};
// Match the installer's trust decision: a payload `grok setup` would refuse
@@ -1015,7 +885,7 @@ pub async fn fetch_setup_report() -> Result<SetupReport, ManagedConfigError> {
/// Run the `grok setup` sync for the current principal. The caller must check
/// [`has_principal`] first and render the no-principal guidance.
pub async fn run_setup() -> SetupOutcome {
match sync_with_budget(SyncBudget::Standard, None).await {
match sync_with_budget(SyncBudget::Standard).await {
// A rejected envelope persisted nothing — reporting Installed would mask a
// fetch the gate is about to refuse.
Ok(SyncOutcome {
crates/codegen/kigi-shell/src/mcp_doctor.rs
+4 -4
View File
@@ -7,7 +7,7 @@ use std::sync::Arc;
use kigi_tools::types::config_source::ConfigSource;
use serde::Serialize;
use crate::auth::GrokComConfig;
use crate::auth::KimiCodeConfig;
use crate::session::managed_mcp;
use crate::session::mcp_servers;
@@ -272,13 +272,13 @@ fn managed_found(
/// Discover managed `grok_com_*` servers if the user has xAI auth on disk.
async fn try_discover_managed_servers() -> (ConfigSourceStatus, Vec<DiscoveredServer>) {
let kigi_home = kigi_tools::util::kigi_home::kigi_home();
let grok_com_config = GrokComConfig::default();
let auth_manager = Arc::new(crate::auth::AuthManager::new(&kigi_home, grok_com_config));
let kimi_code_config = KimiCodeConfig::default();
let auth_manager = Arc::new(crate::auth::AuthManager::new(&kigi_home, kimi_code_config));
let Some(snapshot) = auth_manager.current_or_expired() else {
return managed_skipped("not logged in");
};
if !snapshot.is_managed_mcp_eligible() {
if !snapshot.is_session_auth() {
return managed_skipped(format!("{:?} auth (not xAI OIDC)", snapshot.auth_mode));
}
crates/codegen/kigi-shell/src/remote/agent.rs
+1 -2
View File
@@ -5,7 +5,7 @@
use std::sync::Arc;
use crate::auth::{AuthManager, GrokComConfig};
use crate::auth::{AuthManager, KimiCodeConfig};
use anyhow::{Context, Result, bail};
use serde::de::DeserializeOwned;
@@ -64,7 +64,6 @@ impl SandboxClient {
.context("failed to resolve sandbox auth")?;
let mut builder = builder
.header("Authorization", format!("Bearer {}", auth.key))
.header("X-XAI-Token-Auth", GrokComConfig::default().token_header)
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION);
crates/codegen/kigi-shell/src/remote/chat_models_client.rs
-4
View File
@@ -121,10 +121,6 @@ impl ChatModelsClient {
.post(&url)
.json(&body)
.header("Authorization", format!("Bearer {}", auth.key))
.header(
"X-XAI-Token-Auth",
self.auth.grok_com_config().token_header.clone(),
)
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crates/codegen/kigi-shell/src/remote/client.rs
+16 -37
View File
@@ -1,5 +1,5 @@
//! HTTP client for backend CRUD operations.
use crate::auth::{GrokAuth, GrokComConfig};
use crate::auth::{KimiAuth, KimiCodeConfig};
use crate::session::export::{ExportedMessage, ExportedMetadata, ExportedSession};
use indexmap::IndexMap;
use prod_mc_cli_chat_proxy_types::SubagentBundle;
@@ -16,13 +16,12 @@ pub fn share_url(permission_id: &str) -> String {
}
fn add_cli_chat_proxy_headers_blocking(
builder: reqwest::blocking::RequestBuilder,
auth: &GrokAuth,
auth: &KimiAuth,
alpha_test_key: Option<&str>,
url: &str,
) -> reqwest::blocking::RequestBuilder {
let mut builder = builder
.header("Authorization", format!("Bearer {}", auth.key))
.header("X-XAI-Token-Auth", GrokComConfig::default().token_header)
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION);
if let Some(email) = &auth.email {
@@ -56,7 +55,7 @@ async fn add_bundle_fetch_headers(
Some(am) => am.auth().await.ok(),
None => None,
};
let mut credentials = crate::util::grok_auth_credentials::GrokAuthCredentials::new(
let mut credentials = crate::util::kigi_auth_credentials::KigiAuthCredentials::new(
resolved_auth.as_ref().map(|auth| auth.key.clone()),
);
credentials.deployment_key = deployment_key.map(str::to_owned);
@@ -313,7 +312,7 @@ impl BackendClient {
}
}
/// Attach a live `AuthManager` so every request resolves a fresh token
/// instead of requiring the caller to pass `&GrokAuth`.
/// instead of requiring the caller to pass `&KimiAuth`.
pub fn with_auth_manager(mut self, manager: std::sync::Arc<crate::auth::AuthManager>) -> Self {
let credentials: std::sync::Arc<dyn kigi_auth::AuthCredentialProvider> =
std::sync::Arc::new(
@@ -328,7 +327,7 @@ impl BackendClient {
self
}
/// Resolve auth from the attached `AuthManager`.
async fn resolve_auth(&self) -> Result<GrokAuth, BackendError> {
async fn resolve_auth(&self) -> Result<KimiAuth, BackendError> {
let manager = self
.auth_manager
.as_ref()
@@ -392,9 +391,7 @@ impl BackendClient {
.await?;
Ok(())
}
/// Build auth + identity headers.
/// Must include X-XAI-Token-Auth so nginx auth subrequest routes to authenticate_xai_grok_cli_token.
/// See: crates/codegen/kigi-shell/src/agent/app.rs:run_headless
/// Build auth + identity headers (plain bearer; no token-auth marker).
async fn auth_header_map(&self) -> Result<reqwest::header::HeaderMap, BackendError> {
use reqwest::header::{HeaderMap, HeaderValue};
let auth = self.resolve_auth().await?;
@@ -403,10 +400,6 @@ impl BackendClient {
HeaderValue::from_str(value)
.map_err(|e| BackendError::Auth(format!("invalid {name} header: {e}")))
};
headers.insert(
"X-XAI-Token-Auth",
required(&GrokComConfig::default().token_header, "X-XAI-Token-Auth")?,
);
headers.insert("x-userid", required(&auth.user_id, "x-userid")?);
if let Some(email) = &auth.email
&& let Ok(v) = HeaderValue::from_str(email)
@@ -556,7 +549,7 @@ impl BackendClient {
/// network). 4xx and parse errors are not retried.
pub fn fetch_settings_blocking(
cli_chat_proxy_base_url: &str,
auth: &GrokAuth,
auth: &KimiAuth,
alpha_test_key: Option<&str>,
) -> Option<crate::util::config::RemoteSettings> {
let client = crate::http::shared_blocking_client();
@@ -713,7 +706,7 @@ pub struct FetchModelsResult {
}
pub(crate) fn fetch_models_blocking(
endpoints: &crate::agent::config::EndpointsConfig,
auth: Option<&GrokAuth>,
auth: Option<&KimiAuth>,
fetch_auth: crate::agent::models::ModelFetchAuth,
) -> Result<FetchModelsResult, BackendError> {
let client = crate::http::shared_blocking_client();
@@ -1279,37 +1272,23 @@ mod tests {
let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(format!("{base}/v1"), seen_headers, handle)
}
fn test_auth() -> GrokAuth {
GrokAuth {
fn test_auth() -> KimiAuth {
KimiAuth {
key: "token".to_string(),
auth_mode: crate::auth::AuthMode::Oidc,
auth_mode: crate::auth::AuthMode::OAuth,
create_time: chrono::Utc::now(),
user_id: "user-1".to_string(),
email: Some("test@example.com".to_string()),
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: None,
team_name: None,
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token: None,
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
oidc_issuer: None,
oidc_client_id: None,
expires_in: Some(3600),
scope: None,
token_type: None,
}
}
fn test_auth_manager() -> Arc<crate::auth::AuthManager> {
let dir = tempfile::tempdir().unwrap();
let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::GrokComConfig::default());
let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::KimiCodeConfig::default());
mgr.hot_swap(test_auth());
std::mem::forget(dir);
Arc::new(mgr)
@@ -1336,7 +1315,7 @@ mod tests {
let headers = seen_headers.lock().unwrap();
let headers = headers.last().unwrap();
assert_eq!(headers.authorization.as_deref(), Some("Bearer token"));
assert_eq!(headers.token_auth.as_deref(), Some("xai-grok-cli"));
assert_eq!(headers.token_auth, None, "no token-auth marker header");
assert_eq!(headers.user_id.as_deref(), Some("user-1"));
assert_eq!(headers.email.as_deref(), Some("test@example.com"));
assert_eq!(headers.alpha_test_key, None);
crates/codegen/kigi-shell/src/remote/conversations_client.rs
+4 -8
View File
@@ -2,7 +2,7 @@ use std::sync::Arc;
use serde::{Deserialize, Serialize};
use crate::auth::{AuthManager, GrokAuth};
use crate::auth::{AuthManager, KimiAuth};
const KIGI_WEB_URL: &str = "https://grok.com";
@@ -108,9 +108,9 @@ impl ConversationsClient {
}
}
async fn require_xai_auth(&self) -> Result<GrokAuth, ConvError> {
async fn require_xai_auth(&self) -> Result<KimiAuth, ConvError> {
let auth = self.auth.auth().await.map_err(|_| ConvError::NoOauth)?;
if !auth.is_xai_auth() {
if !auth.is_session_auth() {
return Err(ConvError::NoOauth);
}
Ok(auth)
@@ -119,14 +119,10 @@ impl ConversationsClient {
fn apply_auth_headers(
&self,
builder: reqwest::RequestBuilder,
auth: &GrokAuth,
auth: &KimiAuth,
) -> reqwest::RequestBuilder {
let mut builder = builder
.header("Authorization", format!("Bearer {}", auth.key))
.header(
"X-XAI-Token-Auth",
self.auth.grok_com_config().token_header.clone(),
)
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crates/codegen/kigi-shell/src/remote/pull_smoke_test.rs
+5 -5
View File
@@ -4,17 +4,17 @@
#[cfg(test)]
mod tests {
use crate::auth::GrokAuth;
use crate::auth::KimiAuth;
use crate::remote::client::BackendClient;
use crate::session::storage::{JsonlStorageAdapter, StorageAdapter};
use std::collections::BTreeMap;
use std::sync::Arc;
fn load_prod_auth() -> Option<GrokAuth> {
fn load_prod_auth() -> Option<KimiAuth> {
let path = crate::util::kigi_home::kigi_home().join("auth.json");
let contents = std::fs::read_to_string(&path).ok()?;
let store: BTreeMap<String, GrokAuth> = serde_json::from_str(&contents).ok()?;
let scope = crate::auth::GrokComConfig::default().auth_scope();
let store: BTreeMap<String, KimiAuth> = serde_json::from_str(&contents).ok()?;
let scope = crate::auth::KimiCodeConfig::default().auth_scope();
crate::auth::lookup_auth(&store, &scope)
}
@@ -33,7 +33,7 @@ mod tests {
let auth = load_prod_auth().expect("No auth.json — run `grok login`");
let am = Arc::new(crate::auth::AuthManager::new(
&crate::util::kigi_home::kigi_home(),
crate::auth::GrokComConfig::default(),
crate::auth::KimiCodeConfig::default(),
));
am.hot_swap(auth);
let client = BackendClient::new().with_auth_manager(am.clone());
crates/codegen/kigi-shell/src/remote/workspaces_client.rs
+1 -5
View File
@@ -77,7 +77,7 @@ impl WorkspacesClient {
pub async fn list_workspaces(&self, q: &WsQuery) -> Result<ListWorkspacesPage, WsError> {
let auth = self.auth.auth().await.map_err(|_| WsError::NoOauth)?;
if !auth.is_xai_auth() {
if !auth.is_session_auth() {
return Err(WsError::NoOauth);
}
@@ -98,10 +98,6 @@ impl WorkspacesClient {
.get(&url)
.query(&query)
.header("Authorization", format!("Bearer {}", auth.key))
.header(
"X-XAI-Token-Auth",
self.auth.grok_com_config().token_header.clone(),
)
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs
+6 -81
View File
@@ -328,7 +328,7 @@ impl SessionActor {
.auth_manager
.as_ref()
.and_then(|am| am.current_or_expired())
.filter(|a| a.is_xai_auth())
.filter(|a| a.is_session_auth())
.map(|a| a.user_id),
origin_client: self.origin_client.clone(),
attribution_callback: self.attribution_callback.clone(),
@@ -476,17 +476,11 @@ impl SessionActor {
.and_then(|am| am.current_or_expired().map(|a| a.key.clone()));
let models = self.models_manager.models();
let endpoints = self.models_manager.endpoints();
let disable_api_key_auth = self
.auth_manager
.as_ref()
.map(|am| am.grok_com_config().api_key_auth_disabled())
.unwrap_or(false);
crate::agent::config::resolve_aux_model_sampling_config(
slug,
&models,
&endpoints,
session_key.as_deref(),
disable_api_key_auth,
creds.alpha_test_key.clone(),
creds.client_version.clone(),
)
@@ -678,32 +672,6 @@ impl SessionActor {
)),
);
}
if auth_recovery_eligible
&& crate::auth::devbox_login::is_devbox_environment()
&& let Some(ref am) = self.auth_manager
{
match am.try_devbox_recovery().await {
Ok(auth) => {
tracing::info!(
session_id = % self.session_info.id.0, user_id = % auth.user_id,
"auth recovery: sampler 401, devbox re-mint, retrying"
);
self.prepare_sampler_for_turn().await;
return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit);
}
Err(e) => {
tracing::warn!(
session_id = % self.session_info.id.0, error = % e,
"auth recovery: sampler 401, devbox re-mint failed"
);
kigi_log::unified_log::warn(
"auth recovery: sampler 401, devbox re-mint failed",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!({ "error" : format!("{e}") })),
);
}
}
}
if auth_recovery_eligible && let Some(ref am) = self.auth_manager {
if am.try_recover_unauthorized().await {
tracing::info!(
@@ -762,24 +730,6 @@ impl SessionActor {
.unwrap_or(crate::auth::AuthMode::ApiKey);
let auth_mode_str = format!("{auth_mode:?}");
let client_version = kigi_version::VERSION;
if auth_mode == crate::auth::AuthMode::WebLogin {
let msg = format!(
"{detailed_message}\n\n\
You are using a deprecated authentication method (WebLogin).\n\
This auth method is no longer supported and will cause errors.\n\n\
To fix: run `grok logout` then `grok login` to re-authenticate with OAuth2.\n\n\
Version: {client_version}"
);
self.log_terminal_failure("legacy_auth", error.status_code, &msg);
self.send_xai_notification(XaiSessionUpdate::RetryState(
crate::extensions::notification::RetryState::Failed {
error_type: "legacy_auth".to_string(),
message: msg.clone(),
},
))
.await;
return Err(acp::Error::internal_error().data(msg));
}
let is_model_404 =
error.status_code == Some(404) && detailed_message.contains("does not exist");
let is_auth_401 =
@@ -932,8 +882,10 @@ impl SessionActor {
None,
);
}
use crate::auth::{is_jwt_expired_or_near, parse_jwt_expiration};
const REFRESH_THRESHOLD: chrono::Duration = chrono::Duration::minutes(5);
// BYOK path: pick up an externally rotated per-model key from
// config.toml. Kimi bearers are opaque (no client-side expiry
// probing); a changed on-disk key is adopted, an unchanged one is a
// no-op.
let creds = self.chat_state_handle.get_credentials().await;
let current_key = creds.api_key;
let current_model_id = self
@@ -943,41 +895,14 @@ impl SessionActor {
.map(|c| c.model)
.unwrap_or_default();
let Some(ref key) = current_key else { return };
if !is_jwt_expired_or_near(key, REFRESH_THRESHOLD) {
if let Some(exp) = parse_jwt_expiration(key) {
let remaining_secs = (exp - chrono::Utc::now()).num_seconds();
tracing::debug!(
model = % current_model_id, remaining_secs,
"JWT token valid, no refresh needed"
);
} else {
tracing::debug!(
model = % current_model_id, key_len = key.len(),
"Token is not a JWT, expiry-based refresh not applicable"
);
}
return;
}
let remaining_secs =
parse_jwt_expiration(key).map_or(0, |exp| (exp - chrono::Utc::now()).num_seconds());
tracing::info!(
model = % current_model_id, remaining_secs,
"JWT near expiry, refreshing from config.toml"
);
let Some(new_key) = self.reload_api_key_from_config(&current_model_id) else {
return;
};
if key == &new_key {
tracing::warn!(
model = % current_model_id,
"Config.toml returned same token (not yet rotated by external process?)"
);
return;
}
let new_remaining_secs = parse_jwt_expiration(&new_key)
.map_or(0, |exp| (exp - chrono::Utc::now()).num_seconds());
tracing::info!(
model = % current_model_id, new_remaining_secs, key_len = new_key.len(),
model = % current_model_id, key_len = new_key.len(),
"Refreshed API token from config.toml"
);
let mut creds = self.chat_state_handle.get_credentials().await;
crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs
+36 -122
View File
@@ -1,6 +1,6 @@
use super::support::*;
use super::*;
use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig};
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::sync::mpsc;
@@ -17,12 +17,12 @@ impl crate::auth::refresh::TokenRefresher for AlwaysSucceedRefresher {
_reason: crate::auth::refresh::RefreshReason,
) -> crate::auth::refresh::RefreshOutcome {
self.called.store(true, Ordering::SeqCst);
crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth {
crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth {
key: "refreshed-test-token".to_string(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt-new".into()),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
}))
}
}
@@ -34,13 +34,13 @@ fn auth_manager_with_refresher(
refresher: Arc<dyn crate::auth::refresh::TokenRefresher>,
) -> (tempfile::TempDir, Arc<AuthManager>) {
let dir = tempfile::tempdir().expect("tempdir");
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
am.hot_swap(KimiAuth {
key: "initial-test-key".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
am.set_refresher(refresher);
(dir, am)
@@ -121,13 +121,13 @@ async fn make_actor_with_method_and_credentials(
/// cache hit). The tempdir must outlive the manager (auth.json path).
fn auth_manager_with_valid_token(key: &str) -> (tempfile::TempDir, Arc<AuthManager>) {
let dir = tempfile::tempdir().expect("tempdir");
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
am.hot_swap(KimiAuth {
key: key.into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
(dir, am)
}
@@ -305,12 +305,12 @@ async fn proactive_refresh_makes_per_turn_refresh_a_cache_hit() {
_: crate::auth::refresh::RefreshReason,
) -> crate::auth::refresh::RefreshOutcome {
self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth {
crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth {
key: "proactive-fresh".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt-new".into()),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
}))
}
}
@@ -318,14 +318,12 @@ async fn proactive_refresh_makes_per_turn_refresh_a_cache_hit() {
});
let (_dir, am) = auth_manager_with_refresher(refresher);
let cancel = tokio_util::sync::CancellationToken::new();
am.start_proactive_refresh(cancel.clone());
// Wait for proactive task to fire.
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
// Drive one loop-body iteration directly (the production loop
// ticks on a fixed 60s cadence, far too slow for a unit test).
am.proactive_tick(false).await;
assert!(
call_count.load(Ordering::SeqCst) >= 1,
"proactive task must have fired"
"proactive tick must have refreshed"
);
let count_after_proactive = call_count.load(Ordering::SeqCst);
@@ -350,8 +348,6 @@ async fn proactive_refresh_makes_per_turn_refresh_a_cache_hit() {
Some("proactive-fresh"),
"per-turn refresh must pick up the proactively-refreshed token"
);
cancel.cancel();
})
.await;
}
@@ -370,49 +366,6 @@ fn model_not_found_error() -> kigi_sampler::SamplingErrorInfo {
}
}
/// 404 model-not-found with a legacy WebLogin token appends a
/// "Legacy auth detected" hint to the error message.
#[tokio::test(flavor = "current_thread")]
async fn legacy_auth_hint_on_404_model_not_found() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let dir = tempfile::tempdir().expect("tempdir");
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
key: "legacy-token".into(),
auth_mode: AuthMode::WebLogin,
..GrokAuth::test_default()
});
let (actor, _rx) = make_actor_with_auth_manager(Some(am)).await;
let result = actor.handle_sampling_failure(model_not_found_error()).await;
let err = match result {
Err(e) => e,
Ok(_) => panic!("expected Err from handle_sampling_failure"),
};
let data = err.data.unwrap();
let msg = data.as_str().unwrap();
assert!(
msg.contains("deprecated authentication method"),
"404 with WebLogin must include deprecation message, got: {msg}"
);
assert!(
msg.contains("grok logout"),
"hint must mention `grok logout`, got: {msg}"
);
assert!(
msg.contains("grok login"),
"hint must mention `grok login`, got: {msg}"
);
assert!(
msg.contains("Version:"),
"must show client version, got: {msg}"
);
})
.await;
}
/// Build a 401-shaped error that bypasses step 4b's auth recovery.
///
/// In production, 401s arrive as `SamplingErrorKind::Auth` with
@@ -437,47 +390,6 @@ fn unauthorized_401_error() -> kigi_sampler::SamplingErrorInfo {
}
}
/// 401 Unauthorized with a legacy WebLogin token appends a
/// "Legacy auth detected" hint to the error message.
#[tokio::test(flavor = "current_thread")]
async fn legacy_auth_hint_on_401_unauthorized() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let dir = tempfile::tempdir().expect("tempdir");
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
key: "legacy-token".into(),
auth_mode: AuthMode::WebLogin,
..GrokAuth::test_default()
});
let (actor, _rx) = make_actor_with_auth_manager(Some(am)).await;
let result = actor
.handle_sampling_failure(unauthorized_401_error())
.await;
let err = match result {
Err(e) => e,
Ok(_) => panic!("expected Err from handle_sampling_failure"),
};
let data = err.data.unwrap();
let msg = data.as_str().unwrap();
assert!(
msg.contains("deprecated authentication method"),
"401 with WebLogin must include deprecation message, got: {msg}"
);
assert!(
msg.contains("grok logout"),
"hint must mention `grok logout`, got: {msg}"
);
assert!(
msg.contains("grok login"),
"hint must mention `grok login`, got: {msg}"
);
})
.await;
}
/// 401 with OIDC auth must NOT append the legacy hint.
#[tokio::test(flavor = "current_thread")]
async fn no_legacy_hint_on_401_for_oidc_auth() {
@@ -485,13 +397,13 @@ async fn no_legacy_hint_on_401_for_oidc_auth() {
local
.run_until(async {
let dir = tempfile::tempdir().expect("tempdir");
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
am.hot_swap(KimiAuth {
key: "oidc-token".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
let (actor, _rx) = make_actor_with_auth_manager(Some(am)).await;
@@ -513,7 +425,7 @@ async fn no_legacy_hint_on_401_for_oidc_auth() {
"OIDC auth must NOT trigger WebLogin deprecation on 401, got: {msg}"
);
assert!(
msg.contains("Auth: Oidc"),
msg.contains("Auth: OAuth"),
"OIDC 401 must show auth mode in enriched message, got: {msg}"
);
})
@@ -527,13 +439,13 @@ async fn no_legacy_hint_for_oidc_auth() {
local
.run_until(async {
let dir = tempfile::tempdir().expect("tempdir");
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
am.hot_swap(KimiAuth {
key: "oidc-token".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
let (actor, _rx) = make_actor_with_auth_manager(Some(am)).await;
@@ -553,7 +465,7 @@ async fn no_legacy_hint_for_oidc_auth() {
"OIDC auth must NOT trigger WebLogin deprecation, got: {msg}"
);
assert!(
msg.contains("Auth: Oidc"),
msg.contains("Auth: OAuth"),
"OIDC 404 must show auth mode in enriched message, got: {msg}"
);
assert!(
@@ -627,9 +539,9 @@ async fn sampler_401_session_method_with_stale_api_key_auth_type_still_recovers(
.await;
}
/// Same regression via the `oidc` method id (the other session-based variant).
/// Same regression via the interactive-login method id (the other session-based variant).
#[tokio::test(flavor = "current_thread")]
async fn sampler_401_oidc_method_with_stale_api_key_auth_type_still_recovers() {
async fn sampler_401_login_method_with_stale_api_key_auth_type_still_recovers() {
let local = tokio::task::LocalSet::new();
local
.run_until(async {
@@ -641,7 +553,7 @@ async fn sampler_401_oidc_method_with_stale_api_key_auth_type_still_recovers() {
let (_dir, am) = auth_manager_with_refresher(refresher);
let (actor, _rx) = make_actor_with_method_and_credentials(
Some(am),
"oidc",
"grok.com",
kigi_chat_state::AuthType::ApiKey,
"stale-session-jwt".to_string(),
)
@@ -651,7 +563,7 @@ async fn sampler_401_oidc_method_with_stale_api_key_auth_type_still_recovers() {
assert!(
matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
"oidc method must recover even when auth_type transiently reads ApiKey"
"interactive-login method must recover even when auth_type transiently reads ApiKey"
);
assert!(
called.load(Ordering::SeqCst),
@@ -779,7 +691,9 @@ async fn session_born_on_api_key_recovers_after_oidc_login_without_restart() {
// the shared handle this running actor already holds (no re-spawn).
actor
.auth_method_id
.store(Some(std::sync::Arc::new(acp::AuthMethodId::new("oidc"))));
.store(Some(std::sync::Arc::new(acp::AuthMethodId::new(
"cached_token",
))));
// The gate is recomputed each turn from the shared handle, so the
// flip alone activates the live resolver on the very next turn --
crates/codegen/kigi-shell/src/session/acp_session_tests/goal/goal_classifier_e2e_tests.rs
+1 -1
View File
@@ -2644,7 +2644,7 @@ fn test_auth_manager_for_models() -> std::sync::Arc<crate::auth::AuthManager> {
let tmp = tempfile::tempdir().expect("tempdir");
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
tmp.path(),
crate::auth::GrokComConfig::default(),
crate::auth::KimiCodeConfig::default(),
));
std::mem::forget(tmp);
mgr
crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs
+4 -4
View File
@@ -132,13 +132,13 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
let dir = tempfile::tempdir().unwrap();
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
dir.path(),
crate::auth::GrokComConfig::default(),
crate::auth::KimiCodeConfig::default(),
));
mgr.hot_swap(crate::auth::GrokAuth {
auth_mode: crate::auth::AuthMode::Oidc,
mgr.hot_swap(crate::auth::KimiAuth {
auth_mode: crate::auth::AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..crate::auth::GrokAuth::test_default()
..crate::auth::KimiAuth::test_default()
});
std::mem::forget(dir);
Some(mgr)
crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs
+4 -4
View File
@@ -1256,13 +1256,13 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
let dir = tempfile::tempdir().unwrap();
let mgr = std::sync::Arc::new(crate::auth::AuthManager::new(
dir.path(),
crate::auth::GrokComConfig::default(),
crate::auth::KimiCodeConfig::default(),
));
mgr.hot_swap(crate::auth::GrokAuth {
auth_mode: crate::auth::AuthMode::Oidc,
mgr.hot_swap(crate::auth::KimiAuth {
auth_mode: crate::auth::AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..crate::auth::GrokAuth::test_default()
..crate::auth::KimiAuth::test_default()
});
std::mem::forget(dir);
Some(mgr)
crates/codegen/kigi-shell/src/session/acp_session_tests/media_gen_auth_retry_tests.rs
+11 -11
View File
@@ -1,17 +1,17 @@
use super::*;
use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig};
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
use kigi_tools::types::output::{ToolOutput, ToolRunResult};
use std::sync::atomic::{AtomicUsize, Ordering};
fn succeeding_am() -> Arc<AuthManager> {
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
am.hot_swap(KimiAuth {
key: "expired".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
struct Ok;
#[async_trait::async_trait]
@@ -20,11 +20,11 @@ fn succeeding_am() -> Arc<AuthManager> {
&self,
_: crate::auth::refresh::RefreshReason,
) -> crate::auth::refresh::RefreshOutcome {
crate::auth::refresh::RefreshOutcome::Success(Box::new(GrokAuth {
crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth {
key: "fresh".into(),
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
refresh_token: Some("rt-new".into()),
..GrokAuth::test_default()
..KimiAuth::test_default()
}))
}
}
@@ -36,13 +36,13 @@ fn succeeding_am() -> Arc<AuthManager> {
fn failing_am() -> Arc<AuthManager> {
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
am.hot_swap(KimiAuth {
key: "expired".into(),
auth_mode: AuthMode::Oidc,
auth_mode: AuthMode::OAuth,
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)),
..GrokAuth::test_default()
..KimiAuth::test_default()
});
struct Fail;
#[async_trait::async_trait]
crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs
+3 -3
View File
@@ -156,12 +156,12 @@ async fn actor_with_proxy(
let home = tempfile::tempdir().expect("tempdir");
let auth_manager = Arc::new(crate::auth::AuthManager::new(
home.path(),
crate::auth::GrokComConfig::default(),
crate::auth::KimiCodeConfig::default(),
));
// Valid (1h) token in-memory only — `auth()` fast-paths it without network.
auth_manager.hot_swap(crate::auth::GrokAuth {
auth_manager.hot_swap(crate::auth::KimiAuth {
expires_at: Some(Utc::now() + chrono::Duration::hours(1)),
..crate::auth::GrokAuth::test_default()
..crate::auth::KimiAuth::test_default()
});
let cfg = crate::agent::config::Config {
crates/codegen/kigi-shell/src/session/feedback_manager.rs
+10 -10
View File
@@ -981,23 +981,23 @@ mod tests {
async fn test_is_auth_permanently_failed_reads_auth_manager() {
use crate::agent::feedback_client::FeedbackClient;
use crate::auth::error::RefreshTokenFailedReason;
use crate::auth::{AuthManager, GrokAuth, GrokComConfig};
use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig};
use std::sync::Arc;
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
let client = FeedbackClient::new("http://example/v1", None).with_auth_manager(am.clone());
assert!(!client.is_auth_permanently_failed());
// The verdict is scoped to the live credential's key.
am.hot_swap(GrokAuth {
// The tombstone is scoped to the live credential's refresh token.
am.hot_swap(KimiAuth {
key: "tok".into(),
..GrokAuth::test_default()
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)),
..KimiAuth::test_default()
});
// Use a non-sticky reason: only recoverable verdicts age out (a sticky
// `RefreshTokenRejected` never expires), and this exercises the TTL path.
am.record_permanent_failure("tok".to_string(), RefreshTokenFailedReason::Other.into());
am.record_permanent_failure("rt".to_string(), RefreshTokenFailedReason::Other.into());
assert!(client.is_auth_permanently_failed());
am.force_permanent_failure_aged_out();
@@ -1018,7 +1018,7 @@ mod tests {
#[tokio::test]
async fn test_has_token_refresher_requires_refresher_attached() {
use crate::agent::feedback_client::FeedbackClient;
use crate::auth::{AuthManager, GrokComConfig};
use crate::auth::{AuthManager, KimiCodeConfig};
use std::sync::Arc;
struct NoOpRefresher;
@@ -1035,7 +1035,7 @@ mod tests {
}
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
let bare = FeedbackClient::new("http://example/v1", None);
assert!(!bare.has_token_refresher());
crates/codegen/kigi-shell/src/session/persistence.rs
+2 -5
View File
@@ -1915,11 +1915,8 @@ fn init_remote_sync(
"Writeback storage mode requires authentication. Run 'grok login' first.",
)
})?;
if let Some(auth) = auth_manager.current_or_expired() {
if auth.is_zdr_team() {
tracing::debug!("ZDR team: skipping remote sync");
return Ok(None);
}
if auth_manager.current_or_expired().is_some() {
// ZDR was an xAI team concept; nothing gates remote sync here.
} else {
tracing::warn!(
"writeback: no auth loaded yet, ZDR check skipped (backend enforces server-side)"
crates/codegen/kigi-shell/src/session/unified_list/mod.rs
+5 -7
View File
@@ -565,13 +565,12 @@ mod tests {
fn xai_auth_manager(dir: &std::path::Path) -> std::sync::Arc<crate::auth::AuthManager> {
let am = std::sync::Arc::new(crate::auth::AuthManager::new(
dir,
crate::auth::GrokComConfig::default(),
crate::auth::KimiCodeConfig::default(),
));
am.hot_swap(crate::auth::GrokAuth {
auth_mode: crate::auth::AuthMode::Oidc,
oidc_issuer: Some(crate::auth::xai_oauth2_issuer().to_owned()),
am.hot_swap(crate::auth::KimiAuth {
auth_mode: crate::auth::AuthMode::OAuth,
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
..crate::auth::GrokAuth::test_default()
..crate::auth::KimiAuth::test_default()
});
am
}
@@ -648,9 +647,8 @@ mod tests {
let home = tempfile::tempdir().expect("tempdir");
let auth = std::sync::Arc::new(crate::auth::AuthManager::new(
home.path(),
crate::auth::GrokComConfig::default(),
crate::auth::KimiCodeConfig::default(),
));
auth.set_devbox_env_for_test(false);
let client = ConversationsClient::new(auth);
let mut req = ListReq::default();
force_kind_chat(&mut req);
crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs
+1 -1
View File
@@ -125,7 +125,7 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap<String, bool>) -> SubagentSpawnCon
workspace_ops: kigi_workspace::WorkspaceOps::for_test(),
auth_manager: Arc::new(crate::auth::AuthManager::new(
std::path::Path::new("/tmp/nonexistent-grok-test"),
crate::auth::GrokComConfig::default(),
crate::auth::KimiCodeConfig::default(),
)),
attribution_callback: None,
parent_agent_name: None,
crates/codegen/kigi-shell/src/tier.rs
-58
View File
@@ -1,58 +0,0 @@
//! Subscription-tier classification shared across the shell and the pager.
//!
//! The subscription tier reaches the client as a free-form **display-name
//! string** (from CCP `/settings` `subscription_tier_display`, or the numeric
//! JWT `tier` claim mapped to a display-style string by
//! [`crate::agent::mvp_agent::jwt_tier_claim`]). There is no shared enum, so
//! gating decisions classify the string here in ONE place so the pager's
//! cosmetic slash-command gate and the shell's capability (toolset) gate can't
//! drift apart.
//!
//! "Restricted" tiers are the personal free tier and X Basic — the tiers the
//! server zero-limits on the Imagine and voice endpoints. Everything else
//! (SuperGrok, SuperGrok Heavy/Lite, X Premium/+, and any unknown future name)
//! is unrestricted (**fail-open**).
/// Whether a **known** subscription-tier display name is a gated tier: the free
/// tier (CCP display "Free" or an empty string) or X Basic (CCP display
/// "X Basic"; JWT-claim fallback spelling "x_basic").
///
/// Case-insensitive and whitespace-trimmed. Callers decide the policy for an
/// *absent* tier (`None`): the pager treats absence as restricted (cosmetic,
/// recovers live on the next settings update), while the shell treats absence as
/// unrestricted (fail-open — the server authoritatively enforces per-tier
/// limits, so never withhold a capability on a guess).
pub fn is_restricted_tier_name(tier: &str) -> bool {
let t = tier.trim().to_ascii_lowercase();
t.is_empty() || t == "free" || t == "x basic" || t == "x_basic"
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn restricted_names() {
assert!(is_restricted_tier_name(""));
assert!(is_restricted_tier_name(" "));
assert!(is_restricted_tier_name("Free"));
assert!(is_restricted_tier_name("free"));
assert!(is_restricted_tier_name("X Basic"));
assert!(is_restricted_tier_name("x_basic"));
assert!(is_restricted_tier_name(" X BASIC "));
}
#[test]
fn unrestricted_names() {
assert!(!is_restricted_tier_name("SuperGrok"));
assert!(!is_restricted_tier_name("SuperGrok Heavy"));
assert!(!is_restricted_tier_name("supergrok_lite"));
assert!(!is_restricted_tier_name("X Premium"));
assert!(!is_restricted_tier_name("x_premium_plus"));
// API keys are not free-tier gated.
assert!(!is_restricted_tier_name("api_key"));
assert!(!is_restricted_tier_name("API Key"));
// Unknown future tiers fail open.
assert!(!is_restricted_tier_name("some_new_plan"));
}
}
crates/codegen/kigi-shell/src/trace_classifier/mod.rs
+8 -16
View File
@@ -1073,18 +1073,10 @@ pub async fn resolve_api_key(explicit: Option<&str>, kigi_home: &Path) -> Result
/// * `Err(_)` — refresh attempt failed in a way the operator needs to
/// see (network error, refresh_token rejected by the IdP, etc.).
async fn non_interactive_auth_key(kigi_home: &Path) -> Result<Option<String>> {
use crate::auth::{AuthError, AuthManager, GrokComConfig};
use crate::auth::{AuthError, AuthManager, KimiCodeConfig};
// Production's `try_ensure_fresh_auth` clones the whole config to
// pass into `AuthManager::new` AND clones `auth_provider_command`
// again for `configure_refresher`. We extract the single field we
// need first, then move the rest of `config` into `AuthManager`
// — one `Option<String>` clone instead of one full struct clone
// plus one Option clone.
let config = GrokComConfig::default();
let auth_provider_command = config.auth_provider_command.clone();
let manager = std::sync::Arc::new(AuthManager::new(kigi_home, config));
manager.configure_refresher(auth_provider_command);
let manager = std::sync::Arc::new(AuthManager::new(kigi_home, KimiCodeConfig::default()));
manager.configure_refresher();
match manager.auth().await {
Ok(auth) => {
let trimmed = auth.key.trim();
@@ -2147,7 +2139,7 @@ mod tests {
/// path entirely — useful for "plain key, no refresh wanted"
/// fixtures.
fn write_auth_json(kigi_home: &Path, key: &str) {
let scope = crate::auth::GrokComConfig::default().auth_scope();
let scope = crate::auth::KimiCodeConfig::default().auth_scope();
let body = serde_json::json!({
scope: {
"key": key,
@@ -2172,11 +2164,11 @@ mod tests {
/// returns it via the fast path; the refresher chain is NOT
/// invoked, so no network call fires.
fn write_fresh_oidc_auth_json(kigi_home: &Path, key: &str) {
let scope = crate::auth::GrokComConfig::default().auth_scope();
let scope = crate::auth::KimiCodeConfig::default().auth_scope();
let body = serde_json::json!({
scope: {
"key": key,
"auth_mode": "oidc",
"auth_mode": "oauth",
"create_time": now_offset(0),
"expires_at": now_offset(3600),
"refresh_token": "test-refresh-token",
@@ -2192,11 +2184,11 @@ mod tests {
/// refresh chain has nothing to refresh against, so the auth
/// call fails non-interactively.
fn write_expired_oidc_auth_json_no_refresh(kigi_home: &Path, key: &str) {
let scope = crate::auth::GrokComConfig::default().auth_scope();
let scope = crate::auth::KimiCodeConfig::default().auth_scope();
let body = serde_json::json!({
scope: {
"key": key,
"auth_mode": "oidc",
"auth_mode": "oauth",
"create_time": now_offset(-7200),
"expires_at": now_offset(-3600),
"user_id": "test-user",
crates/codegen/kigi-shell/src/util/grok_auth_credentials.rs → crates/codegen/kigi-shell/src/util/kigi_auth_credentials.rs
+20 -25
View File
@@ -9,10 +9,10 @@ use std::sync::Arc;
/// an `AuthManager` (visibility checks, bundle fetches, tests).
///
/// Deployment key (enterprise) sends bare `Bearer`, routed to management key auth.
/// User token (xAI users) sends `Bearer` + `X-XAI-Token-Auth: xai-grok-cli`.
/// User tokens and deployment keys are both sent as a plain `Bearer`.
/// Deployment key takes precedence when both are present.
#[derive(Clone)]
pub struct GrokAuthCredentials {
pub struct KigiAuthCredentials {
pub user_token: Option<String>,
pub deployment_key: Option<String>,
pub alpha_test_key: Option<String>,
@@ -20,9 +20,9 @@ pub struct GrokAuthCredentials {
/// refresh chain; `resolve()` reads the in-memory cache.
auth_manager: Option<Arc<crate::auth::AuthManager>>,
}
impl std::fmt::Debug for GrokAuthCredentials {
impl std::fmt::Debug for KigiAuthCredentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GrokAuthCredentials")
f.debug_struct("KigiAuthCredentials")
.field(
"user_token",
&self.user_token.as_ref().map(|_| "<redacted>"),
@@ -42,7 +42,7 @@ impl std::fmt::Debug for GrokAuthCredentials {
.finish()
}
}
impl GrokAuthCredentials {
impl KigiAuthCredentials {
/// Static credentials from a snapshot token. No refresh capability.
pub fn new(user_token: Option<String>) -> Self {
Self {
@@ -82,7 +82,7 @@ impl GrokAuthCredentials {
/// Without this, the `resolve_async()` error fallback returns
/// credentials with no token, causing requests to be sent without
/// an Authorization header.
pub fn resolve(&self) -> GrokAuthCredentials {
pub fn resolve(&self) -> KigiAuthCredentials {
if let Some(ref am) = self.auth_manager
&& let Some(auth) = am.current_or_expired()
{
@@ -97,7 +97,7 @@ impl GrokAuthCredentials {
/// (memory -> disk -> active OIDC refresh). Falls back to sync
/// `resolve()` on error so transient refresh failures don't drop
/// the bearer.
pub async fn resolve_async(&self) -> GrokAuthCredentials {
pub async fn resolve_async(&self) -> KigiAuthCredentials {
let Some(ref am) = self.auth_manager else {
return self.clone();
};
@@ -120,12 +120,7 @@ impl GrokAuthCredentials {
let builder = if let Some(ref key) = self.deployment_key {
builder.header("Authorization", format!("Bearer {}", key))
} else if let Some(ref token) = self.user_token {
builder
.header("Authorization", format!("Bearer {}", token))
.header(
obfstr::obfstr!("X-XAI-Token-Auth"),
obfstr::obfstr!("xai-grok-cli"),
)
builder.header("Authorization", format!("Bearer {}", token))
} else {
builder
};
@@ -133,28 +128,28 @@ impl GrokAuthCredentials {
builder
}
}
impl kigi_auth::HttpAuth for GrokAuthCredentials {
impl kigi_auth::HttpAuth for KigiAuthCredentials {
fn apply(&self, builder: RequestBuilder, base_url: &str) -> RequestBuilder {
GrokAuthCredentials::apply(self, builder, base_url)
KigiAuthCredentials::apply(self, builder, base_url)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig};
use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig};
use chrono::{Duration, Utc};
use std::sync::Arc;
fn make_manager_with_token(
expires_at: chrono::DateTime<Utc>,
) -> (Arc<AuthManager>, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let auth = GrokAuth {
let mgr = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
let auth = KimiAuth {
key: "test-bearer-token".into(),
auth_mode: AuthMode::External,
auth_mode: AuthMode::OAuth,
expires_at: Some(expires_at),
create_time: Utc::now(),
..GrokAuth::test_default()
..KimiAuth::test_default()
};
mgr.hot_swap(auth);
(mgr, dir)
@@ -162,14 +157,14 @@ mod tests {
#[test]
fn resolve_returns_token_when_not_expired() {
let (mgr, _dir) = make_manager_with_token(Utc::now() + Duration::hours(1));
let creds = GrokAuthCredentials::new(None).with_auth_manager(mgr);
let creds = KigiAuthCredentials::new(None).with_auth_manager(mgr);
let resolved = creds.resolve();
assert_eq!(resolved.user_token.as_deref(), Some("test-bearer-token"));
}
#[test]
fn resolve_returns_token_during_early_invalidation_window() {
let (mgr, _dir) = make_manager_with_token(Utc::now() + Duration::minutes(3));
let creds = GrokAuthCredentials::new(None).with_auth_manager(mgr.clone());
let creds = KigiAuthCredentials::new(None).with_auth_manager(mgr.clone());
assert!(mgr.current().is_none());
assert!(mgr.current_or_expired().is_some());
assert_eq!(
@@ -179,14 +174,14 @@ mod tests {
}
#[test]
fn resolve_returns_static_token_when_no_auth_manager() {
let creds = GrokAuthCredentials::new(Some("static-token".into()));
let creds = KigiAuthCredentials::new(Some("static-token".into()));
assert_eq!(creds.resolve().user_token.as_deref(), Some("static-token"));
}
#[test]
fn resolve_returns_none_when_no_token_at_all() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
let creds = GrokAuthCredentials::new(None).with_auth_manager(mgr);
let mgr = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
let creds = KigiAuthCredentials::new(None).with_auth_manager(mgr);
assert!(creds.resolve().user_token.is_none());
}
}
crates/codegen/kigi-shell/src/util/mod.rs
+1 -1
View File
@@ -1,7 +1,7 @@
pub mod agent_id;
pub mod config;
pub mod grok_auth_credentials;
pub mod hooks;
pub mod kigi_auth_credentials;
// The foundation utilities live in `kigi-shell-base` (upstream of this
// crate so they build in parallel). Re-exported at the original paths so