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/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(