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+"));
}
}