i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-04 00:17:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(providers): add xAI provider with device code OAuth

Browse Source 
Add a separate `xai` provider type for xAI Grok CLI subscription accounts.
It is independent of the existing `grok` provider, which reverse-proxies
grok.com with browser cookies; behavior of `grok` is unchanged.

Account binding uses the xAI device code flow, so no local callback
listener is needed and headless deployments can bind accounts. Refresh
tokens can also be imported individually or in batches, and are rotated
on refresh.

OAuth requests default to the cli-chat-proxy Responses API; API keys and
compact stay on api.x.ai. Explicit custom gateways are preserved. Only
`openai:responses` and `openai:responses:compact` are exposed; Chat,
Claude and Gemini clients reach the provider through Aether's existing
cross-format conversion rather than new native endpoints.

Upstream Responses payloads are sanitized for what xAI actually rejects:
`previous_response_id` and `metadata.user_id` are dropped, hosted
`tool_choice` is rewritten, `web_search` is restored for converted
clients, `image_generation` is stripped on older Grok conversation
models, unsupported reasoning effort is removed, and requested
`reasoning.encrypted_content` is preserved with a replay policy keyed on
the configured provider type rather than the model name.

Quota refresh reads /user and /billing?format=credits and stores a
structured usage snapshot; a prepaid balance keeps an account selectable
after the weekly allowance is exhausted. API-key accounts skip the
subscription billing surface. The admin UI shows remaining weekly quota
as a labeled bar in the provider drawer and the pool list.

Co-Authored-By: Claude Opus 5 <[email protected]>
This commit is contained in:
stabey
2026-09-14 21:09:03 +08:00
co-authored by Claude Opus 5
parent 60b89cc840
commit e83399db2f
71 changed files with 5495 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs
+5
View File
@@ -583,6 +583,11 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
source_model,
codex_model_capabilities.as_ref(),
);
crate::ai_serving::transport::xai::insert_cli_identity_headers_if_needed(
transport.as_ref(),
prepared.provider_api_format.as_str(),
&mut provider_request_headers,
);
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
redaction.redacted,
apps/aether-gateway/src/ai_serving/planner/standard/deepseek.rs
+24 -1
View File
@@ -13,7 +13,9 @@ pub(crate) fn openai_responses_reasoning_replay_policy(
base_url: &str,
_provider_model: &str,
) -> crate::ai_serving::OpenAiResponsesReasoningReplayPolicy {
if is_deepseek_provider(provider_type, base_url) {
if provider_type.trim().eq_ignore_ascii_case("xai") {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::XaiEncrypted
} else if is_deepseek_provider(provider_type, base_url) {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
} else {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
@@ -238,6 +240,27 @@ mod tests {
openai_responses_reasoning_replay_policy,
};
#[test]
fn xai_reasoning_policy_comes_from_provider_type() {
use crate::ai_serving::OpenAiResponsesReasoningReplayPolicy;
assert_eq!(
openai_responses_reasoning_replay_policy(
"xai",
"https://custom.example/v1",
"grok-4.6"
),
OpenAiResponsesReasoningReplayPolicy::XaiEncrypted
);
assert_eq!(
openai_responses_reasoning_replay_policy(
"openai",
"https://custom.example/v1",
"grok-4.6"
),
OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
);
}
#[test]
fn detects_deepseek_provider_only_by_official_host() {
assert!(!is_deepseek_provider(
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs
+1
View File
@@ -488,6 +488,7 @@ impl ResponsesWebSocketBodyNormalization {
digest.update([match self.reasoning_replay_policy {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds => 0,
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque => 1,
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::XaiEncrypted => 2,
}]);
update_normalization_optional_json_digest(&mut digest, self.model_directive_patch.as_ref());
digest.finalize().into()
apps/aether-gateway/src/ai_serving/pure/mod.rs
+2 -1
View File
@@ -12,7 +12,8 @@ pub(crate) use aether_ai_formats::api::{
apply_codex_openai_responses_websocket_continuation_body_edits_with_source_model_and_capabilities,
apply_codex_openai_special_headers, apply_model_directive_mapping_patch,
apply_model_directive_overrides_from_model, apply_model_directive_overrides_from_request,
apply_openai_responses_compact_special_body_edits, build_chatgpt_web_image_request_body,
apply_openai_responses_compact_special_body_edits, apply_xai_upstream_payload_edits,
apply_xai_upstream_payload_edits_with_client, build_chatgpt_web_image_request_body,
build_codex_model_catalog_metadata, build_codex_openai_image_api_provider_request_body,
build_core_error_body_for_client_format, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_request_body_with_model_directives,
apps/aether-gateway/src/ai_serving/transport.rs
+4
View File
@@ -58,6 +58,10 @@ pub(crate) mod windsurf {
pub(crate) use aether_provider_transport::windsurf::*;
}
pub(crate) mod xai {
pub(crate) use aether_provider_transport::xai::*;
}
pub(crate) use aether_provider_transport::{
append_transport_diagnostics_to_value, apply_codex_fingerprint_convergence,
apply_codex_fingerprint_convergence_with_context, apply_local_auth_config_header_overrides,
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/execution.rs
+2 -1
View File
@@ -74,7 +74,8 @@ fn validate_batch_access_token_import(
) -> Result<(), String> {
if !provider_type_supports_access_token_import(provider_type) {
return Err(
"Access Token 导入仅支持 Claude Code / Codex / ChatGPT Web / Grok Provider".to_string(),
"Access Token 导入仅支持 Claude Code / Codex / ChatGPT Web / Grok / xAI Provider"
.to_string(),
);
}
if provider_type.eq_ignore_ascii_case("claude_code") {
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs
+42 -9
View File
@@ -214,7 +214,12 @@ fn extract_admin_provider_oauth_batch_import_entry(
} else {
let sso_from_cookie = grok_cookie_session_token(provider_type, raw_token);
let token_input = sso_from_cookie.as_deref().unwrap_or(raw_token);
let (refresh_token, access_token) = import_tokens_from_raw_token(token_input);
let (refresh_token, access_token) =
if provider_type.trim().eq_ignore_ascii_case("xai") {
(None, Some(token_input.to_string()))
} else {
import_tokens_from_raw_token(token_input)
};
let (refresh_token, access_token) = normalize_provider_import_tokens(
provider_type,
refresh_token.as_deref(),
@@ -262,6 +267,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
let object = normalized_claude_object.as_ref().unwrap_or(object);
let is_grok = provider_type.trim().eq_ignore_ascii_case("grok");
let is_windsurf = provider_type.trim().eq_ignore_ascii_case("windsurf");
let is_xai = provider_type.trim().eq_ignore_ascii_case("xai");
let is_codex_agent_identity = provider_type.trim().eq_ignore_ascii_case("codex")
&& aether_provider_transport::is_codex_agent_identity_auth_config_value(item);
if is_codex_agent_identity {
@@ -336,14 +342,6 @@ fn extract_admin_provider_oauth_batch_import_entry(
} else {
None
};
let (refresh_token, access_token) = normalize_provider_import_tokens(
provider_type,
refresh_token.as_deref(),
access_token
.as_deref()
.or(session_token.as_deref())
.or(header_bearer_token.as_deref()),
);
let windsurf_api_key = is_windsurf
.then(|| {
coerce_admin_provider_oauth_import_str(
@@ -351,6 +349,22 @@ fn extract_admin_provider_oauth_batch_import_entry(
)
})
.flatten();
let xai_api_key = is_xai
.then(|| {
coerce_admin_provider_oauth_import_str(
object.get("api_key").or_else(|| object.get("apiKey")),
)
})
.flatten();
let (refresh_token, access_token) = normalize_provider_import_tokens(
provider_type,
refresh_token.as_deref(),
access_token
.as_deref()
.or(session_token.as_deref())
.or(header_bearer_token.as_deref())
.or(xai_api_key.as_deref()),
);
let windsurf_token = is_windsurf
.then(|| {
coerce_admin_provider_oauth_import_str(
@@ -1577,4 +1591,23 @@ mod tests {
assert!(entries[1].access_token.is_none());
assert!(entries[1].raw_credentials.is_none());
}
#[test]
fn parses_xai_api_key_json_and_raw_lines_as_access_token() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"xai",
r#"{"api_key":"xai-api-key","email":"[email protected]"}
{"refresh_token":"xai-refresh"}
xai-raw-api-key"#,
);
assert_eq!(entries.len(), 3);
assert!(entries[0].refresh_token.is_none());
assert_eq!(entries[0].access_token.as_deref(), Some("xai-api-key"));
assert_eq!(entries[0].email.as_deref(), Some("[email protected]"));
assert_eq!(entries[1].refresh_token.as_deref(), Some("xai-refresh"));
assert!(entries[1].access_token.is_none());
assert!(entries[2].refresh_token.is_none());
assert_eq!(entries[2].access_token.as_deref(), Some("xai-raw-api-key"));
}
}
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/authorize.rs
+15 -2
View File
@@ -186,10 +186,10 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
));
};
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
if provider_type != "kiro" && provider_type != "windsurf" {
if provider_type != "kiro" && provider_type != "windsurf" && provider_type != "xai" {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"设备授权仅支持 Kiro / Windsurf provider",
"设备授权仅支持 Kiro / Windsurf / xAI provider",
));
}
let Some(principal) = request_context
@@ -219,6 +219,19 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
)
.await;
if provider_type == "xai" {
return super::xai::handle_admin_provider_oauth_xai_device_authorize(
state,
&provider_id,
&provider,
principal,
runtime_endpoint.as_ref(),
request_proxy,
payload.proxy_node_id.as_deref(),
)
.await;
}
if provider_type == "windsurf" {
let session_id = generate_provider_oauth_nonce();
let login_option = payload
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/mod.rs
+1
View File
@@ -2,6 +2,7 @@ mod authorize;
mod lease;
mod poll;
mod session;
mod xai;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/poll.rs
+12
View File
@@ -479,6 +479,18 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
)
.await;
if provider_type == "xai" {
return super::xai::handle_admin_provider_oauth_xai_device_poll(
state,
&provider,
&endpoints,
request_proxy,
session_id,
session,
)
.await;
}
if provider_type == "windsurf" {
return handle_admin_provider_oauth_windsurf_browser_device_poll(
state,
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/xai.rs
+368
View File
@@ -0,0 +1,368 @@
use super::session::attach_admin_provider_oauth_device_poll_terminal_response;
use crate::control::GatewayAdminPrincipalContext;
use crate::handlers::admin::provider::oauth::dispatch::helpers::admin_provider_oauth_key_name_from_auth_config;
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::provisioning::{
provider_oauth_active_api_formats, provider_oauth_key_proxy_value,
};
use crate::handlers::admin::provider::oauth::runtime::spawn_provider_oauth_account_state_refresh_after_update;
use crate::handlers::admin::provider::oauth::state::{
current_unix_secs, generate_provider_oauth_nonce,
};
use crate::handlers::admin::request::AdminAppState;
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
use aether_data::repository::provider_oauth::{
StoredAdminProviderOAuthDeviceSession, KIRO_DEVICE_AUTH_SESSION_TTL_BUFFER_SECS,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
use aether_oauth::core::OAuthError;
use aether_oauth::provider::providers::{
XaiDevicePollOutcome, XaiProviderOAuthAdapter, XAI_CLIENT_ID, XAI_DEVICE_CODE_URL,
XAI_TOKEN_URL,
};
use aether_oauth::provider::ProviderOAuthTransportContext;
use axum::{
body::Body,
http,
response::{IntoResponse, Response},
Json,
};
use serde_json::{json, Value};
pub(super) async fn handle_admin_provider_oauth_xai_device_authorize(
state: &AdminAppState<'_>,
provider_id: &str,
provider: &StoredProviderCatalogProvider,
principal: &GatewayAdminPrincipalContext,
runtime_endpoint: Option<&StoredProviderCatalogEndpoint>,
request_proxy: Option<ProxySnapshot>,
proxy_node_id: Option<&str>,
) -> Result<Response<Body>, GatewayError> {
let device_url = state.provider_oauth_token_url("xai_device", XAI_DEVICE_CODE_URL);
let token_url = state.provider_oauth_token_url("xai", XAI_TOKEN_URL);
let adapter =
XaiProviderOAuthAdapter::default().with_endpoint_overrides(&device_url, &token_url);
let ctx = ProviderOAuthTransportContext {
provider_id: provider_id.to_string(),
provider_type: provider.provider_type.clone(),
endpoint_id: runtime_endpoint.map(|endpoint| endpoint.id.clone()),
key_id: None,
auth_type: Some("oauth".to_string()),
decrypted_api_key: None,
decrypted_auth_config: None,
provider_config: provider.config.clone(),
endpoint_config: runtime_endpoint.and_then(|endpoint| endpoint.config.clone()),
key_config: None,
network: aether_oauth::network::OAuthNetworkContext::provider_operation(
request_proxy.clone(),
),
};
let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state);
let authorization = match adapter.start_device_flow(&executor, &ctx).await {
Ok(authorization) => authorization,
Err(error) => {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
sanitize_xai_oauth_error(&error),
));
}
};
let now_unix_secs = current_unix_secs();
let session_id = generate_provider_oauth_nonce();
let session = StoredAdminProviderOAuthDeviceSession {
session_id: session_id.clone(),
provider_id: provider_id.to_string(),
initiated_by_user_id: principal.user_id.clone(),
initiated_by_session_id: principal.session_id.clone(),
initiated_by_management_token_id: principal.management_token_id.clone(),
region: String::new(),
client_id: XAI_CLIENT_ID.to_string(),
client_secret: String::new(),
device_code: authorization.device_code.clone(),
auth_type: Some("device".to_string()),
social_provider: None,
code_verifier: None,
redirect_uri: Some(token_url),
machine_id: None,
interval: authorization.interval,
expires_at_unix_secs: now_unix_secs.saturating_add(authorization.expires_in),
status: "pending".to_string(),
proxy_node_id: proxy_node_id
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
created_at_unix_ms: now_unix_secs,
key_id: None,
email: None,
replaced: false,
error_msg: None,
};
if let Err(response) = state
.save_provider_oauth_device_session(
&session_id,
&session,
authorization
.expires_in
.saturating_add(KIRO_DEVICE_AUTH_SESSION_TTL_BUFFER_SECS),
)
.await
{
return Ok(response);
}
Ok(Json(json!({
"session_id": session_id,
"user_code": authorization.user_code,
"verification_uri": authorization.verification_uri,
"verification_uri_complete": authorization.verification_uri_complete,
"expires_in": authorization.expires_in,
"interval": authorization.interval,
"auth_type": "device",
}))
.into_response())
}
pub(super) async fn handle_admin_provider_oauth_xai_device_poll(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoints: &[StoredProviderCatalogEndpoint],
request_proxy: Option<ProxySnapshot>,
session_id: &str,
mut session: StoredAdminProviderOAuthDeviceSession,
) -> Result<Response<Body>, GatewayError> {
let token_url = session
.redirect_uri
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.unwrap_or_else(|| state.provider_oauth_token_url("xai", XAI_TOKEN_URL));
let adapter =
XaiProviderOAuthAdapter::default().with_endpoint_overrides(XAI_DEVICE_CODE_URL, token_url);
let ctx = ProviderOAuthTransportContext {
provider_id: provider.id.clone(),
provider_type: provider.provider_type.clone(),
endpoint_id: None,
key_id: None,
auth_type: Some("oauth".to_string()),
decrypted_api_key: None,
decrypted_auth_config: None,
provider_config: provider.config.clone(),
endpoint_config: None,
key_config: None,
network: aether_oauth::network::OAuthNetworkContext::provider_operation(
request_proxy.clone(),
),
};
let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state);
let outcome = match adapter
.poll_device_token(&executor, &ctx, &session.device_code)
.await
{
Ok(outcome) => outcome,
Err(error) => {
return Ok(xai_device_poll_terminal_from_error(
state,
session_id,
&mut session,
&error,
)
.await);
}
};
match outcome {
XaiDevicePollOutcome::Pending => {
Ok(Json(json!({"status": "pending", "replaced": false})).into_response())
}
XaiDevicePollOutcome::SlowDown => {
Ok(Json(json!({"status": "slow_down", "replaced": false})).into_response())
}
XaiDevicePollOutcome::Authorized(result) => {
persist_xai_device_authorization(
state,
provider,
endpoints,
request_proxy,
session_id,
session,
*result,
)
.await
}
}
}
async fn persist_xai_device_authorization(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoints: &[StoredProviderCatalogEndpoint],
request_proxy: Option<ProxySnapshot>,
session_id: &str,
mut session: StoredAdminProviderOAuthDeviceSession,
result: aether_oauth::provider::ProviderOAuthTokenSet,
) -> Result<Response<Body>, GatewayError> {
let access_token = result.token_set.access_token.trim().to_string();
if access_token.is_empty() {
return Ok(Json(json!({
"status": "error",
"error": "xAI token 响应缺少 access_token",
"replaced": false,
}))
.into_response());
}
let mut auth_config = result.auth_config.as_object().cloned().unwrap_or_default();
auth_config.insert("provider_type".to_string(), json!("xai"));
auth_config.insert("auth_method".to_string(), json!("oauth"));
auth_config.insert("using_api".to_string(), json!(false));
let duplicate = match state
.find_duplicate_provider_oauth_key(&provider.id, &auth_config, None)
.await
{
Ok(duplicate) => duplicate,
Err(detail) => {
return Ok(Json(json!({
"status": "error",
"error": detail,
"replaced": false,
}))
.into_response());
}
};
let api_formats = provider_oauth_active_api_formats(endpoints);
let key_proxy = provider_oauth_key_proxy_value(session.proxy_node_id.as_deref());
let expires_at = result.token_set.expires_at_unix_secs;
let email = auth_config
.get("email")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let mut replaced = false;
let persisted_key = if let Some(existing_key) = duplicate {
replaced = true;
match state
.update_existing_provider_oauth_catalog_key(
&existing_key,
&provider.provider_type,
&access_token,
&auth_config,
&api_formats,
key_proxy.clone(),
expires_at,
)
.await?
{
Some(key) => key,
None => {
return Ok(build_internal_control_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"provider oauth write unavailable",
));
}
}
} else {
let key_name = admin_provider_oauth_key_name_from_auth_config(
&provider.provider_type,
&auth_config,
None,
);
match state
.create_provider_oauth_catalog_key(
&provider.id,
&provider.provider_type,
&key_name,
&access_token,
&auth_config,
&api_formats,
key_proxy,
expires_at,
)
.await?
{
Some(key) => key,
None => {
return Ok(build_internal_control_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"provider oauth write unavailable",
));
}
}
};
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
session.status = "authorized".to_string();
session.key_id = Some(persisted_key.id.clone());
session.email = email.clone();
session.replaced = replaced;
session.error_msg = None;
let _ = state
.save_provider_oauth_device_session(session_id, &session, 60)
.await;
Ok(attach_admin_provider_oauth_device_poll_terminal_response(
session_id,
"authorized",
Json(json!({
"status": "authorized",
"key_id": persisted_key.id,
"email": email,
"replaced": replaced,
}))
.into_response(),
))
}
async fn xai_device_poll_terminal_from_error(
state: &AdminAppState<'_>,
session_id: &str,
session: &mut StoredAdminProviderOAuthDeviceSession,
error: &OAuthError,
) -> Response<Body> {
let (status, message) = match error {
OAuthError::InvalidRequest(detail) if detail.to_ascii_lowercase().contains("expired") => {
("expired", "设备码已过期".to_string())
}
OAuthError::InvalidRequest(detail) if detail.to_ascii_lowercase().contains("denied") => {
("error", "用户拒绝授权".to_string())
}
_ => ("error", sanitize_xai_oauth_error(error)),
};
session.status = status.to_string();
session.error_msg = Some(message.clone());
let _ = state
.save_provider_oauth_device_session(session_id, session, 30)
.await;
attach_admin_provider_oauth_device_poll_terminal_response(
session_id,
status,
Json(json!({
"status": status,
"error": message,
"replaced": false,
}))
.into_response(),
)
}
fn sanitize_xai_oauth_error(error: &OAuthError) -> String {
match error {
OAuthError::InvalidRequest(_) => "xAI 设备授权失败: 请求参数无效".to_string(),
OAuthError::HttpStatus { status_code, .. } => {
format!("xAI 设备授权失败: HTTP {status_code}")
}
_ => "xAI 设备授权失败".to_string(),
}
}
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs
+10 -3
View File
@@ -715,7 +715,7 @@ async fn resolve_admin_provider_oauth_single_import_tokens(
if !provider_type_supports_access_token_import(provider_type) {
return Err(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"Access Token 导入仅支持 Claude Code / Codex / ChatGPT Web / Grok Provider",
"Access Token 导入仅支持 Claude Code / Codex / ChatGPT Web / Grok / xAI Provider",
));
}
@@ -867,7 +867,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
flatten_claude_code_credentials_payload(&mut raw_payload);
}
let refresh_token_input = import_payload_string(&raw_payload, "refresh_token", "refreshToken");
let access_token_input = import_payload_string_any(
let mut access_token_input = import_payload_string_any(
&raw_payload,
&[
"access_token",
@@ -879,6 +879,9 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
],
)
.or_else(|| provider_oauth_import_authorization_bearer_token_from_object(&raw_payload));
if provider_type == "xai" && access_token_input.is_none() {
access_token_input = import_payload_string(&raw_payload, "api_key", "apiKey");
}
let imported_expires_at =
import_payload_u64_any(&raw_payload, &["expires_at", "expiresAt", "expired"]);
let (refresh_token_input, access_token_input) = normalize_provider_import_tokens(
@@ -901,7 +904,11 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
if !create_agent_identity && refresh_token_input.is_none() && access_token_input.is_none() {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"Refresh Token、Access Token 或 sso_token 不能为空",
if provider_type == "xai" {
"Refresh Token、Access Token 或 api_key 不能为空"
} else {
"Refresh Token、Access Token 或 sso_token 不能为空"
},
));
}
if !is_fixed_provider_type_for_provider_oauth(&provider_type) {
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/start.rs
+12
View File
@@ -70,6 +70,12 @@ pub(super) async fn handle_admin_provider_oauth_start_key(
"Windsurf 请使用浏览器登录或导入凭据。",
));
}
if provider_type == "xai" {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"xAI 请使用设备授权或导入凭据。",
));
}
let Some(template) = admin_provider_oauth_template(&provider_type) else {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
@@ -167,6 +173,12 @@ pub(super) async fn handle_admin_provider_oauth_start_provider(
"Windsurf 请使用浏览器登录或导入凭据。",
));
}
if provider_type == "xai" {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"xAI 请使用设备授权或导入凭据。",
));
}
let Some(template) = admin_provider_oauth_template(&provider_type) else {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/token_import.rs
+48 -1
View File
@@ -121,6 +121,9 @@ pub(super) fn normalize_provider_import_tokens(
if provider_type == "grok" {
return (None, access_token.or(refresh_token));
}
if provider_type == "xai" {
return (refresh_token, access_token);
}
if provider_type == "claude_code" {
if access_token.is_none() && refresh_token.as_deref().is_some_and(is_claude_access_token) {
return (None, refresh_token);
@@ -237,7 +240,7 @@ pub(super) fn provider_oauth_import_authorization_bearer_token_from_object(
pub(super) fn provider_type_supports_access_token_import(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"claude_code" | "codex" | "chatgpt_web" | "grok"
"claude_code" | "codex" | "chatgpt_web" | "grok" | "xai"
)
}
@@ -331,6 +334,15 @@ pub(super) fn build_provider_access_token_import_auth_config(
auth_config.insert("sso_token".to_string(), json!(access_token));
auth_config.insert("auth_method".to_string(), json!("sso_token"));
}
if provider_type.trim().eq_ignore_ascii_case("xai") {
if refresh_token.is_some() {
auth_config.insert("auth_method".to_string(), json!("oauth"));
auth_config.insert("using_api".to_string(), json!(false));
} else {
auth_config.insert("auth_method".to_string(), json!("api_key"));
auth_config.insert("using_api".to_string(), json!(true));
}
}
auth_config.insert(
"access_token_import_temporary".to_string(),
@@ -532,6 +544,41 @@ mod tests {
);
}
#[test]
fn normalize_xai_import_keeps_refresh_token_separate_from_api_key() {
let (refresh_token, access_token) =
normalize_provider_import_tokens("xai", Some("xai-refresh-token"), None);
assert_eq!(refresh_token.as_deref(), Some("xai-refresh-token"));
assert!(access_token.is_none());
let (refresh_token, access_token) =
normalize_provider_import_tokens("xai", None, Some("xai-api-key"));
assert!(refresh_token.is_none());
assert_eq!(access_token.as_deref(), Some("xai-api-key"));
}
#[test]
fn builds_xai_auth_config_from_api_key_and_oauth_tokens() {
let (api_key_config, _) =
build_provider_access_token_import_auth_config("xai", "xai-api-key", None, None, None);
assert_eq!(api_key_config.get("auth_method"), Some(&json!("api_key")));
assert_eq!(api_key_config.get("using_api"), Some(&json!(true)));
let (oauth_config, _) = build_provider_access_token_import_auth_config(
"xai",
"xai-access-token",
Some("xai-refresh-token"),
None,
None,
);
assert_eq!(oauth_config.get("auth_method"), Some(&json!("oauth")));
assert_eq!(oauth_config.get("using_api"), Some(&json!(false)));
assert_eq!(
oauth_config.get("refresh_token"),
Some(&json!("xai-refresh-token"))
);
}
#[test]
fn flattens_only_claude_ai_oauth_credentials_and_converts_expiry_ms() {
let mut payload = json!({
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/dispatch.rs
+18
View File
@@ -8,6 +8,7 @@ use super::gemini_cli::refresh_gemini_cli_provider_quota_locally;
use super::grok::refresh_grok_provider_quota_locally;
use super::kiro::refresh_kiro_provider_quota_locally;
use super::windsurf::refresh_windsurf_provider_quota_locally;
use super::xai::refresh_xai_provider_quota_locally;
use crate::handlers::admin::request::AdminAppState;
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
@@ -43,6 +44,7 @@ const PROVIDER_QUOTA_REFRESH_HANDLERS: &[(&str, ProviderQuotaRefreshHandler)] =
("grok", refresh_grok_provider_quota_locally_boxed),
("kiro", refresh_kiro_provider_quota_locally_boxed),
("windsurf", refresh_windsurf_provider_quota_locally_boxed),
("xai", refresh_xai_provider_quota_locally_boxed),
];
pub(crate) async fn refresh_provider_pool_quota_locally(
@@ -174,3 +176,19 @@ fn refresh_windsurf_provider_quota_locally_boxed<'a>(
proxy_override,
))
}
fn refresh_xai_provider_quota_locally_boxed<'a>(
state: &'a AdminAppState<'a>,
provider: &'a StoredProviderCatalogProvider,
endpoint: &'a StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> ProviderQuotaRefreshFuture<'a> {
Box::pin(refresh_xai_provider_quota_locally(
state,
provider,
endpoint,
keys,
proxy_override,
))
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/mod.rs
+1
View File
@@ -7,3 +7,4 @@ pub(crate) mod grok;
pub(crate) mod kiro;
pub(crate) mod shared;
pub(crate) mod windsurf;
pub(crate) mod xai;
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+14
View File
@@ -1715,6 +1715,7 @@ fn provider_quota_url_has_allowed_origin(provider_name: &str, value: &str) -> bo
"gemini_cli" => host == "cloudcode-pa.googleapis.com",
"chatgpt_web" | "codex" => host == "chatgpt.com",
"grok" => host == "grok.com",
"xai" => host == "cli-chat-proxy.grok.com",
"windsurf" => host == "server.codeium.com",
"kiro" => kiro_quota_host_is_allowed(host),
_ => false,
@@ -1814,6 +1815,14 @@ mod tests {
),
("codex", "https://chatgpt.com/backend-api/wham/usage"),
("grok", "https://grok.com/rest/rate-limits"),
(
"xai",
"https://cli-chat-proxy.grok.com/v1/billing?format=credits",
),
(
"xai",
"https://cli-chat-proxy.grok.com/v1/user",
),
(
"windsurf",
"https://server.codeium.com/exa.seat_management_pb.SeatManagementService/GetUserStatus",
@@ -1847,6 +1856,11 @@ mod tests {
"https://chatgpt.com.attacker.test/backend-api/wham/usage",
),
("grok", "https://grok.com.attacker.test/rest/rate-limits"),
(
"xai",
"https://cli-chat-proxy.grok.com.attacker.test/v1/billing",
),
("xai", "https://api.x.ai/v1/billing?format=credits"),
("windsurf", "https://server.codeium.com.attacker.test/quota"),
(
"gemini_cli",
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/xai.rs
+313
View File
@@ -0,0 +1,313 @@
use super::shared::{
build_provider_quota_execution_plan, build_quota_snapshot_payload,
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
extract_execution_error_message, oauth_refresh_auto_removed_result,
persist_provider_quota_refresh_state, quota_key_auto_removed,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::parse_xai_billing_response;
use aether_admin::provider::redaction::admin_provider_metadata_bucket_safe_json;
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_provider_pool::{build_xai_pool_billing_request, build_xai_pool_user_request};
use aether_provider_transport::xai::{
extract_xai_user_id_from_auth_config, extract_xai_user_id_from_value, xai_auth_uses_api,
};
use serde_json::{json, Value};
use std::time::{SystemTime, UNIX_EPOCH};
async fn execute_xai_quota_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
spec: aether_provider_pool::ProviderPoolQuotaRequestSpec,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = match proxy_override {
Some(proxy) => Some(proxy.clone()),
None => {
state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = build_provider_quota_execution_plan(
transport,
spec,
proxy,
state.resolve_transport_profile(transport),
timeouts,
);
execute_provider_quota_plan(state, transport, plan, "xai").await
}
fn xai_authorization_from_header(authorization: &(String, String)) -> (String, String) {
authorization.clone()
}
fn enrich_xai_subscription_title(mut metadata: Value, auth_config: Option<&str>) -> Value {
if metadata
.get("subscription_title")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty())
{
return metadata;
}
let Some(config) = auth_config
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| serde_json::from_str::<Value>(value).ok())
else {
return metadata;
};
let title = ["subscription_tier", "subscriptionTier", "tier", "plan"]
.iter()
.find_map(|field| {
config
.get(*field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
});
if let Some(title) = title {
if let Some(object) = metadata.as_object_mut() {
object.insert("subscription_title".to_string(), json!(title));
}
}
metadata
}
pub(crate) async fn refresh_xai_provider_quota_locally(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> Result<Option<serde_json::Value>, GatewayError> {
let mut results = Vec::new();
let mut success_count = 0usize;
let mut failed_count = 0usize;
let mut auto_removed_count = 0usize;
for key in keys {
let transport = match state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
.await?
{
Some(transport) => transport,
None => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "Provider transport snapshot unavailable",
}));
continue;
}
};
if xai_auth_uses_api(
transport.key.auth_type.as_str(),
transport.key.decrypted_auth_config.as_deref(),
) {
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "skipped",
"message": "xAI API Key 账号没有 Grok Build 订阅额度接口，请使用设备授权账号查询额度。",
}));
continue;
}
let authorization = match state.resolve_local_oauth_header_auth(&transport).await? {
Some(auth) => auth,
_ => {
if quota_key_auto_removed(state, &key.id).await? {
auto_removed_count += 1;
results.push(oauth_refresh_auto_removed_result(&key));
continue;
}
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "缺少 OAuth 认证信息，请先授权/刷新 Token",
}));
continue;
}
};
let fallback_user_id =
extract_xai_user_id_from_auth_config(transport.key.decrypted_auth_config.as_deref());
let user_id = match execute_xai_quota_plan(
state,
&transport,
build_xai_pool_user_request(
&transport.key.id,
xai_authorization_from_header(&authorization),
),
proxy_override.as_ref(),
)
.await?
{
ProviderQuotaExecutionOutcome::Response(result) if result.status_code == 200 => result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
.and_then(extract_xai_user_id_from_value)
.or(fallback_user_id),
_ => fallback_user_id,
};
let result = match execute_xai_quota_plan(
state,
&transport,
build_xai_pool_billing_request(
&transport.key.id,
xai_authorization_from_header(&authorization),
user_id.as_deref(),
),
proxy_override.as_ref(),
)
.await?
{
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(_) => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "xAI billing 请求执行失败",
"status_code": 502,
}));
continue;
}
};
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs())
.unwrap_or(0);
let mut metadata_update = None::<serde_json::Value>;
let (mut oauth_invalid_at_unix_secs, mut oauth_invalid_reason) =
quota_refresh_success_invalid_state(&key);
let mut status = "error".to_string();
let mut message = None::<String>;
if result.status_code == 200 {
if let Some(body_json) = result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
{
metadata_update =
parse_xai_billing_response(body_json, now_unix_secs).map(|metadata| {
json!({
"xai": enrich_xai_subscription_title(
metadata,
transport.key.decrypted_auth_config.as_deref(),
)
})
});
if metadata_update.is_some() {
status = "success".to_string();
} else {
status = "no_metadata".to_string();
message = Some("响应中未包含可用的 Grok Build 额度信息".to_string());
}
} else {
status = "no_metadata".to_string();
message = Some("响应中未包含配额信息".to_string());
}
} else {
message = Some(
extract_execution_error_message(&result)
.unwrap_or_else(|| format!("xAI billing 返回状态码 {}", result.status_code)),
);
if result.status_code == 401 || result.status_code == 403 {
let reason = message
.clone()
.unwrap_or_else(|| "账户访问被禁止".to_string());
oauth_invalid_at_unix_secs = Some(now_unix_secs);
oauth_invalid_reason = Some(format!("账户访问被禁止: {reason}"));
status = if result.status_code == 401 {
"unauthorized".to_string()
} else {
"forbidden".to_string()
};
}
}
if !persist_provider_quota_refresh_state(
state,
&key.id,
metadata_update.as_ref(),
oauth_invalid_at_unix_secs,
oauth_invalid_reason,
None,
)
.await?
{
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "Key 状态写入失败",
}));
continue;
}
if status == "success" {
success_count += 1;
} else {
failed_count += 1;
}
let mut payload = serde_json::Map::new();
payload.insert("key_id".to_string(), json!(key.id));
payload.insert("key_name".to_string(), json!(key.name));
payload.insert("status".to_string(), json!(status));
if let Some(message) = message {
payload.insert("message".to_string(), json!(message));
}
if let Some(metadata) = metadata_update.as_ref().and_then(|value| value.get("xai")) {
payload.insert(
"metadata".to_string(),
admin_provider_metadata_bucket_safe_json("xai", Some(metadata)),
);
}
if let Some(quota_snapshot) = build_quota_snapshot_payload(
"xai",
key.status_snapshot.as_ref(),
metadata_update.as_ref(),
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
results.push(serde_json::Value::Object(payload));
}
Ok(Some(json!({
"success": success_count,
"failed": failed_count,
"total": results.len(),
"results": results,
"message": format!("已处理 {} 个 Key", results.len()),
"auto_removed": auto_removed_count,
})))
}
apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs
+32
View File
@@ -932,6 +932,13 @@ fn admin_pool_build_account_quota(
return Some(account_quota);
}
}
"xai" => {
if let Some(account_quota) =
admin_pool_build_kiro_account_quota_from_snapshot(quota_snapshot)
{
return Some(account_quota);
}
}
"chatgpt_web" => {
if let Some(account_quota) =
admin_pool_build_chatgpt_web_account_quota_from_snapshot(quota_snapshot)
@@ -1591,4 +1598,29 @@ mod tests {
Some("Auto剩余 40.0% (60/150) | Heavy剩余 0.0% (0/20)".to_string())
);
}
#[test]
fn xai_account_quota_is_rendered_as_remaining_percent() {
let quota_snapshot = json!({
"provider_type": "xai",
"code": "ok",
"exhausted": false,
"plan_type": "SuperGrok",
"windows": [
{
"code": "usage",
"label": "周额度",
"scope": "account",
"used_ratio": 0.46,
"remaining_ratio": 0.54
}
]
});
let quota_snapshot = quota_snapshot.as_object().unwrap();
assert_eq!(
admin_pool_build_account_quota("xai", Some(quota_snapshot)),
Some("剩余 54.0%".to_string())
);
}
}
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs
+5
View File
@@ -3511,6 +3511,11 @@ async fn provider_query_execute_standard_test_candidate(
codex_model_capabilities.as_ref(),
);
}
crate::provider_transport::insert_cli_identity_headers_if_needed(
&transport,
provider_api_format,
&mut request_headers,
);
if !uses_vertex_query_auth {
if let (Some(auth_header), Some(auth_value)) =
(auth_header.as_deref(), auth_value.as_deref())
apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs
+10 -2
View File
@@ -4,9 +4,9 @@ pub(crate) fn normalize_provider_type_input(value: &str) -> Result<String, Strin
let normalized = value.trim().to_ascii_lowercase();
match normalized.as_str() {
"custom" | "claude_code" | "kiro" | "codex" | "chatgpt_web" | "gemini_cli"
| "antigravity" | "vertex_ai" | "grok" | "windsurf" => Ok(normalized),
| "antigravity" | "vertex_ai" | "grok" | "windsurf" | "xai" => Ok(normalized),
_ => Err(
"provider_type 仅支持 custom / claude_code / kiro / codex / chatgpt_web / gemini_cli / antigravity / vertex_ai / grok / windsurf"
"provider_type 仅支持 custom / claude_code / kiro / codex / chatgpt_web / gemini_cli / antigravity / vertex_ai / grok / windsurf / xai"
.to_string(),
),
}
@@ -405,6 +405,14 @@ mod tests {
);
}
#[test]
fn normalize_provider_type_supports_xai() {
assert_eq!(
normalize_provider_type_input(" xAI ").expect("type should normalize"),
"xai"
);
}
#[test]
fn normalize_api_format_list_dedupes_canonical_formats() {
assert_eq!(
apps/aether-gateway/src/handlers/proxy/websocket/responses/continuation.rs
+33 -1
View File
@@ -90,6 +90,8 @@ pub(super) struct ResponsesWebSocketContinuationRecord {
/// request JSON can never set it.
#[serde(default)]
deepseek_opaque_reasoning_replay: bool,
#[serde(default)]
xai_encrypted_reasoning_replay: bool,
/// A prior turn stored PII sentinels whose restore mapping exists only on
/// the original downstream socket. Such a chain cannot safely resume on a
/// new socket without leaking sentinels, so lookup succeeds but bootstrap
@@ -122,6 +124,10 @@ impl ResponsesWebSocketContinuationRecord {
normalization.reasoning_replay_policy(),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
),
xai_encrypted_reasoning_replay: matches!(
normalization.reasoning_replay_policy(),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::XaiEncrypted
),
has_connection_local_redaction,
responses_lite_static_config,
};
@@ -156,7 +162,9 @@ impl ResponsesWebSocketContinuationRecord {
pub(super) fn reasoning_replay_policy(
&self,
) -> crate::ai_serving::OpenAiResponsesReasoningReplayPolicy {
if self.deepseek_opaque_reasoning_replay {
if self.xai_encrypted_reasoning_replay {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::XaiEncrypted
} else if self.deepseek_opaque_reasoning_replay {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
} else {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
@@ -476,6 +484,7 @@ mod tests {
binding_fingerprint: [7; 32],
normalization_fingerprint: [9; 32],
deepseek_opaque_reasoning_replay: false,
xai_encrypted_reasoning_replay: false,
has_connection_local_redaction: false,
responses_lite_static_config: Some(ResponsesLiteStaticConfig::from_response_create(
&json!({
@@ -714,6 +723,29 @@ mod tests {
assert_eq!(decoded, record());
}
#[test]
fn serialized_record_preserves_xai_replay_policy_and_reads_legacy_records() {
let mut expected = record();
expected.xai_encrypted_reasoning_replay = true;
let mut serialized = serde_json::to_value(&expected).unwrap();
let decoded: ResponsesWebSocketContinuationRecord =
serde_json::from_value(serialized.clone()).unwrap();
assert_eq!(
decoded.reasoning_replay_policy(),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::XaiEncrypted
);
serialized
.as_object_mut()
.unwrap()
.remove("xai_encrypted_reasoning_replay");
let legacy: ResponsesWebSocketContinuationRecord =
serde_json::from_value(serialized).unwrap();
assert_eq!(
legacy.reasoning_replay_policy(),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
);
}
#[test]
fn serialized_record_preserves_only_the_server_derived_reasoning_replay_policy_bit() {
let mut expected = record();
apps/aether-gateway/src/handlers/shared/catalog.rs
+200
View File
@@ -1388,6 +1388,168 @@ fn build_kiro_quota_status_snapshot(
}))
}
fn build_xai_quota_status_snapshot(
upstream_metadata: Option<&Value>,
source: &str,
) -> Option<Value> {
let metadata = provider_quota_metadata_bucket(upstream_metadata, "xai")?;
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
let usage_limit = metadata
.get("usage_limit")
.and_then(admin_provider_quota_pure::coerce_json_f64);
let current_usage = metadata
.get("current_usage")
.and_then(admin_provider_quota_pure::coerce_json_f64);
let remaining = metadata
.get("remaining")
.and_then(admin_provider_quota_pure::coerce_json_f64);
let usage_ratio = metadata
.get("usage_percentage")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.map(|value| (value / 100.0).clamp(0.0, 1.0))
.or_else(|| {
current_usage
.zip(usage_limit)
.and_then(|(current_usage, usage_limit)| {
(usage_limit > 0.0).then_some((current_usage / usage_limit).clamp(0.0, 1.0))
})
});
let remaining_ratio = usage_ratio.map(|value| (1.0 - value).max(0.0));
let next_reset_at = provider_quota_timestamp_unix_secs(metadata.get("next_reset_at"));
let reset_seconds = quota_window_reset_seconds(observed_at_unix_secs, next_reset_at);
let plan_type = metadata
.get("subscription_title")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let period_type = metadata
.get("period_type")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let usage_label = match period_type.as_deref() {
Some("monthly") => "月额度",
Some("weekly") => "周额度",
_ => "额度",
};
let mut windows = Vec::new();
if usage_ratio.is_some()
|| remaining.is_some()
|| usage_limit.is_some()
|| current_usage.is_some()
|| next_reset_at.is_some()
{
windows.push(json!({
"code": "usage",
"label": usage_label,
"scope": "account",
"unit": if usage_limit.is_some() { "usd" } else { "percent" },
"used_ratio": usage_ratio,
"remaining_ratio": remaining_ratio,
"used_value": current_usage,
"remaining_value": remaining,
"limit_value": usage_limit,
"reset_at": next_reset_at,
"reset_seconds": reset_seconds,
}));
}
let prepaid_balance = metadata
.get("prepaid_balance")
.and_then(admin_provider_quota_pure::coerce_json_f64);
if prepaid_balance.is_some_and(|value| value > 0.0) {
windows.push(json!({
"code": "prepaid",
"label": "预付额度",
"scope": "account",
"unit": "usd",
"used_ratio": serde_json::Value::Null,
"remaining_ratio": serde_json::Value::Null,
"remaining_value": prepaid_balance,
"reset_at": serde_json::Value::Null,
"reset_seconds": serde_json::Value::Null,
}));
}
let on_demand_cap = metadata
.get("on_demand_cap")
.and_then(admin_provider_quota_pure::coerce_json_f64);
let on_demand_used = metadata
.get("on_demand_used")
.and_then(admin_provider_quota_pure::coerce_json_f64);
let on_demand_enabled = metadata
.get("on_demand_enabled")
.and_then(admin_provider_quota_pure::coerce_json_bool)
!= Some(false);
if on_demand_enabled && on_demand_cap.is_some_and(|value| value > 0.0) {
let on_demand_remaining = on_demand_cap
.zip(on_demand_used)
.map(|(cap, used)| (cap - used).max(0.0));
let on_demand_ratio = on_demand_cap
.zip(on_demand_used)
.and_then(|(cap, used)| (cap > 0.0).then_some((used / cap).clamp(0.0, 1.0)));
windows.push(json!({
"code": "on_demand",
"label": "按需额度",
"scope": "account",
"unit": "usd",
"used_ratio": on_demand_ratio,
"remaining_ratio": on_demand_ratio.map(|value| (1.0 - value).max(0.0)),
"used_value": on_demand_used,
"remaining_value": on_demand_remaining,
"limit_value": on_demand_cap,
"reset_at": serde_json::Value::Null,
"reset_seconds": serde_json::Value::Null,
}));
}
if windows.is_empty() && plan_type.is_none() && observed_at_unix_secs.is_none() {
return None;
}
let prepaid_available = prepaid_balance.is_some_and(|value| value > 0.0);
let on_demand_available = on_demand_enabled
&& on_demand_cap.is_some_and(|value| value > 0.0)
&& on_demand_used
.zip(on_demand_cap)
.is_some_and(|(used, cap)| used < cap);
let usage_exhausted = remaining.is_some_and(|value| value <= 0.0)
|| usage_ratio.is_some_and(|value| value >= 1.0 - 1e-6);
let exhausted = usage_exhausted && !prepaid_available && !on_demand_available;
let reason = if exhausted {
Some("额度已耗尽".to_string())
} else {
None
};
let label = if exhausted {
Some("额度耗尽")
} else {
None
};
let code = if exhausted { "exhausted" } else { "ok" };
Some(json!({
"version": 2,
"provider_type": "xai",
"code": code,
"label": label,
"reason": reason,
"freshness": "fresh",
"source": source,
"observed_at": observed_at_unix_secs,
"exhausted": exhausted,
"usage_ratio": usage_ratio,
"updated_at": observed_at_unix_secs,
"reset_at": next_reset_at,
"reset_seconds": reset_seconds,
"plan_type": plan_type,
"windows": windows,
}))
}
fn build_chatgpt_web_quota_status_snapshot(
upstream_metadata: Option<&Value>,
source: &str,
@@ -2255,6 +2417,7 @@ pub(crate) fn sync_provider_key_quota_status_snapshot(
let mut quota = match normalized_provider_type.as_str() {
"codex" => build_codex_quota_status_snapshot(upstream_metadata, source),
"kiro" => build_kiro_quota_status_snapshot(upstream_metadata, source),
"xai" => build_xai_quota_status_snapshot(upstream_metadata, source),
"chatgpt_web" => build_chatgpt_web_quota_status_snapshot(upstream_metadata, source),
"windsurf" => build_windsurf_quota_status_snapshot(upstream_metadata, source),
"antigravity" => build_antigravity_quota_status_snapshot(upstream_metadata, source),
@@ -3622,6 +3785,43 @@ mod tests {
assert_eq!(auto.get("used_value"), Some(&json!(90.0)));
}
#[test]
fn provider_key_status_snapshot_payload_backfills_xai_weekly_credits() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"xai": {
"updated_at": 1_778_067_246u64,
"usage_percentage": 46.0,
"period_type": "weekly",
"next_reset_at": 1_778_157_172u64,
"subscription_title": "SuperGrok",
"prepaid_balance": 0.0,
"on_demand_cap": 0.0,
"on_demand_used": 0.0
}
}));
let payload = provider_key_status_snapshot_payload(&key, "xai");
let quota = payload
.get("quota")
.and_then(Value::as_object)
.expect("quota snapshot should be object");
let windows = quota
.get("windows")
.and_then(Value::as_array)
.expect("xai quota windows should exist");
assert_eq!(quota.get("provider_type"), Some(&json!("xai")));
assert_eq!(quota.get("code"), Some(&json!("ok")));
assert_eq!(quota.get("exhausted"), Some(&json!(false)));
assert_eq!(quota.get("plan_type"), Some(&json!("SuperGrok")));
assert_eq!(quota.get("usage_ratio"), Some(&json!(0.46)));
assert_eq!(quota.get("reset_at"), Some(&json!(1_778_157_172u64)));
assert_eq!(windows.len(), 1);
assert_eq!(windows[0].get("code"), Some(&json!("usage")));
assert_eq!(windows[0].get("label"), Some(&json!("周额度")));
}
#[test]
fn provider_key_status_snapshot_payload_backfills_gemini_cli_account_credits() {
let mut key = sample_catalog_key();
apps/aether-gateway/src/provider_key_auth.rs
+17
View File
@@ -172,6 +172,7 @@ fn provider_uses_bearer_oauth_runtime(provider_type: &str) -> bool {
| "antigravity"
| "kiro"
| "windsurf"
| "xai"
)
}
@@ -406,6 +407,22 @@ mod tests {
);
}
#[test]
fn recognizes_xai_oauth_as_bearer_runtime() {
let semantics = provider_key_auth_semantics(&sample_key("oauth"), "xai");
assert!(semantics.oauth_managed());
assert!(semantics.can_refresh_oauth());
assert_eq!(
semantics.credential_kind(),
ProviderKeyCredentialKind::OAuthSession
);
assert_eq!(
semantics.runtime_auth_kind(),
ProviderKeyRuntimeAuthKind::Bearer
);
}
#[test]
fn refresh_capability_requires_stored_refresh_token() {
let semantics = provider_key_auth_semantics(&sample_key("oauth"), "codex");
apps/aether-gateway/src/tests/architecture/admin_provider.rs
+2
View File
@@ -1789,6 +1789,7 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
"pub(crate) mod dispatch;",
"pub(crate) mod kiro;",
"pub(crate) mod shared;",
"pub(crate) mod xai;",
] {
assert!(
quota_mod.contains(pattern),
@@ -1861,6 +1862,7 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
"refresh_antigravity_provider_quota_locally",
"refresh_gemini_cli_provider_quota_locally",
"refresh_chatgpt_web_provider_quota_locally",
"refresh_xai_provider_quota_locally",
] {
assert!(
quota_dispatch.contains(pattern),
apps/aether-gateway/src/tests/architecture/ai_serving.rs
+10
View File
@@ -1452,6 +1452,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
"GeminiCliProviderPoolAdapter",
"KiroProviderPoolAdapter",
"ChatGptWebProviderPoolAdapter",
"XaiProviderPoolAdapter",
"CLAUDE_CODE_PROVIDER_POOL_ADAPTER",
"VERTEX_AI_PROVIDER_POOL_ADAPTER",
"provider_types_for_capability",
@@ -1478,6 +1479,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
"pub mod gemini_cli;",
"pub mod kiro;",
"pub mod chatgpt_web;",
"pub mod xai;",
] {
assert!(
provider_pool_providers.contains(pattern),
@@ -1513,6 +1515,14 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
"crates/aether-provider/pool/src/providers/kiro.rs",
vec!["KiroProviderPoolAdapter", "quota_exhausted_from_bucket"],
),
(
"crates/aether-provider/pool/src/providers/xai.rs",
vec![
"XaiProviderPoolAdapter",
"build_xai_pool_billing_request",
"quota_exhausted_from_bucket",
],
),
(
"crates/aether-provider/pool/src/providers/chatgpt_web.rs",
vec![
apps/aether-gateway/src/tests/control/admin/oauth.rs
+291
View File
@@ -983,6 +983,297 @@ async fn gateway_rejects_generic_oauth_start_for_windsurf_provider_impl() {
);
}
#[test]
fn gateway_rejects_generic_oauth_start_for_xai_provider() {
run_admin_oauth_test(
"gateway_rejects_generic_oauth_start_for_xai_provider",
gateway_rejects_generic_oauth_start_for_xai_provider_impl,
);
}
async fn gateway_rejects_generic_oauth_start_for_xai_provider_impl() {
let mut provider = sample_provider("provider-xai", "xai", 10);
provider.provider_type = "xai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-xai/start",
None,
)
.await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert!(
payload["detail"].as_str().is_some_and(|detail| {
detail.contains("设备授权") || detail.contains("导入凭据")
}),
"payload={payload}"
);
}
#[test]
fn gateway_handles_admin_provider_oauth_device_authorize_for_xai() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_device_authorize_for_xai",
gateway_handles_admin_provider_oauth_device_authorize_for_xai_impl,
);
}
async fn gateway_handles_admin_provider_oauth_device_authorize_for_xai_impl() {
let authorize_hits = Arc::new(Mutex::new(0usize));
let authorize_hits_clone = Arc::clone(&authorize_hits);
let oidc_server = Router::new().fallback(any(move |_request: Request| {
let authorize_hits_inner = Arc::clone(&authorize_hits_clone);
async move {
*authorize_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"device_code": "xai-device-code",
"user_code": "XAI-CODE",
"verification_uri": "https://auth.x.ai/activate",
"verification_uri_complete": "https://auth.x.ai/activate?user_code=XAI-CODE",
"expires_in": 600,
"interval": 5,
}))
}
}));
let mut provider = sample_provider("provider-xai", "xai", 10);
provider.provider_type = "xai".to_string();
let endpoint = sample_endpoint(
"endpoint-xai-responses",
"provider-xai",
"openai:responses",
"https://cli-chat-proxy.grok.com/v1",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (oidc_url, oidc_handle) = start_server(oidc_server).await;
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
))
.with_provider_oauth_token_url_for_tests(
"xai_device",
format!("{oidc_url}/oauth2/device/code"),
)
.with_provider_oauth_token_url_for_tests("xai", format!("{oidc_url}/oauth2/token"));
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-xai/device-authorize",
Some(json!({ "proxy_node_id": "proxy-node-xai" })),
)
.await;
let status = response.status();
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
let session_id = payload["session_id"]
.as_str()
.expect("session_id should exist")
.to_string();
assert_eq!(payload["user_code"], "XAI-CODE");
assert_eq!(payload["verification_uri"], "https://auth.x.ai/activate");
assert_eq!(
payload["verification_uri_complete"],
"https://auth.x.ai/activate?user_code=XAI-CODE"
);
assert_eq!(payload["auth_type"], "device");
assert!(payload.get("callback_required").is_none() || payload["callback_required"] == false);
assert_eq!(*authorize_hits.lock().expect("mutex should lock"), 1);
let stored = state
.load_provider_oauth_device_session_for_tests(&format!("device_auth_session:{session_id}"))
.expect("device session should be stored");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["provider_id"], "provider-xai");
assert_eq!(stored["device_code"], "xai-device-code");
assert_eq!(stored["auth_type"], "device");
assert_eq!(stored["redirect_uri"], format!("{oidc_url}/oauth2/token"));
assert_eq!(stored["proxy_node_id"], "proxy-node-xai");
assert_eq!(stored["status"], "pending");
oidc_handle.abort();
}
#[test]
fn gateway_handles_admin_provider_oauth_device_poll_for_xai() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_device_poll_for_xai",
gateway_handles_admin_provider_oauth_device_poll_for_xai_impl,
);
}
async fn gateway_handles_admin_provider_oauth_device_poll_for_xai_impl() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let access_token = sample_kiro_device_access_token("[email protected]");
let id_token = access_token.clone();
let token_server = Router::new().fallback(any(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
let access_token = access_token.clone();
let id_token = id_token.clone();
async move {
let hit = {
let mut hits = token_hits_inner.lock().expect("mutex should lock");
*hits += 1;
*hits
};
if hit == 1 {
return (
StatusCode::BAD_REQUEST,
Json(json!({ "error": "authorization_pending" })),
)
.into_response();
}
Json(json!({
"access_token": access_token,
"refresh_token": "xai-refresh-token",
"token_type": "Bearer",
"expires_in": 3600,
"id_token": id_token,
}))
.into_response()
}
}));
let mut provider = sample_provider("provider-xai", "xai", 10);
provider.provider_type = "xai".to_string();
let endpoint = sample_endpoint(
"endpoint-xai-responses",
"provider-xai",
"openai:responses",
"https://cli-chat-proxy.grok.com/v1",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (token_url, token_handle) = start_server(token_server).await;
let resolved_token_url = format!("{token_url}/oauth2/token");
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_device_session_entry_for_tests(
"session-xai",
json!({
"provider_id": "provider-xai",
"region": "",
"client_id": "b1a00492-073a-47ea-816f-4c329264a828",
"client_secret": "",
"device_code": "xai-device-code",
"auth_type": "device",
"social_provider": null,
"code_verifier": null,
"redirect_uri": resolved_token_url,
"machine_id": null,
"interval": 5,
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": null,
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
"error_msg": null,
}),
)
.with_provider_oauth_token_url_for_tests("xai", resolved_token_url.clone());
let pending = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-xai/device-poll",
Some(json!({ "session_id": "session-xai" })),
)
.await;
let pending_body = to_bytes(pending.into_body(), usize::MAX)
.await
.expect("pending body should read");
let pending_payload: serde_json::Value =
serde_json::from_slice(&pending_body).expect("pending json should parse");
assert_eq!(
pending_payload["status"], "pending",
"payload={pending_payload}"
);
let authorized = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-xai/device-poll",
Some(json!({ "session_id": "session-xai" })),
)
.await;
let status = authorized.status();
let body = to_bytes(authorized.into_body(), usize::MAX)
.await
.expect("authorized body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["status"], "authorized");
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["replaced"], false);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 2);
let stored = state
.load_provider_oauth_device_session_for_tests("device_auth_session:session-xai")
.expect("device session should persist");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["status"], "authorized");
let key_id = stored["key_id"]
.as_str()
.expect("key_id should be stored")
.to_string();
assert_eq!(payload["key_id"], key_id);
let persisted = provider_catalog_repository
.list_keys_by_ids(std::slice::from_ref(&key_id))
.await
.expect("keys should load")
.into_iter()
.next()
.expect("persisted key should exist");
assert_eq!(persisted.auth_type, "oauth");
let decrypted_auth_config = decrypt_persisted_provider_auth_config(&persisted);
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["provider_type"], "xai");
assert_eq!(auth_config["auth_method"], "oauth");
assert_eq!(auth_config["using_api"], false);
assert_eq!(auth_config["email"], "[email protected]");
token_handle.abort();
}
#[test]
fn gateway_handles_admin_provider_oauth_device_poll_for_windsurf_one_time_token() {
run_admin_oauth_test(