public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Remove Grok OAuth and legacy BYOK support

Browse Source
This commit is contained in:
Ryan Ward
2026-07-28 02:30:22 -05:00
parent 2faeed7ac5
commit 87e0c83e9e
26 changed files with 153 additions and 4140 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/ai/src/api_keys.rs
+11 -495
View File
@@ -1,9 +1,4 @@
use std::time::{Duration, SystemTime};
use galaxyui_core::{Entity, ModelContext, SingletonEntity};
use galaxyui_extras::secure_storage::{self, AppContextExt};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use warp_multi_agent_api as api;
pub use crate::aws_credentials::{AwsCredentials, AwsCredentialsState};
@@ -12,137 +7,13 @@ pub use crate::geap_credentials::{
LoadGeapCredentialsError, GEAP_REFRESH_LEAD_TIME,
};
const SECURE_STORAGE_KEY: &str = "AiApiKeys";
/// Secure-storage key for the connected xAI/Grok subscription's OAuth tokens.
/// Kept separate from [`SECURE_STORAGE_KEY`] because these are OAuth tokens with
/// a refresh lifecycle, not a user-pasted static key.
const GROK_SECURE_STORAGE_KEY: &str = "GrokOAuthTokens";
/// Emitted when user-provided API keys are updated in-memory.
/// Emitted when the manager's stored credentials (AWS Bedrock or Gemini
/// Enterprise) are updated in-memory.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ApiKeyManagerEvent {
KeysUpdated,
}
/// User-provided API keys for AI providers.
///
/// These are used for "Bring Your Own API Key" functionality, allowing
/// users to use their own API keys instead of Warp's.
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct ApiKeys {
pub google: Option<String>,
pub anthropic: Option<String>,
pub openai: Option<String>,
pub open_router: Option<String>,
pub custom_endpoints: Vec<CustomEndpoint>,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct CustomEndpoint {
pub name: String,
pub url: String,
pub api_key: String,
pub models: Vec<CustomEndpointModel>,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct CustomEndpointModel {
pub name: String,
pub alias: Option<String>,
/// Stable identifier used as `ModelConfig.{base,coding,cli_agent,computer_use_agent}` and
/// as the `CustomModelProviders.providers[*].models[*].config_key` on the request wire.
/// Generated as a UUIDv4 at model creation.
pub config_key: String,
}
impl CustomEndpointModel {
/// Picker label: prefer the user-provided alias; fall back to the raw model name
/// so a row is never blank.
pub fn display_label(&self) -> &str {
match self.alias.as_deref() {
Some(alias) if !alias.trim().is_empty() => alias,
_ => &self.name,
}
}
}
impl ApiKeys {
pub fn has_any_key(&self) -> bool {
self.openai.is_some()
|| self.anthropic.is_some()
|| self.google.is_some()
|| self.open_router.is_some()
|| self
.custom_endpoints
.iter()
.any(|endpoint| !endpoint.api_key.trim().is_empty())
}
/// Number of single-provider API keys currently configured (OpenAI,
/// Anthropic, Google, OpenRouter). Custom endpoints are counted separately
/// via `custom_endpoints`.
pub fn provider_key_count(&self) -> usize {
[
&self.openai,
&self.anthropic,
&self.google,
&self.open_router,
]
.into_iter()
.filter(|key| key.as_deref().is_some_and(|v| !v.trim().is_empty()))
.count()
}
}
/// OAuth tokens for a connected xAI / Grok subscription (e.g. SuperGrok).
///
/// Persisted to secure storage under [`GROK_SECURE_STORAGE_KEY`], separate from
/// the BYO [`ApiKeys`] blob because these are OAuth tokens with a refresh
/// lifecycle rather than a user-pasted static key. `crate::grok_subscription`
/// owns refreshing them; this module is the storage and request-injection
/// source of truth that [`ApiKeyManager::api_keys_for_request`] reads from.
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct GrokTokens {
pub access_token: String,
#[serde(default)]
pub refresh_token: Option<String>,
/// Absolute time at which `access_token` expires, if the provider told us.
#[serde(default)]
pub expires_at: Option<SystemTime>,
/// When the user originally connected the subscription (i.e. when the
/// browser OAuth flow completed). Carried over across token refreshes so
/// it keeps reflecting the initial connection, not the latest refresh;
/// surfaced in the settings UI as "Connected on ...". `None` for tokens
/// stored before this field existed.
#[serde(default)]
pub connected_at: Option<SystemTime>,
}
impl GrokTokens {
/// Returns the access token whenever it is non-empty, regardless of
/// expiry. Possibly-expired tokens are still sent so the server stays the
/// final authority on token validity (it rejects truly invalid tokens);
/// `crate::grok_subscription` refreshes (nearly) expired tokens in the
/// background.
pub fn access_token_for_request(&self) -> Option<&str> {
(!self.access_token.trim().is_empty()).then_some(self.access_token.as_str())
}
/// Returns `true` when the token is known to expire within `lead_time` and
/// should be proactively refreshed. Tokens with an unknown expiry never
/// report as needing a refresh (there's no expiry signal to act on).
pub fn needs_refresh(&self, lead_time: Duration) -> bool {
match self.expires_at {
Some(expires_at) => expires_at <= SystemTime::now() + lead_time,
None => false,
}
}
}
/// Controls how AWS credentials are refreshed by [`ApiKeyManager`].
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub enum AwsCredentialsRefreshStrategy {
@@ -159,187 +30,25 @@ pub enum AwsCredentialsRefreshStrategy {
},
}
/// A structure that manages API keys for AI providers.
/// A structure that manages locally-held credentials used to authenticate AI
/// provider requests: AWS Bedrock credentials and Gemini Enterprise (GEAP)
/// credentials.
pub struct ApiKeyManager {
keys: ApiKeys,
/// OAuth tokens for a connected xAI/Grok subscription, if any. Persisted
/// separately from `keys` under [`GROK_SECURE_STORAGE_KEY`];
/// `crate::grok_subscription` keeps these fresh.
grok_tokens: Option<GrokTokens>,
/// Whether background refresh of `grok_tokens` is currently allowed.
/// Mirrors the BYO API key policy, which lives in the app layer; wired in
/// via `ApiKeyManager::set_grok_refresh_allowed` (`crate::grok_subscription`).
#[cfg(not(target_family = "wasm"))]
pub(crate) grok_refresh_allowed: bool,
/// Guards against overlapping Grok token refreshes: the proactive refresh
/// timer and the request-time safety net
/// (`ApiKeyManager::refresh_grok_tokens_if_needed`) can otherwise race.
#[cfg(not(target_family = "wasm"))]
pub(crate) grok_refresh_in_flight: bool,
pub(crate) aws_credentials_state: AwsCredentialsState,
aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy,
/// In-memory Gemini Enterprise (GEAP) credential state.
pub(crate) geap_credentials_state: GeapCredentialsState,
secure_storage_write_version: u64,
grok_secure_storage_write_version: u64,
}
impl ApiKeyManager {
pub fn new(ctx: &mut ModelContext<Self>) -> Self {
let keys = Self::load_keys_from_secure_storage(ctx);
let grok_tokens = Self::load_grok_tokens_from_secure_storage(ctx);
pub fn new(_ctx: &mut ModelContext<Self>) -> Self {
Self {
keys,
grok_tokens,
#[cfg(not(target_family = "wasm"))]
grok_refresh_allowed: false,
#[cfg(not(target_family = "wasm"))]
grok_refresh_in_flight: false,
aws_credentials_state: AwsCredentialsState::Missing,
aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy::default(),
geap_credentials_state: GeapCredentialsState::Missing,
secure_storage_write_version: 0,
grok_secure_storage_write_version: 0,
}
}
pub fn keys(&self) -> &ApiKeys {
&self.keys
}
/// The currently stored xAI/Grok OAuth tokens, if the user has connected a
/// Grok subscription.
pub fn grok_tokens(&self) -> Option<&GrokTokens> {
self.grok_tokens.as_ref()
}
/// Returns `true` when a Grok subscription is connected with a usable OAuth
/// access token.
pub fn has_grok_subscription(&self) -> bool {
self.grok_tokens
.as_ref()
.and_then(GrokTokens::access_token_for_request)
.is_some()
}
/// Returns `true` when the user has any usable BYO credential: a pasted
/// provider or custom-endpoint key, or a connected Grok subscription.
pub fn has_any_key(&self) -> bool {
self.keys.has_any_key() || self.has_grok_subscription()
}
/// Stores (or clears, with `None`) the xAI/Grok OAuth tokens and persists
/// them to secure storage. No-op when the value is unchanged so we don't
/// emit spurious events or schedule redundant keychain writes.
pub fn set_grok_tokens(&mut self, tokens: Option<GrokTokens>, ctx: &mut ModelContext<Self>) {
if self.grok_tokens == tokens {
return;
}
self.grok_tokens = tokens;
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
self.write_grok_tokens_to_secure_storage(ctx);
}
pub fn set_google_key(&mut self, key: Option<String>, ctx: &mut ModelContext<Self>) {
self.keys.google = key;
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
self.write_keys_to_secure_storage(ctx);
}
pub fn set_anthropic_key(&mut self, key: Option<String>, ctx: &mut ModelContext<Self>) {
self.keys.anthropic = key;
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
self.write_keys_to_secure_storage(ctx);
}
pub fn set_openai_key(&mut self, key: Option<String>, ctx: &mut ModelContext<Self>) {
self.keys.openai = key;
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
self.write_keys_to_secure_storage(ctx);
}
pub fn set_open_router_key(&mut self, key: Option<String>, ctx: &mut ModelContext<Self>) {
self.keys.open_router = key;
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
self.write_keys_to_secure_storage(ctx);
}
pub fn add_custom_endpoint(
&mut self,
name: String,
url: String,
api_key: String,
models: Vec<(String, Option<String>, Option<String>)>,
ctx: &mut ModelContext<Self>,
) {
self.keys.custom_endpoints.push(CustomEndpoint {
name,
url,
api_key,
models: models
.into_iter()
.map(|(name, alias, config_key)| CustomEndpointModel {
name,
alias,
config_key: config_key
.filter(|k| !k.is_empty())
.unwrap_or_else(|| Uuid::new_v4().to_string()),
})
.collect(),
});
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
self.write_keys_to_secure_storage(ctx);
}
pub fn save_custom_endpoint(
&mut self,
index: usize,
name: String,
url: String,
api_key: String,
models: Vec<(String, Option<String>, Option<String>)>,
ctx: &mut ModelContext<Self>,
) {
if index >= self.keys.custom_endpoints.len() {
return;
}
self.keys.custom_endpoints[index] = CustomEndpoint {
name,
url,
api_key,
models: models
.into_iter()
.map(|(name, alias, config_key)| CustomEndpointModel {
name,
alias,
config_key: config_key
.filter(|k| !k.is_empty())
.unwrap_or_else(|| Uuid::new_v4().to_string()),
})
.collect(),
};
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
self.write_keys_to_secure_storage(ctx);
}
pub fn remove_custom_endpoint(&mut self, index: usize, ctx: &mut ModelContext<Self>) {
if index >= self.keys.custom_endpoints.len() {
return;
}
self.keys.custom_endpoints.remove(index);
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
self.write_keys_to_secure_storage(ctx);
}
pub fn clear_custom_endpoints(&mut self, ctx: &mut ModelContext<Self>) {
if self.keys.custom_endpoints.is_empty() {
return;
}
self.keys.custom_endpoints.clear();
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
self.write_keys_to_secure_storage(ctx);
}
pub fn set_aws_credentials_state(
&mut self,
state: AwsCredentialsState,
@@ -380,93 +89,14 @@ impl ApiKeyManager {
self.aws_credentials_refresh_strategy = strategy;
}
/// Builds the `CustomModelProviders` registry that ships with every agent request.
///
/// Emits one [`CustomModelProvider`] per configured [`CustomEndpoint`], each populated with
/// all of its [`CustomEndpointModel`]s. The per-model `config_key` is what the server uses
/// to map a `ModelConfig.{base,coding,cli_agent,computer_use_agent}` selection back to a
/// user-provided endpoint, so it MUST be the same UUID we store locally.
///
/// Returns `None` when custom models should not be included or no endpoint has both a
/// non-empty URL and API key.
pub fn custom_model_providers_for_request(
&self,
include_custom_models: bool,
) -> Option<api::request::settings::CustomModelProviders> {
if !include_custom_models {
return None;
}
let providers: Vec<_> = self
.keys
.custom_endpoints
.iter()
.filter(|endpoint| !endpoint.url.trim().is_empty() && !endpoint.api_key.is_empty())
.map(
|endpoint| api::request::settings::custom_model_providers::CustomModelProvider {
base_url: endpoint.url.clone(),
api_key: endpoint.api_key.clone(),
models: endpoint
.models
.iter()
.filter(|m| !m.name.trim().is_empty() && !m.config_key.is_empty())
.map(
|m| api::request::settings::custom_model_providers::CustomModel {
slug: m.name.clone(),
config_key: m.config_key.clone(),
},
)
.collect(),
},
)
.filter(|provider| !provider.models.is_empty())
.collect();
if providers.is_empty() {
None
} else {
Some(api::request::settings::CustomModelProviders { providers })
}
}
/// Builds the `ApiKeys` request payload carrying AWS Bedrock and/or Gemini
/// Enterprise (GEAP) credentials, when applicable. Returns `None` when
/// neither credential type applies to this request.
pub fn api_keys_for_request(
&self,
include_byo_keys: bool,
include_aws_bedrock_credentials: bool,
geap_binding: Option<GeapMintBinding>,
) -> Option<api::request::settings::ApiKeys> {
let anthropic = include_byo_keys
.then(|| self.keys.anthropic.clone())
.flatten()
.unwrap_or_default();
let openai = include_byo_keys
.then(|| self.keys.openai.clone())
.flatten()
.unwrap_or_default();
let google = include_byo_keys
.then(|| self.keys.google.clone())
.flatten()
.unwrap_or_default();
let open_router = include_byo_keys
.then(|| self.keys.open_router.clone())
.flatten()
.unwrap_or_default();
// The connected Grok subscription's OAuth access token is user-provided
// auth, just like a pasted BYO API key, so it respects the same BYO
// policy gate: when BYO keys are disabled (e.g. by workspace policy),
// the token must not be sent. Possibly-expired tokens ARE sent — the
// server is the authority on validity.
let grok_oauth_access_token = include_byo_keys
.then(|| {
self.grok_tokens
.as_ref()
.and_then(GrokTokens::access_token_for_request)
.map(str::to_owned)
})
.flatten()
.unwrap_or_default();
// Also include credentials when running with OIDC-managed Bedrock inference, regardless
// of the per-user setting flag (which only applies to the local credential chain path).
let include_aws = include_aws_bedrock_credentials
@@ -506,130 +136,16 @@ impl ApiKeyManager {
_ => None,
});
if anthropic.is_empty()
&& openai.is_empty()
&& google.is_empty()
&& open_router.is_empty()
&& grok_oauth_access_token.is_empty()
&& aws_credentials.is_none()
&& google_cloud_credentials.is_none()
{
if aws_credentials.is_none() && google_cloud_credentials.is_none() {
None
} else {
Some(api::request::settings::ApiKeys {
anthropic,
openai,
google,
open_router,
grok_oauth_access_token,
allow_use_of_warp_credits: false,
aws_credentials,
google_cloud_credentials,
..Default::default()
})
}
}
fn load_keys_from_secure_storage(ctx: &mut ModelContext<Self>) -> ApiKeys {
let key_json = match ctx.secure_storage().read_value(SECURE_STORAGE_KEY) {
Ok(json) => json,
Err(e) => {
if !matches!(e, secure_storage::Error::NotFound) {
log::error!("Failed to read API keys from secure storage: {e:#}");
}
return ApiKeys::default();
}
};
match serde_json::from_str(&key_json) {
Ok(keys) => keys,
Err(e) => {
log::error!("Failed to deserialize API keys: {e:#}");
ApiKeys::default()
}
}
}
fn write_keys_to_secure_storage(&mut self, ctx: &mut ModelContext<Self>) {
let json = match serde_json::to_string(&self.keys) {
Ok(json) => json,
Err(e) => {
log::error!("Failed to serialize API keys: {e:#}");
return;
}
};
self.secure_storage_write_version += 1;
let write_version = self.secure_storage_write_version;
// Defer the keychain write so it doesn't block the current event
// processing. The in-memory state is already updated and events
// already emitted, so the UI updates immediately while the
// potentially slow platform secure-storage call runs in a
// subsequent main-thread callback. Skip stale callbacks so older
// writes cannot complete after and overwrite a newer payload.
ctx.spawn(async move { json }, move |me, json, ctx| {
if write_version != me.secure_storage_write_version {
return;
}
if let Err(e) = ctx.secure_storage().write_value(SECURE_STORAGE_KEY, &json) {
log::error!("Failed to write API keys to secure storage: {e:#}");
}
});
}
fn load_grok_tokens_from_secure_storage(ctx: &mut ModelContext<Self>) -> Option<GrokTokens> {
let json = match ctx.secure_storage().read_value(GROK_SECURE_STORAGE_KEY) {
Ok(json) => json,
Err(e) => {
if !matches!(e, secure_storage::Error::NotFound) {
log::error!("Failed to read Grok tokens from secure storage: {e:#}");
}
return None;
}
};
match serde_json::from_str(&json) {
Ok(tokens) => Some(tokens),
Err(e) => {
log::error!("Failed to deserialize Grok tokens: {e:#}");
None
}
}
}
fn write_grok_tokens_to_secure_storage(&mut self, ctx: &mut ModelContext<Self>) {
// `Some(json)` writes the tokens; `None` removes the stored entry (the
// user disconnected). Serialize up front so the deferred callback only
// touches the keychain.
let payload = match self.grok_tokens.as_ref().map(serde_json::to_string) {
Some(Ok(json)) => Some(json),
Some(Err(e)) => {
log::error!("Failed to serialize Grok tokens: {e:#}");
return;
}
None => None,
};
self.grok_secure_storage_write_version += 1;
let write_version = self.grok_secure_storage_write_version;
// Defer the keychain write/remove like `write_keys_to_secure_storage`,
// skipping stale callbacks so an older write can't clobber a newer one.
ctx.spawn(async move { payload }, move |me, payload, ctx| {
if write_version != me.grok_secure_storage_write_version {
return;
}
let result = match payload {
Some(ref json) => ctx
.secure_storage()
.write_value(GROK_SECURE_STORAGE_KEY, json),
None => ctx.secure_storage().remove_value(GROK_SECURE_STORAGE_KEY),
};
if let Err(e) = result {
if !matches!(e, secure_storage::Error::NotFound) {
log::error!("Failed to persist Grok tokens to secure storage: {e:#}");
}
}
});
}
}
impl Entity for ApiKeyManager {
crates/ai/src/api_keys_tests.rs
+58 -528
View File
@@ -2,41 +2,20 @@ use std::time::{Duration, SystemTime};
use super::*;
fn make_manager(keys: ApiKeys) -> ApiKeyManager {
make_manager_with_grok(keys, None)
}
fn make_manager_with_grok(keys: ApiKeys, grok_tokens: Option<GrokTokens>) -> ApiKeyManager {
fn make_manager() -> ApiKeyManager {
ApiKeyManager {
keys,
grok_tokens,
#[cfg(not(target_family = "wasm"))]
grok_refresh_allowed: false,
#[cfg(not(target_family = "wasm"))]
grok_refresh_in_flight: false,
aws_credentials_state: AwsCredentialsState::Missing,
aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy::default(),
geap_credentials_state: GeapCredentialsState::Missing,
secure_storage_write_version: 0,
grok_secure_storage_write_version: 0,
}
}
fn make_manager_with_geap(geap_credentials_state: GeapCredentialsState) -> ApiKeyManager {
let mut manager = make_manager(ApiKeys::default());
let mut manager = make_manager();
manager.geap_credentials_state = geap_credentials_state;
manager
}
fn grok_tokens(access_token: &str, expires_in: Option<u64>) -> GrokTokens {
GrokTokens {
access_token: access_token.into(),
refresh_token: Some("refresh".into()),
expires_at: expires_in.map(|secs| SystemTime::now() + Duration::from_secs(secs)),
connected_at: None,
}
}
fn geap_credentials(access_token: &str, expires_in: Option<u64>) -> GeapCredentials {
GeapCredentials::new(
access_token.into(),
@@ -56,8 +35,6 @@ fn geap_binding() -> GeapMintBinding {
}
}
// The expected binding the request build site passes in is the same type as
// the stored `minted_for`, so the attach check is a plain `==`.
fn geap_gate() -> GeapMintBinding {
geap_binding()
}
@@ -70,487 +47,6 @@ fn geap_loaded(access_token: &str, expires_in: Option<u64>) -> GeapCredentialsSt
}
}
fn endpoint(
name: &str,
url: &str,
api_key: &str,
models: &[(&str, Option<&str>)],
) -> CustomEndpoint {
endpoint_with_keys(
name,
url,
api_key,
&models
.iter()
.enumerate()
.map(|(i, (n, a))| (*n, *a, format!("cfg-{i}")))
.collect::<Vec<_>>()
.iter()
.map(|(n, a, k)| (*n, *a, k.as_str()))
.collect::<Vec<_>>(),
)
}
fn endpoint_with_keys(
name: &str,
url: &str,
api_key: &str,
models: &[(&str, Option<&str>, &str)],
) -> CustomEndpoint {
CustomEndpoint {
name: name.into(),
url: url.into(),
api_key: api_key.into(),
models: models
.iter()
.map(|(n, a, cfg)| CustomEndpointModel {
name: (*n).into(),
alias: a.map(|s| s.into()),
config_key: (*cfg).into(),
})
.collect(),
}
}
// ── serde round-trip ────────────────────────────────────────────
#[test]
fn serde_round_trip_empty() {
let keys = ApiKeys::default();
let json = serde_json::to_string(&keys).unwrap();
let deser: ApiKeys = serde_json::from_str(&json).unwrap();
assert_eq!(keys, deser);
}
#[test]
fn serde_round_trip_with_provider_keys() {
let keys = ApiKeys {
openai: Some("sk-openai".into()),
anthropic: Some("sk-ant-abc".into()),
google: Some("AIzaSy123".into()),
open_router: Some("sk-or-xxx".into()),
custom_endpoints: vec![],
};
let json = serde_json::to_string(&keys).unwrap();
let deser: ApiKeys = serde_json::from_str(&json).unwrap();
assert_eq!(keys, deser);
}
#[test]
fn serde_round_trip_with_custom_endpoints() {
let keys = ApiKeys {
openai: None,
anthropic: None,
google: None,
open_router: None,
custom_endpoints: vec![
endpoint("ep1", "https://a.io/v1", "key1", &[("gpt-4", Some("fast"))]),
endpoint(
"ep2",
"https://b.io/v1",
"key2",
&[("llama-70b", None), ("mixtral", Some("mix"))],
),
],
};
let json = serde_json::to_string(&keys).unwrap();
let deser: ApiKeys = serde_json::from_str(&json).unwrap();
assert_eq!(keys, deser);
}
#[test]
fn serde_ignores_unknown_fields() {
let json = r#"{"openai":"sk-x","unknown_field":"value","custom_endpoints":[]}"#;
let keys: ApiKeys = serde_json::from_str(json).unwrap();
assert_eq!(keys.openai, Some("sk-x".into()));
assert!(keys.custom_endpoints.is_empty());
}
// ── has_any_key ─────────────────────────────────────────────────
#[test]
fn has_any_key_false_when_empty() {
assert!(!ApiKeys::default().has_any_key());
}
#[test]
fn has_any_key_true_for_openai_only() {
let keys = ApiKeys {
openai: Some("sk-x".into()),
..Default::default()
};
assert!(keys.has_any_key());
}
#[test]
fn has_any_key_true_for_custom_endpoints_only() {
let keys = ApiKeys {
custom_endpoints: vec![endpoint("ep", "https://a.io", "key", &[("m", None)])],
..Default::default()
};
assert!(keys.has_any_key());
}
#[test]
fn has_any_key_false_for_endpoint_with_empty_api_key() {
let keys = ApiKeys {
custom_endpoints: vec![endpoint("ep", "https://a.io", "", &[("m", None)])],
..Default::default()
};
assert!(!keys.has_any_key());
}
// ── provider_key_count ─────────────────────────────────────────
#[test]
fn provider_key_count_zero_when_empty() {
assert_eq!(ApiKeys::default().provider_key_count(), 0);
}
#[test]
fn provider_key_count_counts_each_provider_key() {
let keys = ApiKeys {
openai: Some("sk-o".into()),
anthropic: Some("sk-a".into()),
google: Some("AIza".into()),
open_router: Some("sk-or".into()),
custom_endpoints: vec![],
};
assert_eq!(keys.provider_key_count(), 4);
}
#[test]
fn provider_key_count_ignores_blank_keys_and_endpoints() {
let keys = ApiKeys {
openai: Some("sk-o".into()),
anthropic: Some(" ".into()),
google: None,
open_router: None,
custom_endpoints: vec![endpoint("ep", "https://a.io", "k", &[("m", None)])],
};
// Only the non-blank OpenAI key counts; the whitespace Anthropic key and the
// custom endpoint are excluded.
assert_eq!(keys.provider_key_count(), 1);
}
// ── custom_model_providers_for_request ──────────────────────────
#[test]
fn custom_model_providers_none_when_empty() {
let mgr = make_manager(ApiKeys::default());
assert!(mgr.custom_model_providers_for_request(true).is_none());
}
#[test]
fn custom_model_providers_none_when_byo_disabled() {
let mgr = make_manager(ApiKeys {
custom_endpoints: vec![endpoint("ep", "https://a.io", "k", &[("m", None)])],
..Default::default()
});
assert!(mgr.custom_model_providers_for_request(false).is_none());
}
#[test]
fn custom_model_providers_populates_single_endpoint() {
let mgr = make_manager(ApiKeys {
custom_endpoints: vec![endpoint_with_keys(
"My EP",
"https://custom.io/v1",
"ep-key",
&[("big-model", Some("alias"), "uuid-1")],
)],
..Default::default()
});
let result = mgr.custom_model_providers_for_request(true).unwrap();
assert_eq!(result.providers.len(), 1);
let p = &result.providers[0];
assert_eq!(p.base_url, "https://custom.io/v1");
assert_eq!(p.api_key, "ep-key");
assert_eq!(p.models.len(), 1);
assert_eq!(p.models[0].slug, "big-model");
assert_eq!(p.models[0].config_key, "uuid-1");
}
#[test]
fn multiple_endpoints_all_serialize() {
let mgr = make_manager(ApiKeys {
custom_endpoints: vec![
endpoint_with_keys(
"ep1",
"https://a.io",
"k1",
&[("gpt-4", Some("fast"), "uuid-a")],
),
endpoint_with_keys(
"ep2",
"https://b.io",
"k2",
&[
("llama-70b", None, "uuid-b"),
("mixtral", Some("mix"), "uuid-c"),
],
),
],
..Default::default()
});
let result = mgr.custom_model_providers_for_request(true).unwrap();
assert_eq!(result.providers.len(), 2);
assert_eq!(result.providers[0].base_url, "https://a.io");
assert_eq!(result.providers[0].models[0].config_key, "uuid-a");
assert_eq!(result.providers[1].base_url, "https://b.io");
assert_eq!(result.providers[1].models.len(), 2);
assert_eq!(result.providers[1].models[0].slug, "llama-70b");
assert_eq!(result.providers[1].models[0].config_key, "uuid-b");
assert_eq!(result.providers[1].models[1].config_key, "uuid-c");
}
#[test]
fn byok_disabled_returns_none_even_with_endpoints() {
let mgr = make_manager(ApiKeys {
custom_endpoints: vec![endpoint("ep", "https://a.io", "k", &[("m", None)])],
..Default::default()
});
assert!(mgr.custom_model_providers_for_request(false).is_none());
}
#[test]
fn empty_api_key_endpoints_are_skipped() {
let mgr = make_manager(ApiKeys {
custom_endpoints: vec![
endpoint_with_keys("empty", "https://a.io", "", &[("m", None, "uuid-x")]),
endpoint_with_keys("ok", "https://b.io", "k", &[("m", None, "uuid-y")]),
],
..Default::default()
});
let result = mgr.custom_model_providers_for_request(true).unwrap();
assert_eq!(result.providers.len(), 1);
assert_eq!(result.providers[0].base_url, "https://b.io");
}
#[test]
fn endpoints_with_only_empty_models_are_skipped() {
let mgr = make_manager(ApiKeys {
custom_endpoints: vec![endpoint_with_keys(
"ep",
"https://a.io",
"k",
&[("", None, "uuid-z")],
)],
..Default::default()
});
assert!(mgr.custom_model_providers_for_request(true).is_none());
}
// ── display_label fallback ─────────────────────────────────────
#[test]
fn display_label_uses_alias_when_present() {
let m = CustomEndpointModel {
name: "raw-name".into(),
alias: Some("My Alias".into()),
config_key: "k".into(),
};
assert_eq!(m.display_label(), "My Alias");
}
#[test]
fn display_label_falls_back_to_name_when_alias_missing() {
let m = CustomEndpointModel {
name: "raw-name".into(),
alias: None,
config_key: "k".into(),
};
assert_eq!(m.display_label(), "raw-name");
}
#[test]
fn display_label_falls_back_to_name_when_alias_is_whitespace() {
let m = CustomEndpointModel {
name: "raw-name".into(),
alias: Some(" ".into()),
config_key: "k".into(),
};
assert_eq!(m.display_label(), "raw-name");
}
// ── api_keys_for_request ────────────────────────────────────────
#[test]
fn api_keys_for_request_none_when_empty() {
let mgr = make_manager(ApiKeys::default());
assert!(mgr.api_keys_for_request(true, false, None).is_none());
}
#[test]
fn api_keys_for_request_populates_provider_keys() {
let mgr = make_manager(ApiKeys {
openai: Some("sk-o".into()),
anthropic: Some("sk-a".into()),
..Default::default()
});
let result = mgr.api_keys_for_request(true, false, None).unwrap();
assert_eq!(result.openai, "sk-o");
assert_eq!(result.anthropic, "sk-a");
assert!(result.google.is_empty());
}
#[test]
fn api_keys_for_request_omits_keys_when_byo_disabled() {
let mgr = make_manager(ApiKeys {
openai: Some("sk-o".into()),
..Default::default()
});
// With BYO disabled and no other credentials, returns None.
assert!(mgr.api_keys_for_request(false, false, None).is_none());
}
#[test]
fn api_keys_for_request_none_for_custom_endpoints_only() {
let mgr = make_manager(ApiKeys {
custom_endpoints: vec![endpoint("ep", "https://a.io", "k", &[("m", None)])],
..Default::default()
});
assert!(mgr.api_keys_for_request(true, false, None).is_none());
}
// ── grok oauth token ────────────────────────────────────────────
#[test]
fn grok_access_token_present_without_expiry() {
let t = GrokTokens {
access_token: "tok".into(),
..Default::default()
};
assert_eq!(t.access_token_for_request(), Some("tok"));
}
#[test]
fn grok_access_token_blank_is_none() {
let t = GrokTokens {
access_token: " ".into(),
..Default::default()
};
assert_eq!(t.access_token_for_request(), None);
}
#[test]
fn grok_access_token_near_expiry_still_sent() {
// Expired tokens are still sent; the server is the authority on validity.
let t = grok_tokens("tok", Some(0));
assert_eq!(t.access_token_for_request(), Some("tok"));
}
#[test]
fn grok_access_token_far_future_is_some() {
let t = grok_tokens("tok", Some(3600));
assert_eq!(t.access_token_for_request(), Some("tok"));
}
#[test]
fn grok_needs_refresh_within_lead_time() {
assert!(grok_tokens("tok", Some(30)).needs_refresh(Duration::from_secs(300)));
assert!(!grok_tokens("tok", Some(3600)).needs_refresh(Duration::from_secs(300)));
// Expired tokens still need a refresh.
assert!(grok_tokens("tok", Some(0)).needs_refresh(Duration::from_secs(300)));
// Unknown expiry never reports as needing refresh.
assert!(!grok_tokens("tok", None).needs_refresh(Duration::from_secs(300)));
}
#[test]
fn api_keys_for_request_includes_grok_token() {
let mgr = make_manager_with_grok(
ApiKeys::default(),
Some(grok_tokens("grok-abc", Some(3600))),
);
let result = mgr.api_keys_for_request(true, false, None).unwrap();
assert_eq!(result.grok_oauth_access_token, "grok-abc");
assert!(result.anthropic.is_empty());
}
#[test]
fn api_keys_for_request_omits_grok_token_when_byo_disabled() {
// The Grok subscription is user-provided auth, so it follows the BYO
// policy gate: with BYO disabled and no other credentials, returns None.
let mgr = make_manager_with_grok(
ApiKeys::default(),
Some(grok_tokens("grok-abc", Some(3600))),
);
assert!(mgr.api_keys_for_request(false, false, None).is_none());
}
#[test]
fn api_keys_for_request_includes_expired_grok_token() {
// Expired tokens are still sent in requests; the server rejects truly
// invalid ones and the background refresh replaces them.
let mgr = make_manager_with_grok(ApiKeys::default(), Some(grok_tokens("grok-abc", Some(0))));
let result = mgr.api_keys_for_request(true, false, None).unwrap();
assert_eq!(result.grok_oauth_access_token, "grok-abc");
}
#[test]
fn has_grok_subscription_false_when_not_connected() {
let mgr = make_manager(ApiKeys::default());
assert!(!mgr.has_grok_subscription());
}
#[test]
fn has_grok_subscription_true_when_connected() {
let mgr = make_manager_with_grok(
ApiKeys::default(),
Some(grok_tokens("grok-abc", Some(3600))),
);
assert!(mgr.has_grok_subscription());
}
#[test]
fn has_grok_subscription_true_for_expired_token() {
// A connected subscription still counts even when its token is past expiry:
// the token is sent anyway and the server is the authority on validity.
let mgr = make_manager_with_grok(ApiKeys::default(), Some(grok_tokens("grok-abc", Some(0))));
assert!(mgr.has_grok_subscription());
}
#[test]
fn has_grok_subscription_false_when_token_blank() {
// A blank token can't be sent, so it does not count as a usable credential.
let mgr = make_manager_with_grok(ApiKeys::default(), Some(grok_tokens(" ", None)));
assert!(!mgr.has_grok_subscription());
}
// ── ApiKeyManager::has_any_key ──────────────────
#[test]
fn manager_has_any_key_false_when_no_keys_and_no_grok() {
let mgr = make_manager(ApiKeys::default());
assert!(!mgr.has_any_key());
}
#[test]
fn manager_has_any_key_true_for_pasted_key_without_grok() {
let mgr = make_manager(ApiKeys {
openai: Some("sk-x".into()),
..Default::default()
});
assert!(mgr.has_any_key());
}
#[test]
fn manager_has_any_key_true_for_connected_grok_without_pasted_key() {
// The crux: a connected Grok subscription counts even with no pasted keys,
// matching how it's sent as a BYO credential on requests.
let mgr = make_manager_with_grok(
ApiKeys::default(),
Some(grok_tokens("grok-abc", Some(3600))),
);
assert!(mgr.has_any_key());
}
#[test]
fn manager_has_any_key_false_for_blank_grok_and_no_keys() {
let mgr = make_manager_with_grok(ApiKeys::default(), Some(grok_tokens(" ", None)));
assert!(!mgr.has_any_key());
}
// ── geap credentials ────────────────────────────────────────────
#[test]
@@ -587,13 +83,10 @@ fn geap_needs_refresh_lead_time_boundaries() {
#[test]
fn api_keys_for_request_includes_geap_token_when_gate_and_binding_match() {
let mgr = make_manager_with_geap(geap_loaded("geap-abc", Some(3600)));
let result = mgr
.api_keys_for_request(false, false, Some(geap_gate()))
.unwrap();
let result = mgr.api_keys_for_request(false, Some(geap_gate())).unwrap();
let credentials = result.google_cloud_credentials.unwrap();
assert_eq!(credentials.access_token, "geap-abc");
// The GEAP token is independent of the BYO key gate.
assert!(result.anthropic.is_empty());
assert!(result.aws_credentials.is_none());
}
#[test]
@@ -602,9 +95,7 @@ fn api_keys_for_request_includes_expired_geap_token() {
// rejects truly invalid ones, which surfaces a recoverable error instead
// of a silent fallback to another route.
let mgr = make_manager_with_geap(geap_loaded("geap-abc", Some(0)));
let result = mgr
.api_keys_for_request(false, false, Some(geap_gate()))
.unwrap();
let result = mgr.api_keys_for_request(false, Some(geap_gate())).unwrap();
assert_eq!(
result.google_cloud_credentials.unwrap().access_token,
"geap-abc"
@@ -616,7 +107,7 @@ fn api_keys_for_request_omits_geap_token_without_gate() {
// No gate (policy off at the call site) ⇒ no GEAP credentials, even when
// a token is loaded.
let mgr = make_manager_with_geap(geap_loaded("geap-abc", Some(3600)));
assert!(mgr.api_keys_for_request(false, false, None).is_none());
assert!(mgr.api_keys_for_request(false, None).is_none());
}
#[test]
@@ -626,19 +117,19 @@ fn api_keys_for_request_omits_geap_token_on_binding_mismatch() {
// A different user (sign-out/account switch).
let mut gate = geap_gate();
gate.user_uid = "someone-else".into();
assert!(mgr.api_keys_for_request(false, false, Some(gate)).is_none());
assert!(mgr.api_keys_for_request(false, Some(gate)).is_none());
// A different audience (admin changed the pool/provider).
let mut gate = geap_gate();
gate.audience = "//iam.googleapis.com/projects/2/locations/global/workloadIdentityPools/other/providers/other".into();
assert!(mgr.api_keys_for_request(false, false, Some(gate)).is_none());
assert!(mgr.api_keys_for_request(false, Some(gate)).is_none());
// A different service account (admin changed impersonation target).
let mut gate = geap_gate();
gate.federation = GeapFederation::ServiceAccount {
email: "other@proj.iam.gserviceaccount.com".into(),
};
assert!(mgr.api_keys_for_request(false, false, Some(gate)).is_none());
assert!(mgr.api_keys_for_request(false, Some(gate)).is_none());
}
#[test]
@@ -648,9 +139,7 @@ fn api_keys_for_request_serves_previous_geap_token_while_refreshing() {
let mgr = make_manager_with_geap(GeapCredentialsState::Refreshing {
previous: Some((geap_credentials("geap-old", Some(10)), geap_binding())),
});
let result = mgr
.api_keys_for_request(false, false, Some(geap_gate()))
.unwrap();
let result = mgr.api_keys_for_request(false, Some(geap_gate())).unwrap();
assert_eq!(
result.google_cloud_credentials.unwrap().access_token,
"geap-old"
@@ -661,9 +150,7 @@ fn api_keys_for_request_serves_previous_geap_token_while_refreshing() {
fn api_keys_for_request_omits_geap_token_during_first_mint() {
// The very first mint has nothing to serve yet.
let mgr = make_manager_with_geap(GeapCredentialsState::Refreshing { previous: None });
assert!(mgr
.api_keys_for_request(false, false, Some(geap_gate()))
.is_none());
assert!(mgr.api_keys_for_request(false, Some(geap_gate())).is_none());
}
#[test]
@@ -679,9 +166,7 @@ fn api_keys_for_request_omits_geap_token_for_non_loaded_states() {
},
] {
let mgr = make_manager_with_geap(state);
assert!(mgr
.api_keys_for_request(false, false, Some(geap_gate()))
.is_none());
assert!(mgr.api_keys_for_request(false, Some(geap_gate())).is_none());
}
}
@@ -692,5 +177,50 @@ fn api_keys_for_request_omits_geap_token_when_previous_binding_mismatches() {
});
let mut gate = geap_gate();
gate.user_uid = "someone-else".into();
assert!(mgr.api_keys_for_request(false, false, Some(gate)).is_none());
assert!(mgr.api_keys_for_request(false, Some(gate)).is_none());
}
// ── aws credentials ─────────────────────────────────────────────
#[test]
fn api_keys_for_request_none_when_nothing_configured() {
let mgr = make_manager();
assert!(mgr.api_keys_for_request(false, None).is_none());
}
#[test]
fn api_keys_for_request_includes_aws_credentials_when_requested() {
let mut mgr = make_manager();
mgr.aws_credentials_state = AwsCredentialsState::Loaded {
credentials: AwsCredentials::new("ak".into(), "sk".into(), None, None),
loaded_at: SystemTime::now(),
};
let result = mgr.api_keys_for_request(true, None).unwrap();
assert_eq!(result.aws_credentials.unwrap().access_key, "ak");
}
#[test]
fn api_keys_for_request_omits_aws_credentials_when_not_requested() {
let mut mgr = make_manager();
mgr.aws_credentials_state = AwsCredentialsState::Loaded {
credentials: AwsCredentials::new("ak".into(), "sk".into(), None, None),
loaded_at: SystemTime::now(),
};
assert!(mgr.api_keys_for_request(false, None).is_none());
}
#[test]
fn api_keys_for_request_includes_aws_credentials_when_oidc_managed_regardless_of_flag() {
let mut mgr = make_manager();
mgr.aws_credentials_state = AwsCredentialsState::Loaded {
credentials: AwsCredentials::new("ak".into(), "sk".into(), None, None),
loaded_at: SystemTime::now(),
};
mgr.aws_credentials_refresh_strategy = AwsCredentialsRefreshStrategy::OidcManaged {
task_id: Some("task-1".into()),
role_arn: "arn:aws:iam::123:role/test".into(),
region: "us-east-1".into(),
};
let result = mgr.api_keys_for_request(false, None).unwrap();
assert_eq!(result.aws_credentials.unwrap().access_key, "ak");
}
crates/ai/src/grok_subscription/mod.rs
-228
View File
@@ -1,228 +0,0 @@
//! Refresh orchestration for a connected xAI / Grok subscription's OAuth
//! tokens.
//!
//! The tokens themselves live in [`ApiKeyManager`] (the request-building
//! source of truth, persisted to secure storage under `GrokOAuthTokens`).
//! This module owns the network-facing refresh lifecycle — converting a
//! [`TokenResponse`] into stored [`GrokTokens`], proactively refreshing the
//! access token shortly before it expires, and rescheduling the next refresh.
//!
//! The Grok subscription is BYO auth, so background refresh follows the BYO
//! API key policy. That policy lives in the app layer (workspace settings),
//! which this crate has no visibility into; the app wires it in via
//! [`ApiKeyManager::set_grok_refresh_allowed`].
//!
//! The network/protocol side of the connect flow (authorize URL, loopback
//! callback server, token exchange/refresh) lives in the [`oauth`] submodule.
pub mod oauth;
use std::time::{Duration, SystemTime};
use galaxyui_core::r#async::Timer;
use galaxyui_core::ModelContext;
use self::oauth::TokenResponse;
use crate::api_keys::{ApiKeyManager, GrokTokens};
/// Refresh the access token this long before its hard expiry so a request
/// never races the expiration. Possibly-expired tokens are still sent (the
/// server is the authority on validity), so this lead time is purely about
/// keeping the token fresh, not about when it stops being sent.
const REFRESH_LEAD_TIME: Duration = Duration::from_secs(5 * 60);
/// Builds [`GrokTokens`] from a token-endpoint [`TokenResponse`], computing the
/// absolute `expires_at` from the relative `expires_in`. Values not present in
/// the response are carried over from `previous`: the refresh token when xAI
/// doesn't return a new one (refresh-token rotation is optional in OAuth 2.0),
/// and `connected_at` so it keeps reflecting the initial connection time
/// (initialized to now when there are no previous tokens, i.e. a fresh
/// connect).
pub fn grok_tokens_from_response(
response: TokenResponse,
previous: Option<&GrokTokens>,
) -> GrokTokens {
let expires_at = response
.expires_in
.and_then(|secs| u64::try_from(secs).ok())
.and_then(|secs| SystemTime::now().checked_add(Duration::from_secs(secs)));
GrokTokens {
access_token: response.access_token,
refresh_token: response
.refresh_token
.or_else(|| previous.and_then(|tokens| tokens.refresh_token.clone())),
expires_at,
connected_at: previous
.and_then(|tokens| tokens.connected_at)
.or_else(|| Some(SystemTime::now())),
}
}
impl ApiKeyManager {
/// Persists freshly obtained tokens (e.g. right after the connect flow) and
/// schedules the next proactive refresh.
pub fn store_grok_tokens(&mut self, response: TokenResponse, ctx: &mut ModelContext<Self>) {
apply_grok_tokens(self, response, ctx);
}
/// Updates whether background refresh of the stored Grok tokens is
/// allowed. The Grok subscription is BYO auth, so refresh follows the same
/// policy gate as request injection ([`Self::api_keys_for_request`]):
/// tokens that can never be sent shouldn't be kept fresh. The policy lives
/// in the app layer, which calls this at startup and whenever the policy
/// may have changed (e.g. team data arriving, or a workspace switch).
///
/// Schedules a refresh on a disabled -> enabled transition (refreshing
/// immediately if the token has already (nearly) expired); in-flight
/// timers re-check the flag when they fire. Repeated calls with an
/// unchanged value are no-ops, so duplicate timers can't pile up.
pub fn set_grok_refresh_allowed(&mut self, allowed: bool, ctx: &mut ModelContext<Self>) {
if self.grok_refresh_allowed == allowed {
return;
}
self.grok_refresh_allowed = allowed;
if allowed {
schedule_grok_token_refresh(self, ctx);
}
}
/// Request-time safety net: kicks off a background refresh of the stored
/// Grok tokens when they are nearing (or already past) expiry, so
/// upcoming requests can authenticate even if the proactive refresh loop
/// never armed or died (e.g. a stale BYO policy at startup, or an earlier
/// failed refresh). The triggering request still carries the currently
/// stored token — the server is the authority on its validity.
///
/// `byo_allowed` is the BYO API key policy as freshly evaluated by the
/// caller at request time. It also re-syncs the stored policy mirror,
/// which can go stale between `TeamsChanged` events; a disabled ->
/// enabled transition re-arms the proactive refresh loop.
pub fn refresh_grok_tokens_if_needed(
&mut self,
byo_allowed: bool,
ctx: &mut ModelContext<Self>,
) {
self.set_grok_refresh_allowed(byo_allowed, ctx);
if !byo_allowed || self.grok_refresh_in_flight {
return;
}
let Some(tokens) = self.grok_tokens() else {
return;
};
if !tokens.needs_refresh(REFRESH_LEAD_TIME) {
return;
}
let Some(refresh_token) = tokens.refresh_token.clone() else {
return;
};
log::info!(
"Grok OAuth token is nearing or past expiry at request time; refreshing in background"
);
spawn_grok_refresh(self, refresh_token, ctx);
}
}
/// Stores the tokens from `response` (carrying over the previous refresh token
/// and connection time when absent) and schedules the next proactive refresh.
fn apply_grok_tokens(
manager: &mut ApiKeyManager,
response: TokenResponse,
ctx: &mut ModelContext<ApiKeyManager>,
) {
let tokens = grok_tokens_from_response(response, manager.grok_tokens());
manager.set_grok_tokens(Some(tokens), ctx);
schedule_grok_token_refresh(manager, ctx);
}
/// Schedules a one-shot proactive refresh [`REFRESH_LEAD_TIME`] before the
/// current token's expiry (immediately if already within that window).
///
/// No-op when there's nothing to refresh against (no tokens, no refresh token,
/// or no known expiry). Reschedules itself after each successful refresh, so a
/// single call establishes an ongoing refresh loop for the lifetime of the
/// connection.
fn schedule_grok_token_refresh(manager: &mut ApiKeyManager, ctx: &mut ModelContext<ApiKeyManager>) {
// When the BYO API key policy is disabled the token is never sent, so
// don't refresh it in the background either. `set_grok_refresh_allowed`
// re-establishes the loop if the policy is later enabled.
if !manager.grok_refresh_allowed {
return;
}
let Some(tokens) = manager.grok_tokens() else {
return;
};
let Some(refresh_token) = tokens.refresh_token.clone() else {
return;
};
let Some(expires_at) = tokens.expires_at else {
// No expiry signal, so there's nothing to schedule against.
return;
};
let now = SystemTime::now();
let fire_at = expires_at.checked_sub(REFRESH_LEAD_TIME).unwrap_or(now);
let delay = fire_at.duration_since(now).unwrap_or(Duration::ZERO);
ctx.spawn(
async move {
Timer::after(delay).await;
},
move |manager, _output, ctx| {
// The BYO policy may have flipped off while we slept;
// `set_grok_refresh_allowed` restarts the loop if it flips back
// on.
if !manager.grok_refresh_allowed {
return;
}
// The stored token may have changed (reconnect/disconnect) while we
// slept; only refresh if our refresh token is still the current one.
let still_current = manager
.grok_tokens()
.and_then(|t| t.refresh_token.as_deref())
== Some(refresh_token.as_str());
if still_current {
spawn_grok_refresh(manager, refresh_token, ctx);
}
},
);
}
/// Kicks off a background token refresh using `refresh_token`, applying the
/// result (which reschedules the next refresh) or logging the failure.
///
/// No-op when a refresh is already in flight, so the proactive timer and the
/// request-time safety net can't issue overlapping refreshes.
fn spawn_grok_refresh(
manager: &mut ApiKeyManager,
refresh_token: String,
ctx: &mut ModelContext<ApiKeyManager>,
) {
if manager.grok_refresh_in_flight {
return;
}
manager.grok_refresh_in_flight = true;
ctx.spawn(
async move { oauth::refresh_access_token(&refresh_token).await },
|manager, result, ctx| {
manager.grok_refresh_in_flight = false;
match result {
Ok(response) => {
log::info!(
"Refreshed Grok OAuth token (expires_in={:?}, has_refresh_token={})",
response.expires_in,
response.refresh_token.is_some(),
);
apply_grok_tokens(manager, response, ctx);
}
Err(err) => {
// Leave the existing (possibly expired) token in place; the
// server remains the authority and will reject it if it's
// truly invalid. The request-time safety net
// (`ApiKeyManager::refresh_grok_tokens_if_needed`) retries
// on the next request.
log::error!("Failed to refresh Grok OAuth token: {err:#}");
}
}
},
);
}
crates/ai/src/grok_subscription/oauth.rs
-457
View File
@@ -1,457 +0,0 @@
//! OAuth flow for connecting an xAI / Grok subscription (e.g. SuperGrok) to
//! Warp, so users can "plug in" their subscription instead of pasting a
//! pay-as-you-go API key.
//!
//! This mirrors the public Grok-CLI desktop OAuth flow: an OAuth 2.0
//! Authorization Code grant with PKCE and a fixed loopback redirect URI. xAI's
//! auth server only accepts the loopback redirect for an allowlisted
//! `client_id` bound to a specific port, so we reuse the Grok-CLI client and
//! bind the callback server to that exact port.
//!
//! Some browsers/networks can't reach the loopback callback (e.g. Private
//! Network Access is blocked), in which case xAI's consent screen instead
//! *displays* the authorization code for the user to paste back into the app.
//! [`OauthAttempt::manual_code_exchange`] supports that fallback by capturing
//! the attempt's PKCE verifier so a pasted code can be exchanged directly,
//! without ever observing the loopback redirect.
//!
//! This module owns only the network/protocol side: building the authorize
//! URL, running the loopback callback server, and exchanging/refreshing tokens
//! at xAI's token endpoint. Persistence of the resulting tokens, proactive
//! refresh scheduling, and injection into the request live in the parent
//! [`crate::grok_subscription`] module (refresh orchestration) and
//! [`crate::api_keys::ApiKeyManager`] (storage + request injection).
use std::io::{ErrorKind, Read, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::time::Duration;
use anyhow::{bail, Context as _};
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine as _;
// `std::time::Instant` is disallowed (no wasm support); `instant::Instant` is a
// drop-in that re-exports the std type on native targets.
use instant::Instant;
use rand::RngCore as _;
use serde::Deserialize;
use sha2::{Digest, Sha256};
const CLIENT_ID: &str = "b1a00492-073a-47ea-816f-4c329264a828";
const AUTHORIZE_URL: &str = "https://auth.x.ai/oauth2/authorize";
const TOKEN_URL: &str = "https://auth.x.ai/oauth2/token";
const SCOPE: &str = "openid profile email offline_access grok-cli:access api:access";
const REDIRECT_HOST: &str = "127.0.0.1";
const REDIRECT_PORT: u16 = 56121;
/// How long we keep the loopback server open waiting for the user to approve
/// the consent screen in their browser.
const CALLBACK_TIMEOUT: Duration = Duration::from_secs(300);
/// How long to nap between non-blocking `accept()` attempts.
const POLL_INTERVAL: Duration = Duration::from_millis(100);
/// xAI's browser consent screen fetches the loopback callback from these
/// origins. Since that request crosses origins (https://accounts.x.ai ->
/// http://127.0.0.1), browsers require CORS and Private Network Access headers
/// before the page can observe the callback response.
const CORS_ALLOWED_ORIGINS: [&str; 2] = ["https://accounts.x.ai", "https://auth.x.ai"];
fn redirect_uri() -> String {
format!("http://{REDIRECT_HOST}:{REDIRECT_PORT}/callback")
}
/// One in-flight OAuth login attempt: the bound loopback callback listener
/// plus the per-attempt PKCE/CSRF secrets, which never leave this module.
///
/// Construct with [`OauthAttempt::start`], open [`OauthAttempt::authorize_url`]
/// in the browser, then await [`OauthAttempt::finish`] to obtain tokens. Tying
/// the secrets to the attempt guarantees the same PKCE verifier and CSRF state
/// are used for both the authorize URL and the code exchange.
pub struct OauthAttempt {
listener: TcpListener,
pkce: PkceParams,
}
impl OauthAttempt {
/// Binds the loopback callback server and generates fresh per-attempt
/// secrets. Call this before opening the browser so a bind failure (e.g.
/// another login already in progress, or Grok-CLI holding the port)
/// surfaces before a browser tab opens.
pub fn start() -> anyhow::Result<Self> {
Ok(Self {
listener: bind_callback_listener()?,
pkce: PkceParams::generate(),
})
}
/// The authorization URL the user's browser should open to begin the flow.
pub fn authorize_url(&self) -> String {
authorize_url(&self.pkce)
}
/// Runs the rest of the browser-based PKCE flow: waits for the loopback
/// callback, validates the CSRF state, and exchanges the authorization
/// code for tokens. Consumes the attempt so its secrets can't be reused.
pub async fn finish(self) -> anyhow::Result<TokenResponse> {
run_oauth_flow(self.listener, self.pkce).await
}
/// Clones the PKCE verifier for the pasted-code fallback while the
/// loopback flow continues racing in parallel.
pub fn manual_code_exchange(&self) -> ManualCodeExchange {
ManualCodeExchange {
verifier: self.pkce.verifier.clone(),
}
}
}
/// Completes OAuth from a manually-pasted authorization code.
///
/// There is no redirect `state` to validate in this out-of-band path; PKCE
/// protects the exchange.
#[derive(Clone)]
pub struct ManualCodeExchange {
verifier: String,
}
impl ManualCodeExchange {
/// Exchanges a user-pasted authorization `code` with the attempt's PKCE verifier.
pub async fn exchange(&self, code: &str) -> anyhow::Result<TokenResponse> {
let code = code.trim();
if code.is_empty() {
bail!("enter the code shown in your browser to finish connecting");
}
exchange_code_for_tokens(code, &self.verifier).await
}
}
/// The per-attempt secrets for one authorization request: the PKCE
/// verifier/challenge pair and the CSRF `state` value.
struct PkceParams {
verifier: String,
challenge: String,
/// CSRF token echoed back on the redirect and validated against the
/// response before the code is exchanged.
state: String,
}
impl PkceParams {
/// Generates a fresh PKCE verifier + S256 challenge and a random CSRF state.
fn generate() -> Self {
let verifier = random_url_safe_token();
let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()));
let state = random_url_safe_token();
Self {
verifier,
challenge,
state,
}
}
}
/// Returns a URL-safe, unpadded base64 string of 32 random bytes. This is used
/// for both the PKCE code verifier (RFC 7636 allows 43-128 chars from the
/// unreserved set) and the CSRF state.
fn random_url_safe_token() -> String {
let mut bytes = [0u8; 32];
rand::thread_rng().fill_bytes(&mut bytes);
URL_SAFE_NO_PAD.encode(bytes)
}
/// Builds the authorization URL the user's browser should open to begin the
/// flow.
fn authorize_url(pkce: &PkceParams) -> String {
let redirect = redirect_uri();
// `plan=generic` opts the consent screen into xAI's generic OAuth plan tier
// (required for loopback OAuth from non-allowlisted clients); `referrer`
// is best-effort attribution in xAI's OAuth logs.
let params: [(&str, &str); 9] = [
("response_type", "code"),
("client_id", CLIENT_ID),
("redirect_uri", &redirect),
("scope", SCOPE),
("code_challenge", &pkce.challenge),
("code_challenge_method", "S256"),
("state", &pkce.state),
("plan", "generic"),
("referrer", "warp"),
];
let query =
serde_urlencoded::to_string(params).expect("static OAuth params are always serializable");
format!("{AUTHORIZE_URL}?{query}")
}
/// The token endpoint's response. Fields beyond `access_token` are optional
/// because xAI does not always return them. Other response fields (e.g.
/// `token_type`, `scope`) are ignored since nothing consumes them.
#[derive(Debug, Deserialize)]
pub struct TokenResponse {
pub access_token: String,
#[serde(default)]
pub refresh_token: Option<String>,
#[serde(default)]
pub expires_in: Option<i64>,
}
/// The authorization code and state captured from the loopback redirect.
struct CallbackData {
code: String,
state: String,
}
/// Binds the loopback callback server to the fixed redirect address.
fn bind_callback_listener() -> anyhow::Result<TcpListener> {
let listener = TcpListener::bind((REDIRECT_HOST, REDIRECT_PORT)).with_context(|| {
format!(
"couldn't bind the Grok OAuth callback server to {REDIRECT_HOST}:{REDIRECT_PORT}. \
Another login may be in progress, or another app (e.g. Grok CLI) is using the port."
)
})?;
listener
.set_nonblocking(true)
.context("failed to set the Grok OAuth callback listener to non-blocking mode")?;
Ok(listener)
}
/// Runs the full browser-based PKCE flow: waits for the loopback callback on a
/// dedicated thread, validates the CSRF state, and exchanges the authorization
/// code for tokens.
async fn run_oauth_flow(listener: TcpListener, pkce: PkceParams) -> anyhow::Result<TokenResponse> {
// The loopback accept loop is blocking, so run it on a dedicated OS thread
// and bridge the result back through a runtime-agnostic async channel.
let (tx, rx) = async_channel::bounded(1);
std::thread::Builder::new()
.name("grok-oauth-callback".to_owned())
.spawn(move || {
// `send_blocking` is disallowed (no wasm support); block this
// dedicated thread on the async `send` instead.
let _ = galaxyui_core::r#async::block_on(
tx.send(wait_for_callback(&listener, CALLBACK_TIMEOUT)),
);
})
.context("failed to spawn the Grok OAuth callback server thread")?;
let callback = rx
.recv()
.await
.context("the Grok OAuth callback server stopped unexpectedly")??;
if callback.state != pkce.state {
bail!("the authorization response state did not match — aborting to prevent CSRF");
}
exchange_code_for_tokens(&callback.code, &pkce.verifier).await
}
/// Blocks (on a non-blocking listener with polling) until the browser hits the
/// redirect URI, returning the captured code and state, or an error on timeout.
fn wait_for_callback(listener: &TcpListener, timeout: Duration) -> anyhow::Result<CallbackData> {
let deadline = Instant::now() + timeout;
loop {
if Instant::now() >= deadline {
bail!("timed out waiting for the Grok authorization callback");
}
match listener.accept() {
Ok((stream, _)) => match handle_callback_connection(stream)? {
Some(data) => return Ok(data),
// Unrelated request (e.g. /favicon.ico); keep waiting.
None => continue,
},
Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
std::thread::sleep(POLL_INTERVAL);
}
Err(e) => {
return Err(anyhow::Error::new(e).context("Grok OAuth callback accept failed"))
}
}
}
}
/// Reads a single HTTP request from the callback connection, writes back a
/// minimal HTML response, and extracts the OAuth parameters.
///
/// Returns `Ok(None)` for requests that aren't the OAuth callback (so the
/// caller keeps listening), `Ok(Some(..))` on a successful callback, and `Err`
/// when the provider reported an error or the callback was malformed.
fn handle_callback_connection(mut stream: TcpStream) -> anyhow::Result<Option<CallbackData>> {
// The accepted stream may inherit the listener's non-blocking flag on some
// platforms; force blocking reads with a timeout so we get the full request
// line without spinning.
stream.set_nonblocking(false).ok();
stream.set_read_timeout(Some(Duration::from_secs(10))).ok();
let mut buf = [0u8; 8192];
let n = stream
.read(&mut buf)
.context("failed to read the Grok OAuth callback request")?;
let request = String::from_utf8_lossy(&buf[..n]);
let origin = request_header(&request, "Origin");
// The request line looks like: "GET /callback?code=...&state=... HTTP/1.1".
let mut request_line_parts = request
.lines()
.next()
.unwrap_or_default()
.split_whitespace();
let method = request_line_parts.next().unwrap_or_default();
let path = request_line_parts.next().unwrap_or_default();
if method == "OPTIONS" && path.starts_with("/callback") {
write_response(&mut stream, "204 No Content", "", origin.as_deref());
return Ok(None);
}
let Some(query) = path
.strip_prefix("/callback")
.and_then(|rest| rest.strip_prefix('?'))
else {
write_response(
&mut stream,
"404 Not Found",
"Not found.",
origin.as_deref(),
);
return Ok(None);
};
let mut code = None;
let mut state = None;
let mut error = None;
let mut error_description = None;
let pairs: Vec<(String, String)> = serde_urlencoded::from_str(query).unwrap_or_default();
for (key, value) in pairs {
match key.as_str() {
"code" => code = Some(value),
"state" => state = Some(value),
"error" => error = Some(value),
"error_description" => error_description = Some(value),
_ => {}
}
}
if let Some(error) = error {
write_response(
&mut stream,
"400 Bad Request",
FAILURE_HTML,
origin.as_deref(),
);
let detail = error_description.unwrap_or(error);
bail!("Grok authorization was denied or failed: {detail}");
}
let (Some(code), Some(state)) = (code, state) else {
write_response(
&mut stream,
"400 Bad Request",
FAILURE_HTML,
origin.as_deref(),
);
bail!("the Grok authorization callback was missing the code or state parameter");
};
write_response(&mut stream, "200 OK", SUCCESS_HTML, origin.as_deref());
Ok(Some(CallbackData { code, state }))
}
fn request_header(request: &str, header_name: &str) -> Option<String> {
request.lines().skip(1).find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case(header_name)
.then(|| value.trim().to_owned())
})
}
/// Writes a minimal HTTP/1.1 response and closes the connection.
fn write_response(stream: &mut TcpStream, status: &str, body: &str, origin: Option<&str>) {
let cors_headers = cors_headers(origin);
let response = format!(
"HTTP/1.1 {status}\r\nContent-Type: text/html; charset=utf-8\r\n\
{cors_headers}Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
body.len()
);
let _ = stream.write_all(response.as_bytes());
let _ = stream.flush();
let _ = stream.shutdown(Shutdown::Both);
}
fn cors_headers(origin: Option<&str>) -> String {
origin
.filter(|origin| CORS_ALLOWED_ORIGINS.contains(origin))
.map(|origin| {
format!(
"Access-Control-Allow-Origin: {origin}\r\n\
Access-Control-Allow-Methods: GET, OPTIONS\r\n\
Access-Control-Allow-Headers: Content-Type\r\n\
Access-Control-Allow-Private-Network: true\r\n\
Vary: Origin\r\n"
)
})
.unwrap_or_default()
}
/// Exchanges the authorization code for OAuth tokens at xAI's token endpoint.
async fn exchange_code_for_tokens(code: &str, verifier: &str) -> anyhow::Result<TokenResponse> {
let redirect = redirect_uri();
let form: [(&str, &str); 5] = [
("grant_type", "authorization_code"),
("code", code),
("redirect_uri", &redirect),
("client_id", CLIENT_ID),
("code_verifier", verifier),
];
post_token_request(&form).await
}
/// Exchanges a previously obtained refresh token for a fresh set of tokens via
/// the OAuth 2.0 `refresh_token` grant. Used to keep the connected Grok
/// subscription's access token valid without re-running the browser flow.
///
/// xAI may or may not return a new `refresh_token`; callers should fall back to
/// the existing one when [`TokenResponse::refresh_token`] is `None` (rotation is
/// optional in OAuth 2.0).
pub async fn refresh_access_token(refresh_token: &str) -> anyhow::Result<TokenResponse> {
let form: [(&str, &str); 3] = [
("grant_type", "refresh_token"),
("refresh_token", refresh_token),
("client_id", CLIENT_ID),
];
post_token_request(&form).await
}
/// POSTs a form-encoded body to xAI's token endpoint and parses the
/// [`TokenResponse`]. Shared by the initial code exchange and refresh grants.
async fn post_token_request<T: serde::Serialize + ?Sized>(
form: &T,
) -> anyhow::Result<TokenResponse> {
let response = http_client::Client::new()
.post(TOKEN_URL)
.form(form)
.send()
.await
.context("failed to send the Grok token request")?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
bail!("Grok token request failed ({status}): {body}");
}
response
.json::<TokenResponse>()
.await
.context("failed to parse the Grok token response")
}
const SUCCESS_HTML: &str = "<!doctype html><html><head><meta charset=\"utf-8\">\
<title>Warp — Grok connected</title></head>\
<body style=\"font-family:system-ui,-apple-system,sans-serif;text-align:center;padding:3rem\">\
<h1>Grok connected</h1><p>You can close this window and return to Warp.</p></body></html>";
const FAILURE_HTML: &str = "<!doctype html><html><head><meta charset=\"utf-8\">\
<title>Warp — Grok authorization failed</title></head>\
<body style=\"font-family:system-ui,-apple-system,sans-serif;text-align:center;padding:3rem\">\
<h1>Authorization failed</h1><p>Something went wrong. Return to Warp and try again.</p></body></html>";
#[cfg(test)]
#[path = "oauth_tests.rs"]
mod tests;
crates/ai/src/grok_subscription/oauth_tests.rs
-57
View File
@@ -1,57 +0,0 @@
use super::*;
#[test]
fn authorize_url_contains_required_params() {
let pkce = PkceParams::generate();
let url = authorize_url(&pkce);
assert!(url.starts_with("https://auth.x.ai/oauth2/authorize?"));
assert!(url.contains("response_type=code"));
assert!(url.contains(&format!("client_id={CLIENT_ID}")));
assert!(url.contains("code_challenge_method=S256"));
assert!(url.contains("scope=openid"));
assert!(url.contains("plan=generic"));
assert!(url.contains("referrer=warp"));
// The redirect URI must be percent-encoded and match the registered value.
assert!(url.contains("redirect_uri=http%3A%2F%2F127.0.0.1%3A56121%2Fcallback"));
// The CSRF state and PKCE challenge are echoed into the URL verbatim
// (both are URL-safe base64, so no percent-encoding is applied).
assert!(url.contains(&format!("state={}", pkce.state)));
assert!(url.contains(&format!("code_challenge={}", pkce.challenge)));
}
#[test]
fn token_response_parses_minimal_and_full() {
let minimal: TokenResponse =
serde_json::from_str(r#"{"access_token":"abc"}"#).expect("minimal response should parse");
assert_eq!(minimal.access_token, "abc");
assert!(minimal.refresh_token.is_none());
assert!(minimal.expires_in.is_none());
// Unconsumed response fields (token_type, scope) are ignored by serde.
let full: TokenResponse = serde_json::from_str(
r#"{"access_token":"a","refresh_token":"r","token_type":"Bearer","expires_in":3600,"scope":"api:access"}"#,
)
.expect("full response should parse");
assert_eq!(full.access_token, "a");
assert_eq!(full.refresh_token.as_deref(), Some("r"));
assert_eq!(full.expires_in, Some(3600));
}
#[test]
fn manual_code_exchange_captures_attempt_verifier() {
let pkce = PkceParams::generate();
let exchange = ManualCodeExchange {
verifier: pkce.verifier.clone(),
};
assert_eq!(exchange.verifier, pkce.verifier);
}
#[test]
fn manual_code_exchange_rejects_blank_code() {
let exchange = ManualCodeExchange {
verifier: "verifier".to_string(),
};
let result = galaxyui_core::r#async::block_on(exchange.exchange(" "));
assert!(result.is_err());
}
crates/ai/src/lib.rs
-2
View File
@@ -2,8 +2,6 @@ pub mod agent;
pub mod api_keys;
pub mod aws_credentials;
pub mod geap_credentials;
#[cfg(not(target_family = "wasm"))]
pub mod grok_subscription;
pub mod llm_id;
pub use llm_id::LLMId;