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/extensions/auth.rs
+9 -64
View File
@@ -21,7 +21,6 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
"x.ai/auth/get_url" => handle_get_url(agent).await,
"x.ai/auth/logout" => handle_logout(agent, args).await,
"x.ai/auth/info" => handle_info(agent),
"x.ai/auth/check_subscription" => handle_check_subscription(agent).await,
_ => Err(acp::Error::method_not_found()),
}
}
@@ -111,7 +110,7 @@ async fn handle_get_url(agent: &MvpAgent) -> ExtResult {
to_raw_response(&serde_json::json!({
"auth_url": auth_url,
// `external_provider` kept for older clients; `mode` is authoritative.
"external_provider": mode.is_some_and(|m| m.is_external_provider()),
"external_provider": false,
"mode": mode.map(|m| m.as_wire_str()),
}))
}
@@ -141,43 +140,16 @@ async fn handle_logout(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
}))
}
/// Single-shot subscription re-check (retry button on paywall screen).
///
/// Calls `retry_subscription_check()`, then returns the updated auth
/// response with gate info so the pager can refresh the gate state.
async fn handle_check_subscription(agent: &MvpAgent) -> ExtResult {
agent.retry_subscription_check().await;
let response = agent.auth_response_with_meta();
to_raw_response(&serde_json::json!({
"authenticated": response.meta.is_some(),
"meta": response.meta,
}))
}
/// Returns current auth method ID, user profile fields, and team/principal
/// metadata.
/// Returns current auth method ID and the account fields the Kimi flow
/// exposes (email/user id are empty until a later feature surfaces them).
fn handle_info(agent: &MvpAgent) -> ExtResult {
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct AuthInfoResponse {
method_id: Option<String>,
email: Option<String>,
first_name: Option<String>,
last_name: Option<String>,
/// `grok-asset://` URL resolved by the Electron protocol handler,
/// or a full `http(s)://` URL passed through unchanged.
profile_image_url: Option<String>,
team_id: Option<String>,
team_name: Option<String>,
team_role: Option<String>,
organization_id: Option<String>,
organization_name: Option<String>,
organization_role: Option<String>,
principal_type: Option<String>,
principal_id: Option<String>,
user_blocked_reason: Option<String>,
team_blocked_reasons: Vec<String>,
coding_data_retention_opt_out: bool,
user_id: Option<String>,
auth_mode: Option<String>,
}
let method_id = agent
@@ -186,40 +158,13 @@ fn handle_info(agent: &MvpAgent) -> ExtResult {
.as_ref()
.map(|m| m.0.to_string());
let auth = agent.auth_manager.current();
let raw_asset_id = auth.as_ref().and_then(|a| a.profile_image_asset_id.clone());
// Return a grok-asset:// URL that the Electron renderer resolves at
// display time via a custom protocol handler. The handler proxies
// through cli-chat-proxy's /asset endpoint; Electron's HTTP cache
// handles reuse. No disk-cache or network call needed here.
let profile_image_url = match raw_asset_id.as_deref().filter(|k| !k.is_empty()) {
Some(key) if key.starts_with("http://") || key.starts_with("https://") => {
Some(key.to_owned())
}
Some(key) => Some(format!("grok-asset:///{key}")),
None => None,
};
to_raw_response(&AuthInfoResponse {
method_id,
email: auth.as_ref().and_then(|a| a.email.clone()),
first_name: auth.as_ref().and_then(|a| a.first_name.clone()),
last_name: auth.as_ref().and_then(|a| a.last_name.clone()),
profile_image_url,
team_id: auth.as_ref().and_then(|a| a.team_id.clone()),
team_name: auth.as_ref().and_then(|a| a.team_name.clone()),
team_role: auth.as_ref().and_then(|a| a.team_role.clone()),
organization_id: auth.as_ref().and_then(|a| a.organization_id.clone()),
organization_name: auth.as_ref().and_then(|a| a.organization_name.clone()),
organization_role: auth.as_ref().and_then(|a| a.organization_role.clone()),
principal_type: auth.as_ref().and_then(|a| a.principal_type.clone()),
principal_id: auth.as_ref().and_then(|a| a.principal_id.clone()),
user_blocked_reason: auth.as_ref().and_then(|a| a.user_blocked_reason.clone()),
team_blocked_reasons: auth
user_id: auth
.as_ref()
.map(|a| a.team_blocked_reasons.clone())
.unwrap_or_default(),
coding_data_retention_opt_out: auth
.as_ref()
.is_some_and(|a| a.coding_data_retention_opt_out),
.map(|a| a.user_id.clone())
.filter(|id| !id.is_empty()),
auth_mode: auth.as_ref().map(|a| format!("{:?}", a.auth_mode)),
})
}
crates/codegen/kigi-shell/src/extensions/auth_gate.rs
+4 -4
View File
@@ -1,17 +1,17 @@
use agent_client_protocol as acp;
use crate::auth::{AuthManager, GrokAuth};
use crate::auth::{AuthManager, KimiAuth};
/// Require xAI auth from a sync context, accepting tokens in the client-side buffer window.
/// Require a Kimi Code session from a sync context, accepting tokens in the client-side buffer window.
pub(crate) fn require_xai_auth(
auth_manager: &AuthManager,
missing_message: &'static str,
non_xai_message: &'static str,
) -> Result<GrokAuth, acp::Error> {
) -> Result<KimiAuth, acp::Error> {
let auth = auth_manager
.current_or_expired()
.ok_or_else(|| acp::Error::auth_required().data(missing_message))?;
if !auth.is_xai_auth() {
if !auth.is_session_auth() {
return Err(acp::Error::auth_required().data(non_xai_message));
}
Ok(auth)
crates/codegen/kigi-shell/src/extensions/billing.rs
-8
View File
@@ -213,10 +213,6 @@ async fn handle_get_billing(agent: &MvpAgent) -> ExtResult {
let credits_resp = crate::http::shared_client()
.get(&credits_url)
.header("Authorization", format!("Bearer {}", auth.key))
.header(
"X-XAI-Token-Auth",
crate::auth::GrokComConfig::default().token_header,
)
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
@@ -304,10 +300,6 @@ async fn handle_get_auto_topup_rule(agent: &MvpAgent) -> ExtResult {
let response = crate::http::shared_client()
.get(&url)
.header("Authorization", format!("Bearer {}", auth.key))
.header(
"X-XAI-Token-Auth",
crate::auth::GrokComConfig::default().token_header,
)
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crates/codegen/kigi-shell/src/extensions/bundle.rs
+7 -21
View File
@@ -484,37 +484,23 @@ mod tests {
.insert("review".to_string(), "# Review skill\n".to_string());
bundle
}
fn test_auth() -> crate::auth::GrokAuth {
crate::auth::GrokAuth {
fn test_auth() -> crate::auth::KimiAuth {
crate::auth::KimiAuth {
key: "token".to_string(),
auth_mode: crate::auth::AuthMode::Oidc,
auth_mode: crate::auth::AuthMode::OAuth,
create_time: chrono::Utc::now(),
user_id: "user-1".to_string(),
email: Some("test@example.com".to_string()),
first_name: None,
last_name: None,
profile_image_asset_id: None,
principal_type: None,
principal_id: None,
team_id: None,
team_name: None,
team_role: None,
organization_id: None,
organization_name: None,
organization_role: None,
user_blocked_reason: None,
team_blocked_reasons: vec![],
coding_data_retention_opt_out: false,
has_grok_code_access: None,
refresh_token: None,
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
oidc_issuer: None,
oidc_client_id: None,
expires_in: Some(3600),
scope: None,
token_type: None,
}
}
fn test_auth_manager() -> Arc<crate::auth::AuthManager> {
let dir = tempfile::tempdir().unwrap();
let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::GrokComConfig::default());
let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::KimiCodeConfig::default());
mgr.hot_swap(test_auth());
std::mem::forget(dir);
Arc::new(mgr)
crates/codegen/kigi-shell/src/extensions/mod.rs
-1
View File
@@ -17,7 +17,6 @@ pub mod memory;
pub mod notification;
pub mod plugins;
pub mod pr;
pub mod privacy;
pub mod prompt_history;
pub mod prompt_meta;
pub mod recap;
crates/codegen/kigi-shell/src/extensions/privacy.rs
-91
View File
@@ -1,91 +0,0 @@
//! `x.ai/privacy/setCodingDataRetention` extension handler.
//!
//! PUTs the new opt-out flag to cli-chat-proxy and updates local auth state
//! to match. The local update is fire-and-forget (best-effort cache refresh).
use agent_client_protocol as acp;
use serde::Deserialize;
use super::{ExtResult, parse_params, to_raw_response};
use crate::agent::MvpAgent;
#[tracing::instrument(skip_all, fields(method = %args.method))]
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
match args.method.as_ref() {
"x.ai/privacy/setCodingDataRetention" => handle_set(agent, args).await,
_ => Err(acp::Error::method_not_found()),
}
}
async fn handle_set(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct Params {
coding_data_retention_opt_out: bool,
}
let params: Params = parse_params(args)?;
let auth = agent.auth_manager.auth().await.map_err(|e| {
tracing::warn!(error = %e, "privacy: auth resolution failed");
acp::Error::auth_required()
.data("Authentication required. Run `grok login` to re-authenticate.")
})?;
let proxy_url = agent.cfg.borrow().endpoints.proxy_url();
let url = format!("{proxy_url}/privacy/coding-data-retention");
let token_header = agent.auth_manager.grok_com_config().token_header.clone();
let body = serde_json::json!({
"codingDataRetentionOptOut": params.coding_data_retention_opt_out,
});
let provider: std::sync::Arc<dyn kigi_auth::AuthCredentialProvider> = std::sync::Arc::new(
crate::auth::credential_provider::ShellAuthCredentialProvider::new(
agent.auth_manager.clone(),
None,
None,
),
);
let client = crate::http::with_auth_retry(crate::http::shared_client(), provider);
let resp = client
.put(&url)
.header("X-XAI-Token-Auth", &token_header)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
)
.json(&body)
.send()
.await
.map_err(|e| acp::Error::internal_error().data(format!("HTTP request failed: {e}")))?;
if !resp.status().is_success() {
let status = resp.status().as_u16();
let body = resp.text().await.unwrap_or_default();
tracing::warn!(status, "setCodingDataRetention request failed");
let friendly = serde_json::from_str::<serde_json::Value>(&body)
.ok()
.and_then(|v| {
v.get("error")
.or_else(|| v.get("message"))
.and_then(|e| e.as_str().map(String::from))
})
.unwrap_or_else(|| format!("server returned HTTP {status}"));
return Err(acp::Error::internal_error().data(friendly));
}
// Update local auth state to reflect the change.
// Use save_without_enrichment to avoid a race: update() spawns a
// background GET /user enrichment that may read stale ACL state
// and overwrite the opt-out flag back to its previous value.
let mut updated = auth.clone();
updated.coding_data_retention_opt_out = params.coding_data_retention_opt_out;
let _ = agent.auth_manager.save_without_enrichment(updated).await;
to_raw_response(&serde_json::json!({
"codingDataRetentionOptOut": params.coding_data_retention_opt_out,
}))
}
crates/codegen/kigi-shell/src/extensions/session_admin.rs
+3 -15
View File
@@ -111,10 +111,7 @@ async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR
// Send a SessionSummaryGenerated notification so the TUI updates its title
notify_session_title(agent, session_id, &req.title).await;
if agent.is_writeback_storage()
&& let Some(auth) = agent.current_auth()
&& !auth.is_zdr_team()
{
if agent.is_writeback_storage() && agent.current_auth().is_some() {
use crate::remote::client::BackendClient;
use crate::session::export::ExportedMetadata;
@@ -133,15 +130,7 @@ async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR
// Hook 2: update session replica with summary (fire-and-forget)
if let Some(client) = agent.session_registry_client() {
let sid = req.session_id.to_string();
let title = if agent
.auth_manager
.current_or_expired()
.is_some_and(|a| a.is_zdr_team())
{
None
} else {
Some(req.title.clone())
};
let title = Some(req.title.clone());
tokio::spawn(async move {
let update = crate::agent::session_registry_client::UpdateRequest {
summary: title,
@@ -248,8 +237,7 @@ async fn handle_session_delete(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR
// For writeback storage (non-ZDR): remote delete is authoritative for
// the cloud history and runs first; on failure no local bits are
// touched so the pager does not remove the row or toast success.
let needs_remote =
agent.is_writeback_storage() && agent.current_auth().is_some_and(|a| !a.is_zdr_team());
let needs_remote = agent.is_writeback_storage() && agent.current_auth().is_some();
// Shared delete: remote-first, then local disk + FTS eviction.
// Mirrored by the `grok sessions delete <id>` CLI path.
crates/codegen/kigi-shell/src/extensions/share.rs
+9 -17
View File
@@ -45,13 +45,6 @@ async fn handle_share_session(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRe
);
}
// Only block for ZDR teams (hard data-retention policy), not for
// coding-data-retention opt-out — sharing is user-initiated.
if auth.is_zdr_team() {
return Err(acp::Error::invalid_params()
.data("Session sharing is disabled for your team's data retention policy"));
}
// Find session info by searching through summaries
let summaries = list_summaries(None).await.map_err(|e| {
acp::Error::internal_error().data(format!("Failed to list sessions: {}", e))
@@ -94,7 +87,7 @@ async fn handle_share_session(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRe
fn require_xai_auth_for_share(
auth_manager: &crate::auth::AuthManager,
) -> Result<crate::auth::GrokAuth, acp::Error> {
) -> Result<crate::auth::KimiAuth, acp::Error> {
super::auth_gate::require_xai_auth(
auth_manager,
"Authentication required to share session",
@@ -105,8 +98,8 @@ fn require_xai_auth_for_share(
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::GrokComConfig;
use crate::auth::{AuthMode, GrokAuth};
use crate::auth::KimiCodeConfig;
use crate::auth::{AuthMode, KimiAuth};
use chrono::{Duration, Utc};
use std::sync::Arc;
use tempfile::tempdir;
@@ -117,7 +110,7 @@ mod tests {
let dir = tempdir().expect("tempdir for share auth test");
let mgr = Arc::new(crate::auth::AuthManager::new(
dir.path(),
GrokComConfig::default(),
KimiCodeConfig::default(),
));
let expires_at = Utc::now() + ttl;
@@ -126,9 +119,8 @@ mod tests {
// Only OIDC tokens against https://auth.x.ai (or the local-dev equivalent)
// return true from is_xai_auth(). This is required for the share tests to
// exercise the happy path through require_xai_auth_for_share.
let auth = GrokAuth {
auth_mode: AuthMode::Oidc,
oidc_issuer: Some("https://auth.x.ai".to_string()),
let auth = KimiAuth {
auth_mode: AuthMode::OAuth,
key: "test-key".into(),
expires_at: Some(expires_at),
create_time: Utc::now() - Duration::hours(1),
@@ -168,7 +160,7 @@ mod tests {
let dir = tempdir().expect("tempdir");
let mgr = Arc::new(crate::auth::AuthManager::new(
dir.path(),
GrokComConfig::default(),
KimiCodeConfig::default(),
));
assert!(require_xai_auth_for_share(&mgr).is_err());
}
@@ -178,12 +170,12 @@ mod tests {
let dir = tempdir().expect("tempdir");
let mgr = Arc::new(crate::auth::AuthManager::new(
dir.path(),
GrokComConfig::default(),
KimiCodeConfig::default(),
));
// API key is the simplest non-xAI credential (External and enterprise OIDC
// are also rejected the same way).
let non_xai = GrokAuth {
let non_xai = KimiAuth {
auth_mode: AuthMode::ApiKey,
key: "xai-test-key".into(),
create_time: Utc::now(),