feat(providers): add xAI Grok subscription OAuth (device-code) + per-provider session auth
First subscription-OAuth provider beyond Kimi Code (26th registry variant). Log in with a Grok/SuperGrok/X subscription via RFC-8628 device-code OAuth (auth.x.ai), then use it against api.x.ai/v1 — reusing the existing xai wire (ChatCompletions + OpenAI listing + Passthrough + restrict + models_dev_id xai). Sourced from Pi (earendil-works/pi auth/oauth/xai.ts): client b1a00492..., scope 'openid profile email offline_access grok-cli:access api:access', standard Bearer (no x-xai-token-auth). Foundation (generalizes Kigi's Kimi-singleton OAuth to per-provider, root cause, not a patch): - Registry: OAuthConfig on PlatformSpec (client_id/host/device+token paths/scope/scope_key); XAI_OAUTH_CONFIG + XAI_GROK_SPEC (uses_oauth, method id 'xai-grok', an interactive login after kimi-code). - Generic device-code wire (auth/oauth_device.rs) + GenericDeviceRefresher, sharing the RFC-8628 core with Kimi; Kimi's bespoke flow is byte-identical (X-Msh headers, KIMI_CODE_OAUTH_SCOPE, keyring gating unchanged). - Per-provider AuthManager via a process-global pool (auth/oauth_registry.rs): build-on-demand with start_proactive_refresh, keyed by scope. The session resolves the AuthManager for the ACTIVE model's platform for bearer/refresh/ 401-recovery/api_key — an oauth-platform model always uses its OWN token, never the primary. - Live /models under OAuth; base routes oauth().is_some() -> platform.base_url() (kimi-code stays on proxy_url). Security: adversarial review + a systematic token-leak audit found and closed FIVE channels where the primary Kimi token could reach api.x.ai (bearer resolver, api_key stamping, aux summary/classifier/image-describe models, and subagent model-override). Each fix routes through the platform-aware resolver (the oauth model's pooled token or None, NEVER the primary) and is revert-to-red verified. No access/refresh token is ever logged. Registry at 26; picker updated (xai-grok interactive login row); TUI context-window already auto-updates per model. Full gate green (234 suites, fmt, clippy -D warnings, deny). GPT/Claude/Grok officially permit third-party subscription use.
This commit is contained in:
@@ -13,15 +13,47 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::auth::kimi_oauth::{
|
||||
DeviceAuthorization, DevicePollResult, poll_device_token, request_device_authorization,
|
||||
};
|
||||
use kigi_models::OAuthConfig;
|
||||
|
||||
use crate::auth::kimi_oauth::{DeviceAuthorization, DevicePollResult};
|
||||
use crate::auth::{AuthChannels, AuthManager, AuthUrlInfo, AuthUrlMode, KimiAuth};
|
||||
|
||||
/// Extra wait added to the poll interval when the server answers `slow_down`
|
||||
/// (OAuth-standard device-flow backpressure).
|
||||
const SLOW_DOWN_INCREMENT_SECS: u64 = 5;
|
||||
|
||||
/// The wire behind a device-code login. The `Kimi` arm calls the bespoke Kimi
|
||||
/// Code wire (X-Msh headers, `/api/oauth/*`) verbatim — byte-identical to the
|
||||
/// pre-generalization path; the `Generic` arm drives a registry
|
||||
/// [`OAuthConfig`] provider (xai-grok) through [`crate::auth::oauth_device`].
|
||||
enum DeviceFlowBackend<'a> {
|
||||
Kimi { host: &'a str },
|
||||
Generic(&'a OAuthConfig),
|
||||
}
|
||||
|
||||
impl DeviceFlowBackend<'_> {
|
||||
async fn request(&self) -> anyhow::Result<DeviceAuthorization> {
|
||||
match self {
|
||||
Self::Kimi { host } => {
|
||||
crate::auth::kimi_oauth::request_device_authorization(host).await
|
||||
}
|
||||
Self::Generic(cfg) => {
|
||||
crate::auth::oauth_device::request_device_authorization(cfg).await
|
||||
}
|
||||
}
|
||||
}
|
||||
async fn poll(&self, device_code: &str) -> anyhow::Result<DevicePollResult> {
|
||||
match self {
|
||||
Self::Kimi { host } => {
|
||||
crate::auth::kimi_oauth::poll_device_token(host, device_code).await
|
||||
}
|
||||
Self::Generic(cfg) => {
|
||||
crate::auth::oauth_device::poll_device_token(cfg, device_code).await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Outcome of one full poll loop over a single device authorization.
|
||||
enum PollLoopOutcome {
|
||||
/// Access token issued.
|
||||
@@ -40,11 +72,29 @@ pub async fn run_device_code_login_channels(
|
||||
host: &str,
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
channels: &mut Option<AuthChannels>,
|
||||
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||
run_device_code_login_backend(DeviceFlowBackend::Kimi { host }, auth_manager, channels).await
|
||||
}
|
||||
|
||||
/// Device-code login for a GENERIC [`OAuthConfig`] provider (xai-grok). Same
|
||||
/// TUI/CLI presentation as the Kimi login; only the wire differs.
|
||||
pub async fn run_device_code_login_generic(
|
||||
oauth: &OAuthConfig,
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
channels: &mut Option<AuthChannels>,
|
||||
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||
run_device_code_login_backend(DeviceFlowBackend::Generic(oauth), auth_manager, channels).await
|
||||
}
|
||||
|
||||
async fn run_device_code_login_backend(
|
||||
backend: DeviceFlowBackend<'_>,
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
channels: &mut Option<AuthChannels>,
|
||||
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||
let interactive_tui = channels.is_some();
|
||||
let mut channels = channels.take();
|
||||
loop {
|
||||
let device_auth = request_device_authorization(host).await?;
|
||||
let device_auth = backend.request().await?;
|
||||
let display_uri = device_auth.verification_uri_complete.clone();
|
||||
|
||||
if interactive_tui {
|
||||
@@ -62,7 +112,7 @@ pub async fn run_device_code_login_channels(
|
||||
prompt_on_stderr(&device_auth).await;
|
||||
}
|
||||
|
||||
match complete_device_code_login(host, &device_auth).await? {
|
||||
match complete_device_code_login(&backend, &device_auth).await? {
|
||||
PollLoopOutcome::Done(auth) => {
|
||||
let auth = auth_manager
|
||||
.update(*auth)
|
||||
@@ -112,7 +162,7 @@ async fn prompt_on_stderr(device_auth: &DeviceAuthorization) {
|
||||
/// Poll the token endpoint until the user approves, the device code expires
|
||||
/// (→ [`PollLoopOutcome::Restart`]), or the wire fails.
|
||||
async fn complete_device_code_login(
|
||||
host: &str,
|
||||
backend: &DeviceFlowBackend<'_>,
|
||||
device_auth: &DeviceAuthorization,
|
||||
) -> anyhow::Result<PollLoopOutcome> {
|
||||
let mut poll_interval = std::time::Duration::from_secs(device_auth.interval.max(1) as u64);
|
||||
@@ -120,7 +170,7 @@ async fn complete_device_code_login(
|
||||
// Sleep first: an immediate poll on a fresh code only returns
|
||||
// authorization_pending (and risks slow_down).
|
||||
tokio::time::sleep(poll_interval).await;
|
||||
match poll_device_token(host, &device_auth.device_code).await? {
|
||||
match backend.poll(&device_auth.device_code).await? {
|
||||
DevicePollResult::Success(auth) => {
|
||||
tracing::info!("auth: device login authorized");
|
||||
return Ok(PollLoopOutcome::Done(auth));
|
||||
|
||||
@@ -67,6 +67,37 @@ pub async fn run_auth_flow(
|
||||
run_auth_flow_inner(auth_manager, kimi_code_config, reauth, false, channels).await
|
||||
}
|
||||
|
||||
/// Login flow for a GENERIC device-code OAuth provider (xai-grok): use a valid
|
||||
/// cached session unless re-authing, otherwise run the generic device flow
|
||||
/// (persisting under the provider's own scope via `auth_manager`). Unlike the
|
||||
/// Kimi flow this does not run the silent-refresh dance — the device flow's
|
||||
/// `AuthManager::update` persists a fresh token set directly.
|
||||
pub async fn run_oauth_provider_flow(
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
oauth: &'static kigi_models::OAuthConfig,
|
||||
reauth: bool,
|
||||
channels: Option<AuthChannels>,
|
||||
) -> anyhow::Result<(KimiAuth, bool)> {
|
||||
tracing::info!(
|
||||
scope_key = oauth.scope_key,
|
||||
reauth,
|
||||
"auth: starting generic oauth login"
|
||||
);
|
||||
if reauth {
|
||||
auth_manager.clear()?;
|
||||
}
|
||||
if !reauth && let Some(auth) = auth_manager.current() {
|
||||
tracing::info!(
|
||||
scope_key = oauth.scope_key,
|
||||
"auth: using cached oauth session"
|
||||
);
|
||||
return Ok((auth, false));
|
||||
}
|
||||
let mut channels = channels;
|
||||
crate::auth::device_code::run_device_code_login_generic(oauth, auth_manager, &mut channels)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_auth_flow_inner(
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
_kimi_code_config: &KimiCodeConfig,
|
||||
|
||||
@@ -130,8 +130,9 @@ fn with_device_headers(
|
||||
}
|
||||
|
||||
/// Defend against control characters / non-https redirects from a
|
||||
/// compromised or mis-configured OAuth host.
|
||||
fn validate_verification_uri(uri: &str) -> anyhow::Result<()> {
|
||||
/// compromised or mis-configured OAuth host. Shared with the generic
|
||||
/// device-code wire ([`super::oauth_device`]).
|
||||
pub(crate) fn validate_verification_uri(uri: &str) -> anyhow::Result<()> {
|
||||
if uri.chars().any(|c| c.is_ascii_control()) {
|
||||
anyhow::bail!("Server returned invalid verification URI");
|
||||
}
|
||||
|
||||
@@ -385,6 +385,59 @@ impl AuthManager {
|
||||
)
|
||||
}
|
||||
|
||||
/// Build a manager for a GENERIC device-code OAuth provider (xai-grok),
|
||||
/// scoped to `oauth.scope_key`. Unlike [`Self::new`] this path is
|
||||
/// file-store only (no keyring — that is gated to the default Kimi install)
|
||||
/// and ignores the Kimi-specific `KIGI_AUTH` inline-credential env; it
|
||||
/// otherwise shares the same multi-scope `auth.json` (honoring
|
||||
/// `KIGI_AUTH_PATH`). The refresher is selected from the scope by
|
||||
/// [`super::refresh::build_refresher`].
|
||||
pub(crate) fn new_oauth_provider(
|
||||
kigi_home: &Path,
|
||||
oauth: &'static kigi_models::OAuthConfig,
|
||||
) -> Self {
|
||||
let scope = oauth.scope_key.to_owned();
|
||||
let path = std::env::var("KIGI_AUTH_PATH")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|_| kigi_home.join("auth.json"));
|
||||
|
||||
let (auth, disk_state) = match read_auth_json(&path) {
|
||||
Ok(map) => {
|
||||
let found = lookup_auth(&map, &scope);
|
||||
let state = if found.is_some() {
|
||||
DiskAuthState::Ok
|
||||
} else {
|
||||
DiskAuthState::EntryMissing
|
||||
};
|
||||
(found, state)
|
||||
}
|
||||
Err(e) => {
|
||||
let state = if e.kind() == std::io::ErrorKind::NotFound {
|
||||
DiskAuthState::FileMissing
|
||||
} else {
|
||||
DiskAuthState::Unreadable
|
||||
};
|
||||
(None, state)
|
||||
}
|
||||
};
|
||||
kigi_log::unified_log::info(
|
||||
"AuthManager::new_oauth_provider",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope": &scope,
|
||||
"found": auth.is_some(),
|
||||
"is_expired": auth.as_ref().map(is_expired),
|
||||
})),
|
||||
);
|
||||
Self::assemble(
|
||||
auth,
|
||||
path,
|
||||
scope,
|
||||
KimiCodeConfig::default(),
|
||||
Some(disk_state),
|
||||
)
|
||||
}
|
||||
|
||||
/// Single field-assembly point for [`Self::new`]'s two construction paths
|
||||
/// (inline `KIGI_AUTH` vs. on-disk `auth.json`), which differ only in the
|
||||
/// threaded fields. One literal means a newly added field can't be silently
|
||||
@@ -789,6 +842,13 @@ impl AuthManager {
|
||||
&self.kimi_code_config
|
||||
}
|
||||
|
||||
/// The auth.json / keyring scope key this manager persists under
|
||||
/// (`oauth/kimi-code` for Kimi, `oauth/xai` for xai-grok, …). Drives the
|
||||
/// refresher selection in [`super::refresh::build_refresher`].
|
||||
pub(crate) fn scope(&self) -> &str {
|
||||
&self.scope
|
||||
}
|
||||
|
||||
/// Handle notified after every successful token refresh.
|
||||
///
|
||||
/// Used by [`ModelsManager`] to trigger model catalog recovery
|
||||
|
||||
@@ -8,6 +8,8 @@ mod flow;
|
||||
pub(crate) mod kimi_oauth;
|
||||
pub(crate) mod manager;
|
||||
mod model;
|
||||
pub(crate) mod oauth_device;
|
||||
pub(crate) mod oauth_registry;
|
||||
pub(crate) mod recovery;
|
||||
pub(crate) mod refresh;
|
||||
mod storage;
|
||||
@@ -17,7 +19,8 @@ pub(crate) use flow::try_ensure_session_noninteractive;
|
||||
pub use flow::{
|
||||
AuthChannels, AuthUrlInfo, AuthUrlMode, LogoutResult, ensure_authenticated,
|
||||
ensure_authenticated_or_noninteractive, perform_logout, run_auth_flow,
|
||||
run_auth_flow_with_stderr_bridge, run_cli_login, run_cli_logout, try_ensure_fresh_auth,
|
||||
run_auth_flow_with_stderr_bridge, run_cli_login, run_cli_logout, run_oauth_provider_flow,
|
||||
try_ensure_fresh_auth,
|
||||
};
|
||||
mod meta;
|
||||
pub use device::device_headers;
|
||||
|
||||
@@ -0,0 +1,423 @@
|
||||
//! Generic RFC-8628 device-code OAuth wire, driven by a registry
|
||||
//! [`kigi_models::OAuthConfig`] (xai-grok today; Copilot/Claude later).
|
||||
//!
|
||||
//! Three `application/x-www-form-urlencoded` POSTs against `{auth_host}`:
|
||||
//!
|
||||
//! - `POST {device_path}` — form `client_id` + `scope` + the optional
|
||||
//! `extra_device_field` (e.g. `referrer=kigi`)
|
||||
//! - `POST {token_path}` (poll) — form `client_id` + `device_code` +
|
||||
//! `grant_type=urn:ietf:params:oauth:grant-type:device_code`
|
||||
//! - `POST {token_path}` (refresh) — form `client_id` +
|
||||
//! `grant_type=refresh_token` + `refresh_token`, with the same exponential
|
||||
//! backoff / status handling as the Kimi wire.
|
||||
//!
|
||||
//! Unlike [`super::kimi_oauth`] this sends NO X-Msh device headers — just the
|
||||
//! shared kigi `User-Agent` and `Accept: application/json`. Access/refresh
|
||||
//! tokens are NEVER logged (only non-secret events: requested, poll succeeded,
|
||||
//! refreshed).
|
||||
|
||||
use kigi_models::OAuthConfig;
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::kimi_oauth::{
|
||||
DeviceAuthorization, DevicePollResult, RefreshError, TokenResponse, validate_verification_uri,
|
||||
};
|
||||
|
||||
const DEVICE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code";
|
||||
const REFRESH_GRANT_TYPE: &str = "refresh_token";
|
||||
/// Refresh retry budget over the retryable statuses / network blips.
|
||||
const MAX_REFRESH_RETRIES: u32 = 3;
|
||||
/// HTTP statuses worth retrying a refresh for (kimi-cli parity).
|
||||
const RETRYABLE_REFRESH_STATUSES: [u16; 5] = [429, 500, 502, 503, 504];
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct DeviceAuthorizationResponse {
|
||||
user_code: String,
|
||||
device_code: String,
|
||||
#[serde(default)]
|
||||
verification_uri: Option<String>,
|
||||
/// Optional here (the Kimi wire requires it): Pi's xAI response may omit
|
||||
/// `verification_uri_complete` and carry only `verification_uri`.
|
||||
#[serde(default)]
|
||||
verification_uri_complete: Option<String>,
|
||||
#[serde(default)]
|
||||
expires_in: Option<i64>,
|
||||
#[serde(default)]
|
||||
interval: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Default)]
|
||||
struct OAuthErrorBody {
|
||||
#[serde(default)]
|
||||
error: Option<String>,
|
||||
#[serde(default)]
|
||||
error_description: Option<String>,
|
||||
}
|
||||
|
||||
fn oauth_url(host: &str, path: &str) -> String {
|
||||
format!("{}{path}", host.trim_end_matches('/'))
|
||||
}
|
||||
|
||||
/// The device-authorization form fields: `client_id`, `scope`, and the
|
||||
/// optional non-standard `extra_device_field`.
|
||||
fn device_form(cfg: &OAuthConfig) -> Vec<(&'static str, &'static str)> {
|
||||
let mut form = vec![("client_id", cfg.client_id), ("scope", cfg.scope)];
|
||||
if let Some((name, value)) = cfg.extra_device_field {
|
||||
form.push((name, value));
|
||||
}
|
||||
form
|
||||
}
|
||||
|
||||
/// `POST {auth_host}{device_path}` — start a device login.
|
||||
pub(crate) async fn request_device_authorization(
|
||||
cfg: &OAuthConfig,
|
||||
) -> anyhow::Result<DeviceAuthorization> {
|
||||
let url = oauth_url(cfg.auth_host, cfg.device_path);
|
||||
tracing::info!(url = %url, "auth: requesting device authorization (generic oauth)");
|
||||
let resp = crate::http::shared_client()
|
||||
.post(&url)
|
||||
.header("Accept", "application/json")
|
||||
.form(&device_form(cfg))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
tracing::warn!(%status, "auth: device authorization failed (generic oauth)");
|
||||
anyhow::bail!("Device authorization failed (HTTP {status}): {body}");
|
||||
}
|
||||
let parsed: DeviceAuthorizationResponse = resp.json().await?;
|
||||
|
||||
if !parsed
|
||||
.user_code
|
||||
.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == '-')
|
||||
{
|
||||
anyhow::bail!("Server returned invalid user_code format (expected [A-Z0-9-])");
|
||||
}
|
||||
// Pi forces the displayed URI to https; we require a valid https (or
|
||||
// localhost) verification target, preferring the pre-filled complete form.
|
||||
let verification_uri_complete = parsed
|
||||
.verification_uri_complete
|
||||
.clone()
|
||||
.or_else(|| parsed.verification_uri.clone())
|
||||
.ok_or_else(|| anyhow::anyhow!("Server returned no verification URI"))?;
|
||||
validate_verification_uri(&verification_uri_complete)?;
|
||||
if let Some(ref uri) = parsed.verification_uri {
|
||||
validate_verification_uri(uri)?;
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
user_code = %parsed.user_code,
|
||||
interval = parsed.interval.unwrap_or(5),
|
||||
expires_in = ?parsed.expires_in,
|
||||
"auth: device authorization issued (generic oauth)"
|
||||
);
|
||||
Ok(DeviceAuthorization {
|
||||
user_code: parsed.user_code,
|
||||
device_code: parsed.device_code,
|
||||
verification_uri: parsed.verification_uri.filter(|u| !u.is_empty()),
|
||||
verification_uri_complete,
|
||||
expires_in: parsed.expires_in.filter(|&e| e > 0),
|
||||
interval: parsed.interval.unwrap_or(5),
|
||||
})
|
||||
}
|
||||
|
||||
/// One poll of `POST {auth_host}{token_path}` with the device grant.
|
||||
pub(crate) async fn poll_device_token(
|
||||
cfg: &OAuthConfig,
|
||||
device_code: &str,
|
||||
) -> anyhow::Result<DevicePollResult> {
|
||||
let url = oauth_url(cfg.auth_host, cfg.token_path);
|
||||
let resp = crate::http::shared_client()
|
||||
.post(&url)
|
||||
.header("Accept", "application/json")
|
||||
.form(&[
|
||||
("client_id", cfg.client_id),
|
||||
("device_code", device_code),
|
||||
("grant_type", DEVICE_GRANT_TYPE),
|
||||
])
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Token polling request failed: {e}"))?;
|
||||
|
||||
let status = resp.status();
|
||||
if status.is_server_error() {
|
||||
anyhow::bail!("Token polling server error: {status}");
|
||||
}
|
||||
let body = resp.bytes().await?;
|
||||
if status.is_success() {
|
||||
if let Ok(tokens) = serde_json::from_slice::<TokenResponse>(&body) {
|
||||
tracing::info!("auth: device poll succeeded, access token issued (generic oauth)");
|
||||
return Ok(DevicePollResult::Success(Box::new(tokens.into_auth())));
|
||||
}
|
||||
tracing::warn!(
|
||||
"auth: device poll returned 200 without access_token; continuing (generic oauth)"
|
||||
);
|
||||
return Ok(DevicePollResult::Pending {
|
||||
error: "missing_access_token".to_owned(),
|
||||
description: None,
|
||||
});
|
||||
}
|
||||
let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default();
|
||||
let error = err.error.unwrap_or_else(|| "unknown_error".to_owned());
|
||||
if error == "expired_token" {
|
||||
tracing::info!(
|
||||
"auth: device code expired; restarting device authorization (generic oauth)"
|
||||
);
|
||||
return Ok(DevicePollResult::Expired);
|
||||
}
|
||||
tracing::debug!(error = %error, "auth: device poll pending (generic oauth)");
|
||||
Ok(DevicePollResult::Pending {
|
||||
error,
|
||||
description: err.error_description,
|
||||
})
|
||||
}
|
||||
|
||||
/// `POST {auth_host}{token_path}` with `grant_type=refresh_token`. Retries the
|
||||
/// retryable statuses / network errors with exponential backoff; 401/403
|
||||
/// returns immediately as [`RefreshError::Unauthorized`].
|
||||
pub(crate) async fn refresh_token(
|
||||
cfg: &OAuthConfig,
|
||||
refresh_token: &str,
|
||||
) -> Result<super::model::KimiAuth, RefreshError> {
|
||||
let url = oauth_url(cfg.auth_host, cfg.token_path);
|
||||
let mut last_error = String::from("no attempt made");
|
||||
for attempt in 0..MAX_REFRESH_RETRIES {
|
||||
if attempt > 0 {
|
||||
let backoff = std::time::Duration::from_secs(1 << (attempt - 1));
|
||||
tracing::warn!(
|
||||
attempt,
|
||||
backoff_secs = backoff.as_secs(),
|
||||
last_error = %last_error,
|
||||
"auth: retrying token refresh (generic oauth)"
|
||||
);
|
||||
tokio::time::sleep(backoff).await;
|
||||
}
|
||||
tracing::info!(attempt, "auth: token refresh attempt (generic oauth)");
|
||||
let send_result = crate::http::shared_client()
|
||||
.post(&url)
|
||||
.header("Accept", "application/json")
|
||||
.form(&[
|
||||
("client_id", cfg.client_id),
|
||||
("grant_type", REFRESH_GRANT_TYPE),
|
||||
("refresh_token", refresh_token),
|
||||
])
|
||||
.send()
|
||||
.await;
|
||||
|
||||
let resp = match send_result {
|
||||
Ok(resp) => resp,
|
||||
Err(e) => {
|
||||
last_error = format!("network error: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let status = resp.status().as_u16();
|
||||
let body = resp.bytes().await.unwrap_or_default();
|
||||
if status == 401 || status == 403 {
|
||||
let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default();
|
||||
return Err(RefreshError::Unauthorized {
|
||||
status,
|
||||
description: err
|
||||
.error_description
|
||||
.unwrap_or_else(|| "Token refresh unauthorized.".to_owned()),
|
||||
});
|
||||
}
|
||||
if status == 200 {
|
||||
return match serde_json::from_slice::<TokenResponse>(&body) {
|
||||
Ok(tokens) => Ok(tokens.into_auth()),
|
||||
Err(e) => Err(RefreshError::Fatal {
|
||||
status,
|
||||
description: format!("malformed token payload: {e}"),
|
||||
}),
|
||||
};
|
||||
}
|
||||
let err: OAuthErrorBody = serde_json::from_slice(&body).unwrap_or_default();
|
||||
let description = err
|
||||
.error_description
|
||||
.unwrap_or_else(|| format!("Token refresh failed (HTTP {status})."));
|
||||
if RETRYABLE_REFRESH_STATUSES.contains(&status) {
|
||||
last_error = description;
|
||||
continue;
|
||||
}
|
||||
return Err(RefreshError::Fatal {
|
||||
status,
|
||||
description,
|
||||
});
|
||||
}
|
||||
Err(RefreshError::Exhausted { last_error })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_models::XAI_OAUTH_CONFIG;
|
||||
use wiremock::matchers::{body_string_contains, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
/// An OAuthConfig pointed at a mock server (copies XAI's client_id/scope/
|
||||
/// paths but overrides the host).
|
||||
fn mock_cfg(host: &'static str) -> OAuthConfig {
|
||||
OAuthConfig {
|
||||
auth_host: host,
|
||||
..XAI_OAUTH_CONFIG
|
||||
}
|
||||
}
|
||||
|
||||
fn token_json(access: &str, refresh: &str) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"access_token": access,
|
||||
"refresh_token": refresh,
|
||||
"expires_in": 3600,
|
||||
"scope": "grok-cli:access",
|
||||
"token_type": "bearer",
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn device_authorization_sends_client_scope_and_referrer() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/device/code"))
|
||||
.and(body_string_contains(
|
||||
"client_id=b1a00492-073a-47ea-816f-4c329264a828",
|
||||
))
|
||||
.and(body_string_contains("scope=openid"))
|
||||
.and(body_string_contains("referrer=kigi"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"user_code": "GROK-1234",
|
||||
"device_code": "dev-xai-1",
|
||||
"verification_uri": "https://x.ai/device",
|
||||
"verification_uri_complete": "https://x.ai/device?user_code=GROK-1234",
|
||||
"expires_in": 900,
|
||||
"interval": 5,
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let auth = request_device_authorization(&mock_cfg(host)).await.unwrap();
|
||||
assert_eq!(auth.user_code, "GROK-1234");
|
||||
assert_eq!(auth.device_code, "dev-xai-1");
|
||||
assert_eq!(
|
||||
auth.verification_uri_complete,
|
||||
"https://x.ai/device?user_code=GROK-1234"
|
||||
);
|
||||
assert_eq!(auth.expires_in, Some(900));
|
||||
}
|
||||
|
||||
/// A response with only `verification_uri` (no `_complete`) still yields a
|
||||
/// valid display URI.
|
||||
#[tokio::test]
|
||||
async fn device_authorization_falls_back_to_verification_uri() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/device/code"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"user_code": "GROK-9",
|
||||
"device_code": "d",
|
||||
"verification_uri": "https://x.ai/device",
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let auth = request_device_authorization(&mock_cfg(host)).await.unwrap();
|
||||
assert_eq!(auth.verification_uri_complete, "https://x.ai/device");
|
||||
assert_eq!(auth.interval, 5, "default interval");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn poll_success_builds_auth_with_expiry() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.and(body_string_contains("grant_type=urn"))
|
||||
.and(body_string_contains("device_code=dev-xai-1"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(token_json("grok-at", "grok-rt")),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = poll_device_token(&mock_cfg(host), "dev-xai-1")
|
||||
.await
|
||||
.unwrap();
|
||||
let DevicePollResult::Success(auth) = result else {
|
||||
panic!("expected success, got {result:?}");
|
||||
};
|
||||
assert_eq!(auth.key, "grok-at");
|
||||
assert_eq!(auth.refresh_token.as_deref(), Some("grok-rt"));
|
||||
assert_eq!(auth.expires_in, Some(3600));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn poll_maps_authorization_pending_to_pending() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(400)
|
||||
.set_body_json(serde_json::json!({ "error": "authorization_pending" })),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = poll_device_token(&mock_cfg(host), "dev-xai-1")
|
||||
.await
|
||||
.unwrap();
|
||||
match result {
|
||||
DevicePollResult::Pending { error, .. } => assert_eq!(error, "authorization_pending"),
|
||||
other => panic!("expected pending, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_success_round_trip() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.and(body_string_contains("grant_type=refresh_token"))
|
||||
.and(body_string_contains("refresh_token=grok-rt-old"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(token_json("grok-at-new", "grok-rt-new")),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let auth = refresh_token(&mock_cfg(host), "grok-rt-old").await.unwrap();
|
||||
assert_eq!(auth.key, "grok-at-new");
|
||||
assert_eq!(auth.refresh_token.as_deref(), Some("grok-rt-new"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_401_maps_to_unauthorized() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(401)
|
||||
.set_body_json(serde_json::json!({ "error_description": "refresh revoked" })),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let err = refresh_token(&mock_cfg(host), "grok-rt-dead")
|
||||
.await
|
||||
.unwrap_err();
|
||||
match err {
|
||||
RefreshError::Unauthorized {
|
||||
status,
|
||||
description,
|
||||
} => {
|
||||
assert_eq!(status, 401);
|
||||
assert_eq!(description, "refresh revoked");
|
||||
}
|
||||
other => panic!("expected Unauthorized, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,257 @@
|
||||
//! Process-global per-provider OAuth `AuthManager` pool for INFERENCE-time auth.
|
||||
//!
|
||||
//! A session binds its primary (Kimi / first-party) [`AuthManager`] for the
|
||||
//! subscription path, but a `uses_oauth` platform that carries an
|
||||
//! [`kigi_models::OAuthConfig`] (xai-grok today) needs its OWN scope-keyed
|
||||
//! manager for every per-turn decision — bearer resolution, proactive /
|
||||
//! on-expiry refresh, and 401 recovery. Reusing the Kimi manager for a grok
|
||||
//! turn would transmit the Kimi subscription bearer to `api.x.ai` (a
|
||||
//! cross-provider leak, guaranteed 401) and, without proactive refresh, would
|
||||
//! 401 every turn once the ~1h grok token expired until a process restart.
|
||||
//!
|
||||
//! The pool is the SINGLE SOURCE OF TRUTH: one long-lived `AuthManager` per
|
||||
//! generic-oauth scope, each wired with the SAME lifecycle as the primary Kimi
|
||||
//! manager (`configure_refresher()` + `start_proactive_refresh()`) so the
|
||||
//! on-disk token stays fresh and a 401 recovers via the provider's own manager.
|
||||
//! Managers are built ON DEMAND: the first grok turn (or model switch) reads the
|
||||
//! on-disk token via [`global_manager_for`], so a login that lands AFTER a
|
||||
//! session spawned self-heals — there is no frozen per-session snapshot to go
|
||||
//! stale. [`manager_for_model`] routes a managed catalog key to the pool (oauth
|
||||
//! platform) or to the session's primary (everything else).
|
||||
//!
|
||||
//! SECURITY: access/refresh tokens and resolved bearers are NEVER logged here.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use parking_lot::Mutex;
|
||||
|
||||
use crate::auth::AuthManager;
|
||||
|
||||
/// Process-wide pool of live per-scope OAuth managers.
|
||||
///
|
||||
/// Auth is process-global (one user), so a single manager per scope is correct
|
||||
/// and lets the proactive-refresh task start exactly once per scope no matter
|
||||
/// how many sessions spawn. Keyed by the OAuth `scope_key` (`oauth/xai`, …).
|
||||
fn oauth_manager_pool() -> &'static Mutex<HashMap<&'static str, Arc<AuthManager>>> {
|
||||
static POOL: OnceLock<Mutex<HashMap<&'static str, Arc<AuthManager>>>> = OnceLock::new();
|
||||
POOL.get_or_init(|| Mutex::new(HashMap::new()))
|
||||
}
|
||||
|
||||
/// Get-or-create the process-global manager for `oauth`, wiring the same
|
||||
/// refresher + proactive-refresh lifecycle as the primary Kimi manager the
|
||||
/// FIRST time a scope is seen. The manager reads the on-disk token at
|
||||
/// construction (thereafter kept fresh by the proactive-refresh loop), so a
|
||||
/// grok login that lands after this scope was first built is adopted on the
|
||||
/// manager's own refresh tick — no session ever needs re-spawning.
|
||||
///
|
||||
/// MUST be called from within a Tokio runtime (the proactive-refresh loop
|
||||
/// spawns a task, mirroring the primary).
|
||||
pub(crate) fn global_manager_for(
|
||||
kigi_home: &Path,
|
||||
oauth: &'static kigi_models::OAuthConfig,
|
||||
) -> Arc<AuthManager> {
|
||||
let mut pool = oauth_manager_pool().lock();
|
||||
if let Some(existing) = pool.get(oauth.scope_key) {
|
||||
return existing.clone();
|
||||
}
|
||||
let manager = Arc::new(AuthManager::new_oauth_provider(kigi_home, oauth));
|
||||
manager.configure_refresher();
|
||||
// Never-cancelled token = process-lifetime, matching the api-server /
|
||||
// per-session eager-refresh sites that pass a fresh token.
|
||||
manager.start_proactive_refresh(tokio_util::sync::CancellationToken::new());
|
||||
pool.insert(oauth.scope_key, manager.clone());
|
||||
manager
|
||||
}
|
||||
|
||||
/// The `AuthManager` that governs INFERENCE auth for `managed_key`
|
||||
/// (`{platform}/{model}`, e.g. `xai-grok/grok-4-latest`).
|
||||
///
|
||||
/// A generic device-code OAuth platform routes to ITS OWN scope-keyed manager
|
||||
/// from the process-global pool ([`global_manager_for`], built on demand from
|
||||
/// the on-disk token); every other key (Kimi, API-key platforms, `[model.*]`
|
||||
/// entries, or an unprefixed bare id) routes to `primary`.
|
||||
///
|
||||
/// The pool is the single source of truth — there is no per-session snapshot to
|
||||
/// freeze at spawn, so a grok login that happens AFTER a session spawned is
|
||||
/// resolved correctly on the next grok turn. A grok key NEVER resolves to
|
||||
/// `primary`: even before the user logs into grok the pooled manager simply
|
||||
/// holds no token (its bearer / api_key is then `None`), so the Kimi
|
||||
/// subscription bearer can never reach a third-party host — fail-fast, never a
|
||||
/// silent fallback to the Kimi manager.
|
||||
pub(crate) fn manager_for_model(
|
||||
kigi_home: &Path,
|
||||
managed_key: &str,
|
||||
primary: Option<&Arc<AuthManager>>,
|
||||
) -> Option<Arc<AuthManager>> {
|
||||
if let Some((platform, _)) = kigi_models::parse_managed_model_key(managed_key)
|
||||
&& let Some(oauth) = platform.oauth()
|
||||
{
|
||||
return Some(global_manager_for(kigi_home, oauth));
|
||||
}
|
||||
primary.cloned()
|
||||
}
|
||||
|
||||
/// The SESSION token (the raw bearer/key string) that governs INFERENCE auth
|
||||
/// for `managed_key`, resolved by the model's OWN platform. Thin wrapper over
|
||||
/// [`manager_for_model`] used by the aux-model and subagent-override wire paths
|
||||
/// so a `{platform}/{model}` key never receives the primary token of a
|
||||
/// DIFFERENT provider.
|
||||
///
|
||||
/// A generic device-code OAuth platform (xai-grok) draws its token from ITS OWN
|
||||
/// pooled manager; when that provider has no stored session the result is
|
||||
/// `None` — NEVER the primary Kimi key. Every other key routes to `primary` and
|
||||
/// yields the primary's current-or-expired token, byte-identical to reading it
|
||||
/// directly. SECURITY: the resolved token is never logged.
|
||||
pub(crate) fn session_key_for_model(
|
||||
kigi_home: &Path,
|
||||
managed_key: &str,
|
||||
primary: Option<&Arc<AuthManager>>,
|
||||
) -> Option<String> {
|
||||
manager_for_model(kigi_home, managed_key, primary)
|
||||
.and_then(|am| am.current_or_expired())
|
||||
.map(|a| a.key)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::auth::KimiCodeConfig;
|
||||
use crate::auth::{AuthMode, KimiAuth};
|
||||
|
||||
/// A Kimi manager holding a fixed in-memory bearer, standing in for a
|
||||
/// session's primary. The `TempDir` is returned so the caller keeps it
|
||||
/// alive; the token is read from memory (`current_or_expired`), so disk
|
||||
/// contents are irrelevant to the assertion.
|
||||
fn primary_with_token(key: &str) -> (tempfile::TempDir, Arc<AuthManager>) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let manager = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
|
||||
manager.hot_swap(KimiAuth {
|
||||
key: key.to_string(),
|
||||
auth_mode: AuthMode::OAuth,
|
||||
..KimiAuth::test_default()
|
||||
});
|
||||
(dir, manager)
|
||||
}
|
||||
|
||||
fn xai_oauth() -> &'static kigi_models::OAuthConfig {
|
||||
kigi_models::PlatformId::XaiGrok
|
||||
.oauth()
|
||||
.expect("xai-grok carries an OAuthConfig")
|
||||
}
|
||||
|
||||
/// A non-OAuth managed key (moonshot-cn/…) and an unprefixed bare id both
|
||||
/// route to the primary Kimi manager — the Kimi / first-party path is
|
||||
/// untouched and never consults the pool (no runtime needed).
|
||||
#[test]
|
||||
fn non_oauth_and_bare_models_route_to_primary() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
for key in ["moonshot-cn/kimi-k2", "kimi-k2-0905-preview"] {
|
||||
let resolved = manager_for_model(home.path(), key, Some(&kimi))
|
||||
.expect("non-oauth key routes to the primary");
|
||||
assert!(
|
||||
Arc::ptr_eq(&resolved, &kimi),
|
||||
"{key} must resolve to the primary manager"
|
||||
);
|
||||
assert_eq!(resolved.current_or_expired().unwrap().key, "kimi-tok");
|
||||
}
|
||||
}
|
||||
|
||||
/// The primary being `None` (test / BYOK sessions) still yields `None` for a
|
||||
/// non-oauth key, never a panic — and without touching the pool.
|
||||
#[test]
|
||||
fn none_primary_is_passed_through_for_non_oauth() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
assert!(manager_for_model(home.path(), "kimi-k2", None).is_none());
|
||||
}
|
||||
|
||||
/// An `xai-grok/<model>` turn resolves to the process-global pooled xai
|
||||
/// manager, NEVER the primary Kimi manager — the pool is the single source.
|
||||
#[tokio::test]
|
||||
async fn grok_model_resolves_to_pooled_xai_manager_not_kimi() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi))
|
||||
.expect("grok model resolves to the pooled xai manager");
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &kimi),
|
||||
"grok model must NOT resolve to the Kimi manager"
|
||||
);
|
||||
assert!(
|
||||
Arc::ptr_eq(&resolved, &global_manager_for(home.path(), xai_oauth())),
|
||||
"grok model must resolve to the process-global pooled xai manager"
|
||||
);
|
||||
}
|
||||
|
||||
/// Facet B guard: the resolver routes purely by the model's platform, with
|
||||
/// no auth-method input — so even when the session's primary is a Kimi
|
||||
/// (session) manager holding "kimi-tok", a grok model never resolves that
|
||||
/// Kimi token.
|
||||
#[tokio::test]
|
||||
async fn grok_model_under_kimi_primary_never_yields_kimi_token() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-fast", Some(&kimi))
|
||||
.expect("grok model resolves to its own pooled manager regardless of primary");
|
||||
assert!(!Arc::ptr_eq(&resolved, &kimi));
|
||||
assert_ne!(
|
||||
resolved.current_or_expired().map(|a| a.key),
|
||||
Some("kimi-tok".to_string()),
|
||||
"the Kimi bearer must never be what a grok turn resolves"
|
||||
);
|
||||
}
|
||||
|
||||
/// Fail-fast: a grok key resolves to the pooled xai manager (never the Kimi
|
||||
/// primary) even with no stored grok session in the pool — the pooled
|
||||
/// manager then simply holds no token, so nothing (least of all the Kimi
|
||||
/// bearer) is sent to api.x.ai.
|
||||
#[tokio::test]
|
||||
async fn grok_never_falls_back_to_kimi_primary() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let resolved = manager_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi))
|
||||
.expect("grok routes to the pooled xai manager, not None");
|
||||
assert!(
|
||||
!Arc::ptr_eq(&resolved, &kimi),
|
||||
"an OAuth platform must never fall back to the primary Kimi manager"
|
||||
);
|
||||
}
|
||||
|
||||
/// `session_key_for_model`: a non-oauth / bare key yields the primary Kimi
|
||||
/// token exactly as reading it directly would — byte-identical to the
|
||||
/// pre-fix aux/override wire path (no runtime / pool touched).
|
||||
#[test]
|
||||
fn session_key_for_non_oauth_is_the_primary_token() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
for key in ["moonshot-cn/kimi-k2", "kimi-k2-0905-preview"] {
|
||||
assert_eq!(
|
||||
session_key_for_model(home.path(), key, Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"{key} (non-oauth) must yield the primary token unchanged"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// LEAK guard (aux-model + subagent-override token routing): a grok key with
|
||||
/// a Kimi primary NEVER yields the primary Kimi token — it draws from the
|
||||
/// pooled xai manager (its own token, or `None`). This is the exact source
|
||||
/// the aux `session_key` and the override `session_key` now use.
|
||||
#[tokio::test]
|
||||
async fn session_key_for_grok_is_never_the_kimi_primary() {
|
||||
let (_kd, kimi) = primary_with_token("kimi-tok");
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
assert_ne!(
|
||||
session_key_for_model(home.path(), "xai-grok/grok-4-latest", Some(&kimi)),
|
||||
Some("kimi-tok".to_string()),
|
||||
"a grok aux/override model must never receive the primary Kimi session token"
|
||||
);
|
||||
// Even with `None` primary the routing is unchanged: grok → pool, never a panic.
|
||||
assert_ne!(
|
||||
session_key_for_model(home.path(), "xai-grok/grok-4-fast", None),
|
||||
Some("kimi-tok".to_string()),
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
//! Generic device-code token refresher: drives `POST {token_path}` with
|
||||
//! `grant_type=refresh_token` for any [`kigi_models::OAuthConfig`] provider
|
||||
//! (xai-grok today) through the [`TokenRefresher`] seam.
|
||||
//!
|
||||
//! Structurally identical to [`super::kimi_refresher::KimiRefresher`] — same
|
||||
//! sibling-adoption + post-401 grace — but the wire call goes through
|
||||
//! [`crate::auth::oauth_device`] (no X-Msh headers) instead of the Kimi wire.
|
||||
//! Access/refresh tokens are NEVER logged.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use kigi_models::OAuthConfig;
|
||||
|
||||
use crate::auth::error::RefreshTokenFailedReason;
|
||||
use crate::auth::kimi_oauth::RefreshError;
|
||||
use crate::auth::manager::RefreshReason;
|
||||
use crate::auth::oauth_device::{self};
|
||||
|
||||
use super::{AuthSnapshot, RefreshOutcome, TokenRefresher};
|
||||
|
||||
/// Grace period after a 401/403 before concluding the refresh token is dead:
|
||||
/// a concurrent instance may still be persisting its rotated token.
|
||||
const POST_UNAUTHORIZED_GRACE: std::time::Duration = std::time::Duration::from_secs(1);
|
||||
|
||||
pub(crate) struct GenericDeviceRefresher {
|
||||
auth: Arc<dyn AuthSnapshot>,
|
||||
cfg: &'static OAuthConfig,
|
||||
}
|
||||
|
||||
impl GenericDeviceRefresher {
|
||||
pub(crate) fn new(auth: Arc<dyn AuthSnapshot>, cfg: &'static OAuthConfig) -> Self {
|
||||
Self { auth, cfg }
|
||||
}
|
||||
|
||||
/// Post-401 sibling check: wait a beat, re-read the persisted credential,
|
||||
/// and adopt it when its refresh token differs from the rejected one.
|
||||
async fn adopt_rotation_after_unauthorized(&self, tried_rt: &str) -> Option<RefreshOutcome> {
|
||||
tokio::time::sleep(POST_UNAUTHORIZED_GRACE).await;
|
||||
let latest = self.auth.read_disk_auth()?;
|
||||
let latest_rt = latest.refresh_token.as_deref()?;
|
||||
if latest_rt == tried_rt {
|
||||
return None;
|
||||
}
|
||||
kigi_log::unified_log::info(
|
||||
"auth.refresh.adopted_rotation_after_401",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope_key": self.cfg.scope_key,
|
||||
"adopted_rt_prefix": crate::auth::token_suffix(latest_rt),
|
||||
"rejected_rt_prefix": crate::auth::token_suffix(tried_rt),
|
||||
})),
|
||||
);
|
||||
Some(RefreshOutcome::success(latest))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl TokenRefresher for GenericDeviceRefresher {
|
||||
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome {
|
||||
tracing::info!(
|
||||
?reason,
|
||||
scope_key = self.cfg.scope_key,
|
||||
"auth: generic refresh attempt"
|
||||
);
|
||||
|
||||
let disk_auth = self.auth.read_disk_auth();
|
||||
|
||||
// Sibling short-circuit: a valid persisted token whose key differs from
|
||||
// in-memory means another process refreshed already — adopt directly.
|
||||
if let Some(ref d) = disk_auth
|
||||
&& !crate::auth::is_expired(d)
|
||||
&& self.auth.current().map(|a| a.key).as_deref() != Some(&d.key)
|
||||
{
|
||||
kigi_log::unified_log::info(
|
||||
"auth.refresh.adopted_sibling_token",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope_key": self.cfg.scope_key,
|
||||
"disk_key_prefix": crate::auth::token_suffix(&d.key),
|
||||
})),
|
||||
);
|
||||
return RefreshOutcome::success(d.clone());
|
||||
}
|
||||
|
||||
let Some(auth) = super::resolve_refresh_credential(self.auth.as_ref(), disk_auth, reason)
|
||||
else {
|
||||
tracing::warn!(
|
||||
?reason,
|
||||
"auth: no credential available for refresh (generic)"
|
||||
);
|
||||
return RefreshOutcome::transient("no token with refresh_token available");
|
||||
};
|
||||
let Some(refresh_token) = auth.refresh_token.clone() else {
|
||||
tracing::warn!(
|
||||
?reason,
|
||||
"auth: resolved credential has no refresh token (generic)"
|
||||
);
|
||||
return RefreshOutcome::transient("credential has no refresh token");
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
rt_prefix = crate::auth::token_suffix(&refresh_token),
|
||||
expires_at = ?auth.expires_at,
|
||||
"auth: sending refresh_token grant (generic oauth)"
|
||||
);
|
||||
|
||||
match oauth_device::refresh_token(self.cfg, &refresh_token).await {
|
||||
Ok(new_auth) => {
|
||||
kigi_log::unified_log::info(
|
||||
"auth.refresh.token_rotated",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope_key": self.cfg.scope_key,
|
||||
"new_key_prefix": crate::auth::token_suffix(&new_auth.key),
|
||||
"expires_at": new_auth.expires_at.map(|e| e.to_rfc3339()),
|
||||
})),
|
||||
);
|
||||
RefreshOutcome::success(new_auth)
|
||||
}
|
||||
Err(RefreshError::Unauthorized {
|
||||
status,
|
||||
description,
|
||||
}) => {
|
||||
tracing::warn!(status, %description, "auth: refresh token rejected (generic)");
|
||||
if let Some(adopted) = self.adopt_rotation_after_unauthorized(&refresh_token).await
|
||||
{
|
||||
return adopted;
|
||||
}
|
||||
kigi_log::unified_log::warn(
|
||||
"auth.refresh.unauthorized",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope_key": self.cfg.scope_key,
|
||||
"status": status,
|
||||
"description": description,
|
||||
"rt_prefix": crate::auth::token_suffix(&refresh_token),
|
||||
})),
|
||||
);
|
||||
RefreshOutcome::permanent(
|
||||
RefreshTokenFailedReason::RefreshTokenRejected,
|
||||
Some(refresh_token),
|
||||
)
|
||||
}
|
||||
Err(
|
||||
e @ (RefreshError::Exhausted { .. }
|
||||
| RefreshError::Fatal { .. }
|
||||
| RefreshError::Local(_)),
|
||||
) => {
|
||||
tracing::warn!(error = %e, "auth: refresh attempt failed (transient, generic)");
|
||||
kigi_log::unified_log::warn(
|
||||
"auth.refresh.transient_wire_failure",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"scope_key": self.cfg.scope_key,
|
||||
"error": format!("{e}"),
|
||||
})),
|
||||
);
|
||||
RefreshOutcome::transient(format!("token refresh failed: {e}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::auth::model::KimiAuth;
|
||||
use chrono::{Duration, Utc};
|
||||
use kigi_models::XAI_OAUTH_CONFIG;
|
||||
use parking_lot::Mutex;
|
||||
use wiremock::matchers::{body_string_contains, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
struct FakeSnapshot {
|
||||
current: Mutex<Option<KimiAuth>>,
|
||||
disk: Mutex<Option<KimiAuth>>,
|
||||
}
|
||||
impl FakeSnapshot {
|
||||
fn new(current: Option<KimiAuth>, disk: Option<KimiAuth>) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
current: Mutex::new(current),
|
||||
disk: Mutex::new(disk),
|
||||
})
|
||||
}
|
||||
}
|
||||
impl AuthSnapshot for FakeSnapshot {
|
||||
fn current(&self) -> Option<KimiAuth> {
|
||||
self.current
|
||||
.lock()
|
||||
.clone()
|
||||
.filter(|a| !crate::auth::is_expired(a))
|
||||
}
|
||||
fn expired_auth(&self) -> Option<KimiAuth> {
|
||||
self.current.lock().clone().filter(crate::auth::is_expired)
|
||||
}
|
||||
fn read_disk_auth(&self) -> Option<KimiAuth> {
|
||||
self.disk.lock().clone()
|
||||
}
|
||||
fn is_expired(&self) -> bool {
|
||||
self.current
|
||||
.lock()
|
||||
.as_ref()
|
||||
.is_some_and(crate::auth::is_expired)
|
||||
}
|
||||
}
|
||||
|
||||
fn expired_session(key: &str, rt: &str) -> KimiAuth {
|
||||
KimiAuth {
|
||||
key: key.into(),
|
||||
refresh_token: Some(rt.into()),
|
||||
expires_at: Some(Utc::now() - Duration::hours(1)),
|
||||
expires_in: Some(3600),
|
||||
..KimiAuth::test_default()
|
||||
}
|
||||
}
|
||||
|
||||
fn mock_cfg(host: &'static str) -> OAuthConfig {
|
||||
OAuthConfig {
|
||||
auth_host: host,
|
||||
..XAI_OAUTH_CONFIG
|
||||
}
|
||||
}
|
||||
|
||||
/// A successful refresh rotates the token via the generic wire.
|
||||
#[tokio::test]
|
||||
async fn refresh_success_returns_rotated_token() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.and(body_string_contains("refresh_token=grok-rt-old"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"access_token": "grok-at-new",
|
||||
"refresh_token": "grok-rt-new",
|
||||
"expires_in": 3600,
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let stale = expired_session("grok-at-old", "grok-rt-old");
|
||||
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
|
||||
let cfg: &'static OAuthConfig = Box::leak(Box::new(mock_cfg(host)));
|
||||
let refresher = GenericDeviceRefresher::new(snap, cfg);
|
||||
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
|
||||
let RefreshOutcome::Success(new_auth) = outcome else {
|
||||
panic!("expected success, got {outcome:?}");
|
||||
};
|
||||
assert_eq!(new_auth.key, "grok-at-new");
|
||||
assert_eq!(new_auth.refresh_token.as_deref(), Some("grok-rt-new"));
|
||||
}
|
||||
|
||||
/// A 401 on refresh (with no sibling rotation) tombstones the rejected
|
||||
/// refresh token as a permanent failure.
|
||||
#[tokio::test]
|
||||
async fn unauthorized_is_permanent_failure() {
|
||||
let server = MockServer::start().await;
|
||||
let host: &'static str = Box::leak(server.uri().into_boxed_str());
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth2/token"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(401)
|
||||
.set_body_json(serde_json::json!({ "error_description": "revoked" })),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let stale = expired_session("grok-at-old", "grok-rt-dead");
|
||||
let snap = FakeSnapshot::new(Some(stale.clone()), Some(stale));
|
||||
let cfg: &'static OAuthConfig = Box::leak(Box::new(mock_cfg(host)));
|
||||
let refresher = GenericDeviceRefresher::new(snap, cfg);
|
||||
let outcome = refresher.refresh(RefreshReason::PreRequest).await;
|
||||
let RefreshOutcome::PermanentFailure {
|
||||
error,
|
||||
rejected_refresh_token,
|
||||
} = outcome
|
||||
else {
|
||||
panic!("expected permanent failure, got {outcome:?}");
|
||||
};
|
||||
assert_eq!(error.reason, RefreshTokenFailedReason::RefreshTokenRejected);
|
||||
assert_eq!(rejected_refresh_token.as_deref(), Some("grok-rt-dead"));
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
mod generic_refresher;
|
||||
mod kimi_refresher;
|
||||
|
||||
use std::sync::Arc;
|
||||
@@ -6,6 +7,7 @@ use crate::auth::manager::AuthManager;
|
||||
pub(crate) use crate::auth::manager::RefreshReason;
|
||||
use crate::auth::model::KimiAuth;
|
||||
|
||||
pub(crate) use generic_refresher::GenericDeviceRefresher;
|
||||
pub(crate) use kimi_refresher::KimiRefresher;
|
||||
|
||||
/// Read-only view of `AuthManager` for refreshers. Enforces the
|
||||
@@ -120,8 +122,20 @@ pub(crate) trait TokenRefresher: Send + Sync {
|
||||
async fn refresh(&self, reason: RefreshReason) -> RefreshOutcome;
|
||||
}
|
||||
|
||||
/// Build the production refresher against `kigi_env::oauth_host()`.
|
||||
/// Build the production refresher for this manager's scope. A scope that maps
|
||||
/// to a generic device-code [`kigi_models::OAuthConfig`] (xai-grok) gets the
|
||||
/// [`GenericDeviceRefresher`]; every other scope — Kimi Code, whose registry
|
||||
/// `oauth` field is `None` by design — gets the bespoke [`KimiRefresher`]
|
||||
/// against `kigi_env::oauth_host()`.
|
||||
pub(crate) fn build_refresher(auth_manager: Arc<AuthManager>) -> Arc<dyn TokenRefresher> {
|
||||
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
|
||||
Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host()))
|
||||
match kigi_models::oauth_config_for_scope_key(auth_manager.scope()) {
|
||||
Some(cfg) => {
|
||||
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
|
||||
Arc::new(GenericDeviceRefresher::new(snapshot, cfg))
|
||||
}
|
||||
None => {
|
||||
let snapshot: Arc<dyn AuthSnapshot> = auth_manager;
|
||||
Arc::new(KimiRefresher::new(snapshot, kigi_env::oauth_host()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user