Merge remote-tracking branch 'origin/main' into fix/gemini-cli-v1internal

# Conflicts:
#	apps/aether-gateway/src/ai_serving/transport.rs
#	apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs
#	apps/aether-gateway/src/handlers/shared/catalog.rs
#	crates/aether-admin/src/provider/quota.rs
#	crates/aether-model-fetch/src/strategy.rs
#	crates/aether-provider-pool/src/lib.rs
#	crates/aether-provider-pool/src/service.rs
#	crates/aether-provider-transport/src/provider_types.rs
#	frontend/src/features/providers/components/ProviderDetailDrawer.vue
#	frontend/src/utils/__tests__/providerKeyQuota.spec.ts
#	frontend/src/utils/providerKeyQuota.ts
#	frontend/src/views/admin/PoolManagement.vue
This commit is contained in:
Mas0nShi
2026-05-22 17:13:57 +08:00
220 changed files with 30324 additions and 2776 deletions
Generated
+17
View File
@@ -858,6 +858,15 @@ dependencies = [
"tracing",
]
[[package]]
name = "backon"
version = "1.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
dependencies = [
"fastrand",
]
[[package]]
name = "base64"
version = "0.21.7"
@@ -1637,6 +1646,12 @@ dependencies = [
"regex",
]
[[package]]
name = "fastrand"
version = "2.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6"
[[package]]
name = "filedescriptor"
version = "0.8.3"
@@ -3389,8 +3404,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e37ec3fd44bea2ec947ba6cc7634d7999a6590aca7c35827c250bc0de502bda6"
dependencies = [
"arc-swap",
"backon",
"bytes",
"combine",
"futures-channel",
"futures-util",
"itoa",
"num-bigint",
+1 -1
View File
@@ -82,7 +82,7 @@ hmac = "0.12"
http = "1"
pbkdf2 = { version = "0.12", default-features = false, features = ["hmac"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls", "http2", "socks"] }
redis = { version = "0.28", default-features = false, features = ["tokio-comp", "script", "streams"] }
redis = { version = "0.28", default-features = false, features = ["tokio-comp", "script", "streams", "connection-manager"] }
regex = "1"
rustls = { version = "0.23", features = ["ring"] }
semver = "1"
+5
View File
@@ -335,6 +335,11 @@ if ! command -v curl >/dev/null 2>&1; then
exit 1
fi
if [ -z "$${RUSTC_WRAPPER:-}" ] && command -v sccache >/dev/null 2>&1; then
export RUSTC_WRAPPER="$$(command -v sccache)"
echo "=> 启用 Rust 编译缓存: $${RUSTC_WRAPPER}"
fi
if ! ensure_dev_infra; then
exit 1
fi
+5 -4
View File
@@ -7,10 +7,11 @@ pub(crate) use crate::handlers::admin::{
provider_quota_refresh_endpoint_for_provider, provider_type_supports_quota_refresh,
reconcile_admin_fixed_provider_template_endpoints,
refresh_provider_oauth_account_state_after_update, refresh_provider_pool_quota_locally,
update_existing_provider_oauth_catalog_key, AdminAppState,
AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError, AdminRequestContext,
AdminRouteRequest, AdminRouteResponse, AdminRouteResult, AdminStatsTimeRange,
AdminStatsUsageFilter, OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
store_admin_provider_ops_balance_cache, update_existing_provider_oauth_catalog_key,
AdminAppState, AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError,
AdminRequestContext, AdminRouteRequest, AdminRouteResponse, AdminRouteResult,
AdminStatsTimeRange, AdminStatsUsageFilter, OAUTH_ACCOUNT_BLOCK_PREFIX,
OAUTH_REQUEST_FAILED_PREFIX,
};
use crate::handlers::admin::{
+5 -5
View File
@@ -44,11 +44,11 @@ pub(crate) use aether_ai_formats::api::{
build_core_error_body_for_client_format, convert_standard_chat_response,
core_error_background_report_kind, core_error_default_client_api_format,
core_success_background_report_kind, encode_kiro_sse_events,
implicit_sync_finalize_report_kind, is_core_error_finalize_kind,
normalize_provider_private_report_context, normalize_provider_private_response_value,
provider_private_response_allows_sync_finalize, resolve_claude_stream_spec,
resolve_claude_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
extract_provider_private_stream_error_body, implicit_sync_finalize_report_kind,
is_core_error_finalize_kind, normalize_provider_private_report_context,
normalize_provider_private_response_value, provider_private_response_allows_sync_finalize,
resolve_claude_stream_spec, resolve_claude_sync_spec, resolve_gemini_stream_spec,
resolve_gemini_sync_spec, resolve_local_image_stream_spec, resolve_local_image_sync_spec,
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
resolve_openai_embedding_sync_spec, sanitize_request_path_and_query, AiControlPlanRequest,
CanonicalContentPart, CanonicalStreamEvent, CanonicalStreamFrame, ClaudeClientEmitter,
@@ -85,7 +85,7 @@ fn injects_stable_prompt_cache_key_for_codex_requests() {
assert_eq!(
body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"53363264-dbb0-5f9d-b9c7-3e92c45c5bdf"
);
}
@@ -24,9 +24,13 @@ use crate::ai_serving::transport::kiro::{
use crate::ai_serving::transport::{
build_grok_browser_headers, build_grok_upstream_url, build_kiro_cross_format_upstream_url,
build_openai_image_headers, build_openai_image_upstream_url,
build_standard_provider_request_headers, openai_image_transport_unsupported_reason,
resolve_grok_session_auth, resolve_openai_image_auth, GrokHeaderInput,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
build_standard_provider_request_headers, build_windsurf_cascade_headers,
build_windsurf_cascade_request_body, build_windsurf_cascade_upstream_url,
is_windsurf_provider_transport,
local_windsurf_request_transport_unsupported_reason_with_network,
openai_image_transport_unsupported_reason, resolve_grok_session_auth,
resolve_openai_image_auth, GrokHeaderInput, ProviderOpenAiImageHeadersInput,
StandardProviderRequestHeadersInput, GROK_CHAT_PATH, WINDSURF_ENVELOPE_NAME,
};
use crate::ai_serving::{
build_openai_image_request_body_from_gemini_image_request, gemini_request_is_image_generation,
@@ -61,6 +65,134 @@ fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
)
}
fn provider_preserves_claude_thinking_signatures(provider_type: &str, base_url: &str) -> bool {
let provider_type = provider_type.trim().to_ascii_lowercase();
let base_url = base_url.trim().to_ascii_lowercase();
let is_bedrock_runtime_url = base_url.contains("bedrock-runtime")
&& (base_url.contains("amazonaws.com")
|| base_url.contains("amazonaws.com.cn")
|| base_url.contains("api.aws"));
matches!(
provider_type.as_str(),
"anthropic" | "claude_code" | "bedrock" | "aws_bedrock" | "amazon_bedrock"
) || base_url.contains("api.anthropic.com")
|| is_bedrock_runtime_url
}
fn sanitize_claude_thinking_block(block: Value) -> (Option<Value>, bool) {
let Some(object) = block.as_object() else {
return (Some(block), false);
};
let block_type = object
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
match block_type {
"thinking" => {
let thinking_text = object
.get("thinking")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if thinking_text.is_empty() {
(None, true)
} else {
(
Some(serde_json::json!({
"type": "text",
"text": thinking_text,
})),
true,
)
}
}
"redacted_thinking" => (None, true),
_ => (Some(block), false),
}
}
fn sanitize_claude_message_content_for_non_native_thinking(content: &mut Value) -> bool {
const OMITTED_THINKING_TEXT: &str = "Previous thinking omitted.";
if content.is_object() {
let original = std::mem::take(content);
let (sanitized, changed) = sanitize_claude_thinking_block(original);
if changed {
*content = sanitized.unwrap_or_else(|| {
serde_json::json!({
"type": "text",
"text": OMITTED_THINKING_TEXT,
})
});
}
return changed;
}
let Some(blocks) = content.as_array_mut() else {
return false;
};
let original_blocks = std::mem::take(blocks);
let mut changed = false;
let mut sanitized_blocks = Vec::with_capacity(original_blocks.len());
for block in original_blocks {
let (sanitized, block_changed) = sanitize_claude_thinking_block(block);
changed |= block_changed;
if let Some(sanitized) = sanitized {
sanitized_blocks.push(sanitized);
}
}
if changed && sanitized_blocks.is_empty() {
sanitized_blocks.push(serde_json::json!({
"type": "text",
"text": OMITTED_THINKING_TEXT,
}));
}
*blocks = sanitized_blocks;
changed
}
fn sanitize_claude_request_thinking_signatures_for_non_native(body_json: &mut Value) -> bool {
body_json
.get_mut("messages")
.and_then(Value::as_array_mut)
.map(|messages| {
messages.iter_mut().fold(false, |changed, message| {
let is_assistant = message
.get("role")
.and_then(Value::as_str)
.is_some_and(|role| role.trim().eq_ignore_ascii_case("assistant"));
if !is_assistant {
return changed;
}
let content_changed = message
.get_mut("content")
.is_some_and(sanitize_claude_message_content_for_non_native_thinking);
changed || content_changed
})
})
.unwrap_or(false)
}
fn apply_non_native_claude_thinking_signature_compat(
provider_request_body: &mut Value,
provider_api_format: &str,
transport: &GatewayProviderTransportSnapshot,
) {
if crate::ai_serving::normalize_api_format_alias(provider_api_format) != "claude:messages" {
return;
}
if provider_preserves_claude_thinking_signatures(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
) {
return;
}
let _ = sanitize_claude_request_thinking_signatures_for_non_native(provider_request_body);
}
pub(crate) async fn resolve_local_standard_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
@@ -198,11 +330,18 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
return None;
}
if let Some(skip_reason) = crate::ai_serving::request_pair_transport_unsupported_reason(
transport,
spec_metadata.api_format,
provider_api_format,
) {
let is_windsurf_cascade =
provider_api_format == "openai:chat" && is_windsurf_provider_transport(transport);
let transport_unsupported_reason = if is_windsurf_cascade {
local_windsurf_request_transport_unsupported_reason_with_network(transport)
} else {
crate::ai_serving::request_pair_transport_unsupported_reason(
transport,
spec_metadata.api_format,
provider_api_format,
)
};
if let Some(skip_reason) = transport_unsupported_reason {
mark_skipped_local_standard_candidate(
state,
input,
@@ -321,7 +460,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
provider_api_format,
parts.uri.path(),
upstream_is_stream,
if is_kiro_claude_cli {
if is_kiro_claude_cli || is_windsurf_cascade {
None
} else {
transport.endpoint.body_rules.as_ref()
@@ -379,6 +518,11 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
.await;
return None;
}
apply_non_native_claude_thinking_signature_compat(
&mut provider_request_body,
provider_api_format,
transport,
);
if let Some(mapping) =
crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model(
state,
@@ -422,6 +566,11 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
.await;
return None;
}
apply_non_native_claude_thinking_signature_compat(
&mut provider_request_body,
provider_api_format,
transport,
);
}
if let Some(kiro_auth) = kiro_auth.as_ref() {
@@ -443,6 +592,24 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
)
.await;
}
if is_windsurf_cascade {
return build_windsurf_cross_format_payload_parts(
state,
parts,
trace_id,
body_json,
input,
attempt,
transport,
provider_api_format,
prepared_candidate.mapped_model,
prepared_candidate.auth_header,
prepared_candidate.auth_value,
provider_request_body,
upstream_is_stream,
)
.await;
}
let upstream_url = match crate::ai_serving::planner::standard::build_standard_upstream_url(
parts,
@@ -542,6 +709,121 @@ fn apply_transport_request_body_semantics(
)
}
#[allow(clippy::too_many_arguments)]
async fn build_windsurf_cross_format_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
original_body_json: &serde_json::Value,
input: &LocalStandardDecisionInput,
attempt: &LocalStandardCandidateAttempt,
transport: &Arc<GatewayProviderTransportSnapshot>,
provider_api_format: &str,
mapped_model: String,
auth_header: String,
auth_value: String,
openai_chat_request_body: Value,
upstream_is_stream: bool,
) -> Option<LocalStandardCandidatePayloadParts> {
let candidate = &attempt.eligible.candidate;
let effective_headers = input.effective_headers(&parts.headers);
let provider_request_body = match build_windsurf_cascade_request_body(
&openai_chat_request_body,
&mapped_model,
&auth_value,
transport.endpoint.body_rules.as_ref(),
Some(effective_headers),
upstream_is_stream,
) {
Some(body) => body,
None => {
mark_skipped_local_standard_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
&openai_chat_request_body,
provider_api_format,
provider_api_format,
),
)
.await;
return None;
}
};
let upstream_url = match build_windsurf_cascade_upstream_url(
transport.endpoint.base_url.as_str(),
parts.uri.query(),
) {
Some(url) => url,
None => {
mark_skipped_local_standard_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"upstream_url_missing",
CandidateFailureDiagnostic::upstream_url_missing(
provider_api_format,
provider_api_format,
"standard_family_windsurf_url",
),
)
.await;
return None;
}
};
let provider_request_headers = match build_windsurf_cascade_headers(
effective_headers,
&provider_request_body,
original_body_json,
transport.endpoint.header_rules.as_ref(),
&auth_header,
&auth_value,
upstream_is_stream,
) {
Some(headers) => headers,
None => {
mark_skipped_local_standard_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
provider_api_format,
provider_api_format,
"standard_family_windsurf_headers",
),
)
.await;
return None;
}
};
Some(LocalStandardCandidatePayloadParts {
auth_header,
auth_value,
mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
upstream_is_stream,
envelope_name: Some(WINDSURF_ENVELOPE_NAME),
transport: Arc::clone(transport),
transport_profile: None,
})
}
async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
@@ -799,3 +1081,95 @@ async fn build_kiro_cross_format_payload_parts(
transport_profile: None,
})
}
#[cfg(test)]
mod tests {
use super::{
provider_preserves_claude_thinking_signatures,
sanitize_claude_request_thinking_signatures_for_non_native,
};
use serde_json::json;
#[test]
fn sanitizes_historical_claude_thinking_for_non_native_relays() {
let mut body = json!({
"model": "claude-opus-4-1",
"messages": [{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "I should keep this short.",
"signature": "sig_123"
},
{
"type": "redacted_thinking",
"data": "opaque"
},
{
"type": "text",
"text": "Done."
}
]
}]
});
assert!(sanitize_claude_request_thinking_signatures_for_non_native(
&mut body
));
assert_eq!(body["messages"][0]["content"][0]["type"], json!("text"));
assert_eq!(
body["messages"][0]["content"][0]["text"],
json!("I should keep this short.")
);
assert_eq!(body["messages"][0]["content"].as_array().unwrap().len(), 2);
assert_eq!(body["messages"][0]["content"][1]["text"], json!("Done."));
}
#[test]
fn inserts_placeholder_when_only_redacted_thinking_would_remain() {
let mut body = json!({
"model": "claude-opus-4-1",
"messages": [{
"role": "assistant",
"content": [{
"type": "redacted_thinking",
"data": "opaque"
}]
}]
});
assert!(sanitize_claude_request_thinking_signatures_for_non_native(
&mut body
));
assert_eq!(body["messages"][0]["content"][0]["type"], json!("text"));
assert_eq!(
body["messages"][0]["content"][0]["text"],
json!("Previous thinking omitted.")
);
}
#[test]
fn official_claude_providers_preserve_thinking_signatures() {
assert!(provider_preserves_claude_thinking_signatures(
"anthropic",
"https://relay.example.com"
));
assert!(provider_preserves_claude_thinking_signatures(
"custom",
"https://api.anthropic.com"
));
assert!(provider_preserves_claude_thinking_signatures(
"aws",
"https://bedrock-runtime.us-east-1.amazonaws.com"
));
assert!(provider_preserves_claude_thinking_signatures(
"amazon_bedrock",
"https://relay.example.com"
));
assert!(!provider_preserves_claude_thinking_signatures(
"openai",
"https://relay.example.com"
));
}
}
@@ -210,7 +210,7 @@ fn local_openai_responses_compact_wrapper_strips_include_for_codex_requests() {
assert_eq!(provider_request_body["instructions"], "");
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"3d2e2842-74cb-55dd-803a-b8940b3500c2"
);
}
@@ -356,7 +356,7 @@ fn injects_codex_prompt_cache_key_for_openai_responses_cross_format_requests() {
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"b4dfeb75-b105-544c-a706-39b92f0bddb0"
);
}
@@ -386,6 +386,6 @@ fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() {
assert_eq!(
provider_request_body["prompt_cache_key"],
"172c39e6-c0a0-5a70-8b63-e0f8e0d185a3"
"4ee6ea6e-3ac6-5a18-8cb8-1f8b956419e5"
);
}
@@ -4,7 +4,7 @@ use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_contracts::ResolvedTransportProfile;
use serde_json::Value;
use serde_json::{json, Value};
use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
@@ -16,9 +16,10 @@ use crate::ai_serving::planner::common::{
request_requires_body_stream_field, OPENAI_CHAT_STREAM_PLAN_KIND,
};
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
build_local_openai_chat_upstream_url, request_body_build_failure_extra_data,
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_upstream_url,
build_local_openai_chat_request_body, build_local_openai_chat_upstream_url,
request_body_build_failure_extra_data,
};
use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth;
use crate::ai_serving::transport::kiro::{
@@ -27,10 +28,18 @@ use crate::ai_serving::transport::kiro::{
KIRO_ENVELOPE_NAME,
};
use crate::ai_serving::transport::local_openai_chat_transport_unsupported_reason;
use crate::ai_serving::transport::windsurf::{
build_windsurf_cascade_headers, build_windsurf_cascade_request_body,
build_windsurf_cascade_upstream_url, is_windsurf_provider_transport,
local_windsurf_request_transport_unsupported_reason_with_network,
resolve_windsurf_cascade_auth, WINDSURF_ENVELOPE_NAME,
};
use crate::ai_serving::transport::{
build_grok_browser_headers, build_grok_upstream_url, build_kiro_cross_format_upstream_url,
build_standard_provider_request_headers, GrokHeaderInput, StandardProviderRequestHeadersInput,
GROK_CHAT_PATH,
build_openai_image_headers, build_openai_image_upstream_url,
build_standard_provider_request_headers, openai_image_transport_unsupported_reason,
resolve_openai_image_auth, GrokHeaderInput, ProviderOpenAiImageHeadersInput,
StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
};
use crate::ai_serving::{
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
@@ -331,6 +340,25 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
}));
}
if provider_api_format == "openai:chat" && is_windsurf_provider_transport(transport) {
return build_windsurf_openai_chat_payload_parts(
state,
parts,
trace_id,
body_json,
input,
eligible,
candidate_index,
candidate_id,
decision_kind,
report_kind,
transport,
upstream_is_stream,
redaction.redacted,
)
.await;
}
if provider_api_format == "openai:chat" {
if let Some(skip_reason) = local_openai_chat_transport_unsupported_reason(transport) {
mark_skipped_local_openai_chat_candidate(
@@ -498,6 +526,20 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
};
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
if provider_api_format == "openai:image" {
return resolve_openai_chat_to_openai_image_payload_parts(
state,
parts,
trace_id,
body_json,
input,
eligible,
candidate_index,
candidate_id,
upstream_is_stream,
)
.await;
}
let Some(conversion_kind) =
request_conversion_kind("openai:chat", provider_api_format.as_str())
@@ -789,6 +831,628 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
}))
}
#[allow(clippy::too_many_arguments)]
async fn resolve_openai_chat_to_openai_image_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
upstream_is_stream: bool,
) -> Result<Option<LocalOpenAiChatCandidatePayloadParts>, GatewayError> {
let candidate = &eligible.candidate;
let transport = &eligible.transport;
let provider_api_format = "openai:image";
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, provider_api_format)
{
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return Ok(None);
}
let prepared_candidate = match prepare_header_authenticated_candidate(
crate::ai_serving::PlannerAppState::new(state),
transport,
candidate,
resolve_openai_image_auth(transport),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "openai_chat_image_bridge",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return Ok(None);
}
};
let is_chatgpt_web = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web");
let Some((mut provider_request_body, image_request_summary)) = (if is_chatgpt_web {
build_chatgpt_web_image_provider_body_from_openai_chat_body(
body_json,
&input.requested_model,
)
} else {
build_openai_image_provider_body_from_openai_chat_body(
body_json,
&input.requested_model,
upstream_is_stream,
)
}) else {
mark_skipped_local_openai_chat_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(body_json, "openai:chat", provider_api_format),
)
.await;
return Ok(None);
};
if !is_chatgpt_web {
apply_codex_openai_responses_special_body_edits(
&mut provider_request_body,
transport.provider.provider_type.as_str(),
provider_api_format,
transport.endpoint.body_rules.as_ref(),
Some(candidate.key_id.as_str()),
);
}
let upstream_url = if is_chatgpt_web {
chatgpt_web_image_internal_url(&transport.endpoint.base_url)
} else {
build_openai_image_upstream_url(
transport,
Some("/v1/images/generations"),
parts.uri.query(),
)
};
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
})
else {
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
"openai:chat",
provider_api_format,
"openai_chat_image_bridge_headers",
),
)
.await;
return Ok(None);
};
if is_chatgpt_web {
provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else {
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
}
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", provider_api_format);
Ok(Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
execution_strategy,
conversion_mode,
report_kind: "openai_chat_stream_success".to_string(),
envelope_name: None,
transport: Arc::clone(transport),
request_redacted: false,
transport_profile: None,
image_request_summary: Some(image_request_summary),
}))
}
fn build_openai_image_provider_body_from_openai_chat_body(
body_json: &Value,
requested_model: &str,
upstream_is_stream: bool,
) -> Option<(Value, Value)> {
let (prompt, images) = collect_openai_chat_image_prompt_and_images(body_json)?;
let operation = if images.is_empty() {
"generate"
} else {
"edit"
};
let mut image_options = serde_json::Map::new();
copy_openai_chat_image_option(body_json, &mut image_options, "size");
copy_openai_chat_image_option(body_json, &mut image_options, "quality");
copy_openai_chat_image_option(body_json, &mut image_options, "background");
copy_openai_chat_image_option(body_json, &mut image_options, "output_format");
copy_openai_chat_image_option(body_json, &mut image_options, "output_compression");
copy_openai_chat_image_option(body_json, &mut image_options, "moderation");
copy_openai_chat_image_option(body_json, &mut image_options, "input_fidelity");
copy_openai_chat_image_option(body_json, &mut image_options, "partial_images");
let input = if images.is_empty() {
serde_json::json!([{
"role": "user",
"content": prompt,
}])
} else {
let mut content = vec![serde_json::json!({
"type": "input_text",
"text": prompt,
})];
content.extend(images);
serde_json::json!([{
"role": "user",
"content": content,
}])
};
let mut body = serde_json::Map::new();
if let Some(model) = body_json
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
let requested_model = requested_model.trim();
(!requested_model.is_empty()).then_some(requested_model)
})
{
body.insert("model".to_string(), Value::String(model.to_string()));
}
body.insert("input".to_string(), input);
let mut image_tool = image_options.clone();
image_tool.insert(
"type".to_string(),
Value::String("image_generation".to_string()),
);
body.insert(
"tools".to_string(),
Value::Array(vec![Value::Object(image_tool)]),
);
if upstream_is_stream {
body.insert("stream".to_string(), Value::Bool(true));
}
if let Some(user) = body_json
.get("user")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
body.insert("user".to_string(), Value::String(user.to_string()));
}
let mut summary = serde_json::Map::new();
summary.insert(
"operation".to_string(),
Value::String(operation.to_string()),
);
for key in ["output_format", "partial_images", "size", "quality"] {
if let Some(value) = image_options.get(key) {
summary.insert(key.to_string(), value.clone());
}
}
Some((Value::Object(body), Value::Object(summary)))
}
fn build_chatgpt_web_image_provider_body_from_openai_chat_body(
body_json: &Value,
requested_model: &str,
) -> Option<(Value, Value)> {
let (prompt, images) = collect_openai_chat_image_prompt_and_images(body_json)?;
let operation = if images.is_empty() {
"generate"
} else {
"edit"
};
let size = body_json
.get("size")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("1024x1024");
let output_format = body_json
.get("output_format")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("png");
let quality = body_json
.get("quality")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("medium");
let model = body_json
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or_else(|| requested_model.trim());
let web_model = body_json
.get("web_model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("gpt-5-5-thinking");
let image_urls = openai_image_inputs_as_urls(&images);
let body = json!({
"operation": operation,
"model": if model.is_empty() { "gpt-image-2" } else { model },
"web_model": web_model,
"prompt": prompt,
"size": size,
"ratio": chatgpt_web_ratio_for_size(size),
"output_format": output_format,
"images": image_urls,
});
let summary = json!({
"operation": operation,
"output_format": output_format,
"size": size,
"quality": quality,
});
Some((body, summary))
}
fn copy_openai_chat_image_option(
body_json: &Value,
image_options: &mut serde_json::Map<String, Value>,
key: &str,
) {
if let Some(value) = body_json.get(key) {
image_options.insert(key.to_string(), value.clone());
}
}
fn collect_openai_chat_image_prompt_and_images(body_json: &Value) -> Option<(String, Vec<Value>)> {
let messages = body_json.get("messages").and_then(Value::as_array)?;
let mut prompt_parts = Vec::new();
let mut images = Vec::new();
for message in messages.iter().filter_map(Value::as_object) {
let role = message
.get("role")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
let content = message.get("content");
if matches!(role, "system" | "developer" | "user") {
if let Some(text) = crate::ai_serving::extract_openai_text_content(content)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
{
prompt_parts.push(text);
}
}
if role == "user" {
collect_openai_chat_image_inputs(content, &mut images);
}
}
let prompt = prompt_parts.join("\n").trim().to_string();
(!prompt.is_empty()).then_some((prompt, images))
}
fn collect_openai_chat_image_inputs(content: Option<&Value>, images: &mut Vec<Value>) {
let Some(parts) = content.and_then(Value::as_array) else {
return;
};
for part in parts.iter().filter_map(Value::as_object) {
let part_type = part
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if matches!(part_type, "image_url" | "input_image") {
if let Some(url) = part
.get("image_url")
.and_then(|value| {
value
.as_str()
.or_else(|| value.get("url").and_then(Value::as_str))
})
.map(str::trim)
.filter(|value| !value.is_empty())
{
images.push(serde_json::json!({
"type": "input_image",
"image_url": url,
}));
} else if let Some(file_id) = part
.get("file_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
images.push(serde_json::json!({
"type": "input_image",
"file_id": file_id,
}));
}
}
}
}
fn openai_image_inputs_as_urls(images: &[Value]) -> Vec<Value> {
images
.iter()
.filter_map(|image| {
image
.get("image_url")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| Value::String(value.to_string()))
})
.collect()
}
fn chatgpt_web_ratio_for_size(size: &str) -> String {
let Some((width, height)) = size.split_once('x') else {
return "1:1".to_string();
};
let Ok(width) = width.trim().parse::<u64>() else {
return "1:1".to_string();
};
let Ok(height) = height.trim().parse::<u64>() else {
return "1:1".to_string();
};
if width == 0 || height == 0 {
return "1:1".to_string();
}
let divisor = gcd(width, height);
format!("{}:{}", width / divisor, height / divisor)
}
fn gcd(mut left: u64, mut right: u64) -> u64 {
while right != 0 {
let next = left % right;
left = right;
right = next;
}
left.max(1)
}
fn chatgpt_web_image_internal_url(base_url: &str) -> String {
let base_url = base_url.trim().trim_end_matches('/');
let base_url = if base_url.is_empty() {
"https://chatgpt.com"
} else {
base_url
};
format!("{base_url}/__aether/chatgpt-web-image")
}
#[allow(clippy::too_many_arguments)]
async fn build_windsurf_openai_chat_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
decision_kind: &str,
report_kind: &str,
transport: &Arc<GatewayProviderTransportSnapshot>,
upstream_is_stream: bool,
request_redacted: bool,
) -> Result<Option<LocalOpenAiChatCandidatePayloadParts>, GatewayError> {
let planner_state = crate::ai_serving::PlannerAppState::new(state);
let candidate = &eligible.candidate;
if let Some(skip_reason) =
local_windsurf_request_transport_unsupported_reason_with_network(transport)
{
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return Ok(None);
}
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
transport,
candidate,
resolve_windsurf_cascade_auth(transport)
.or_else(|| resolve_local_openai_bearer_auth(transport)),
OauthPreparationContext {
trace_id,
api_format: "openai:chat",
operation: "openai_chat_windsurf_cascade",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return Ok(None);
}
};
let Some(provider_request_body) = build_windsurf_cascade_request_body(
body_json,
&prepared_candidate.mapped_model,
&prepared_candidate.auth_value,
transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
upstream_is_stream,
) else {
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
CandidateFailureDiagnostic::envelope_build_failed(
"openai:chat",
"openai:chat",
"openai_chat_windsurf_cascade",
),
)
.await;
return Ok(None);
};
let Some(upstream_url) = build_windsurf_cascade_upstream_url(
transport.endpoint.base_url.as_str(),
parts.uri.query(),
) else {
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"upstream_url_missing",
CandidateFailureDiagnostic::upstream_url_missing(
"openai:chat",
"openai:chat",
"openai_chat_windsurf_url",
),
)
.await;
return Ok(None);
};
let mut provider_request_headers = match build_windsurf_cascade_headers(
&parts.headers,
&provider_request_body,
body_json,
transport.endpoint.header_rules.as_ref(),
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
upstream_is_stream,
) {
Some(headers) => headers,
None => {
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
"openai:chat",
"openai:chat",
"openai_chat_windsurf_headers",
),
)
.await;
return Ok(None);
}
};
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
request_redacted,
);
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", "openai:chat");
let resolved_report_kind =
if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND || !upstream_is_stream {
report_kind.to_string()
} else {
"openai_chat_sync_finalize".to_string()
};
Ok(Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format: "openai:chat".to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
execution_strategy,
conversion_mode,
report_kind: resolved_report_kind,
envelope_name: Some(WINDSURF_ENVELOPE_NAME),
transport: Arc::clone(transport),
request_redacted,
transport_profile: None,
image_request_summary: None,
}))
}
#[allow(clippy::too_many_arguments)]
async fn build_kiro_openai_chat_cross_format_payload_parts(
state: &AppState,
@@ -1012,3 +1676,71 @@ fn redaction_mask_error_to_gateway_error(error: RedactionMaskError) -> GatewayEr
},
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn chatgpt_web_chat_image_bridge_body_uses_internal_web_shape() {
let body_json = json!({
"model": "gpt-image-2",
"messages": [
{"role": "system", "content": "Use crisp vector-like shapes."},
{
"role": "user",
"content": [
{"type": "text", "text": "Draw a glass city"},
{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}
]
}
],
"size": "1536x1024",
"output_format": "webp",
"web_model": "gpt-5-image-test"
});
let (provider_body, summary) =
build_chatgpt_web_image_provider_body_from_openai_chat_body(&body_json, "gpt-image-2")
.expect("chat image body should convert");
assert_eq!(provider_body["operation"], "edit");
assert_eq!(provider_body["model"], "gpt-image-2");
assert_eq!(provider_body["web_model"], "gpt-5-image-test");
assert_eq!(
provider_body["prompt"],
"Use crisp vector-like shapes.\nDraw a glass city"
);
assert_eq!(provider_body["size"], "1536x1024");
assert_eq!(provider_body["ratio"], "3:2");
assert_eq!(provider_body["output_format"], "webp");
assert_eq!(provider_body["images"][0], "https://example.com/ref.png");
assert_eq!(summary["operation"], "edit");
assert_eq!(summary["output_format"], "webp");
}
#[test]
fn openai_chat_image_bridge_body_injects_image_generation_tool() {
let body_json = json!({
"model": "gpt-image-2",
"messages": [
{"role": "user", "content": "Draw a glass city"}
],
"size": "1024x1024",
"output_format": "png"
});
let (provider_body, summary) =
build_openai_image_provider_body_from_openai_chat_body(&body_json, "gpt-image-2", true)
.expect("chat image body should convert");
assert_eq!(provider_body["tools"][0]["type"], "image_generation");
assert_eq!(provider_body["tools"][0]["size"], "1024x1024");
assert_eq!(provider_body["tools"][0]["output_format"], "png");
assert_eq!(provider_body["model"], "gpt-image-2");
assert_eq!(provider_body["stream"], true);
assert_eq!(provider_body["input"][0]["content"], "Draw a glass city");
assert_eq!(summary["operation"], "generate");
assert_eq!(summary["output_format"], "png");
}
}
@@ -39,10 +39,13 @@ use crate::ai_serving::transport::kiro::{
use crate::ai_serving::transport::{
build_grok_browser_headers, build_grok_upstream_url, build_kiro_cross_format_upstream_url,
build_openai_image_headers, build_openai_image_upstream_url,
build_standard_provider_request_headers,
local_standard_transport_unsupported_reason_with_network,
build_standard_provider_request_headers, build_windsurf_cascade_headers,
build_windsurf_cascade_request_body, build_windsurf_cascade_upstream_url,
is_windsurf_provider_transport, local_standard_transport_unsupported_reason_with_network,
local_windsurf_request_transport_unsupported_reason_with_network,
openai_image_transport_unsupported_reason, resolve_openai_image_auth, GrokHeaderInput,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
WINDSURF_ENVELOPE_NAME,
};
use crate::ai_serving::{
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
@@ -128,6 +131,8 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
)
.await;
}
let is_windsurf_cascade =
provider_api_format == "openai:chat" && is_windsurf_provider_transport(transport);
let same_format = api_format_alias_matches(provider_api_format, &client_api_format);
let conversion_kind = request_conversion_kind(spec_metadata.api_format, provider_api_format);
@@ -139,6 +144,8 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
local_kiro_request_transport_unsupported_reason_with_network(transport)
} else if same_format {
local_standard_transport_unsupported_reason_with_network(transport, provider_api_format)
} else if is_windsurf_cascade {
local_windsurf_request_transport_unsupported_reason_with_network(transport)
} else {
match conversion_kind {
Some(_) if is_antigravity && provider_api_format == "gemini:generate_content" => None,
@@ -302,7 +309,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
if is_kiro_claude_cli {
if is_kiro_claude_cli || is_windsurf_cascade {
None
} else {
transport.endpoint.body_rules.as_ref()
@@ -319,7 +326,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
force_body_stream_field,
transport.provider.provider_type.as_str(),
provider_api_format,
if is_kiro_claude_cli {
if is_kiro_claude_cli || is_windsurf_cascade {
None
} else {
transport.endpoint.body_rules.as_ref()
@@ -447,6 +454,27 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
)
.await;
}
if is_windsurf_cascade {
return build_windsurf_openai_responses_payload_parts(
state,
parts,
trace_id,
body_json,
input,
eligible,
candidate_index,
candidate_id,
spec_metadata.api_format,
transport,
provider_api_format,
mapped_model,
auth_header,
auth_value,
provider_request_body,
upstream_is_stream,
)
.await;
}
let Some(upstream_url) = (if is_grok && is_grok_text_provider_api_format(provider_api_format) {
Some(build_grok_upstream_url(transport, GROK_CHAT_PATH))
@@ -619,6 +647,130 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
})
}
#[allow(clippy::too_many_arguments)]
async fn build_windsurf_openai_responses_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
original_body_json: &serde_json::Value,
input: &LocalOpenAiResponsesDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
client_api_format: &str,
transport: &Arc<GatewayProviderTransportSnapshot>,
provider_api_format: &str,
mapped_model: String,
auth_header: String,
auth_value: String,
openai_chat_request_body: Value,
upstream_is_stream: bool,
) -> Option<LocalOpenAiResponsesCandidatePayloadParts> {
let candidate = &eligible.candidate;
let effective_headers = input.effective_headers(&parts.headers);
let provider_request_body = match build_windsurf_cascade_request_body(
&openai_chat_request_body,
&mapped_model,
&auth_value,
transport.endpoint.body_rules.as_ref(),
Some(effective_headers),
upstream_is_stream,
) {
Some(body) => body,
None => {
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
CandidateFailureDiagnostic::envelope_build_failed(
client_api_format,
provider_api_format,
"openai_responses_windsurf_cascade",
),
)
.await;
return None;
}
};
let upstream_url = match build_windsurf_cascade_upstream_url(
transport.endpoint.base_url.as_str(),
parts.uri.query(),
) {
Some(url) => url,
None => {
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"upstream_url_missing",
CandidateFailureDiagnostic::upstream_url_missing(
client_api_format,
provider_api_format,
"openai_responses_windsurf_url",
),
)
.await;
return None;
}
};
let provider_request_headers = match build_windsurf_cascade_headers(
effective_headers,
&provider_request_body,
original_body_json,
transport.endpoint.header_rules.as_ref(),
&auth_header,
&auth_value,
upstream_is_stream,
) {
Some(headers) => headers,
None => {
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
client_api_format,
provider_api_format,
"openai_responses_windsurf_headers",
),
)
.await;
return None;
}
};
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(client_api_format, provider_api_format);
Some(LocalOpenAiResponsesCandidatePayloadParts {
auth_header,
auth_value,
mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
execution_strategy,
conversion_mode,
is_antigravity: false,
envelope_name: Some(WINDSURF_ENVELOPE_NAME),
upstream_is_stream,
transport: Arc::clone(transport),
transport_profile: None,
image_request_summary: None,
})
}
fn api_format_alias_matches(left: &str, right: &str) -> bool {
crate::ai_serving::api_format_alias_matches(left, right)
}
@@ -54,6 +54,10 @@ pub(crate) mod vertex {
pub(crate) use aether_provider_transport::vertex::*;
}
pub(crate) mod windsurf {
pub(crate) use aether_provider_transport::windsurf::*;
}
pub(crate) use aether_provider_transport::{
append_transport_diagnostics_to_value, apply_local_body_rules,
apply_local_body_rules_with_request_headers, apply_local_header_rules,
@@ -73,13 +77,16 @@ pub(crate) use aether_provider_transport::{
build_standard_plan_fallback_openai_responses_url, build_standard_provider_request_headers,
build_transport_request_url, build_transport_request_url_for_request_body,
build_video_create_headers, build_video_create_request_body, build_video_create_upstream_url,
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
classify_same_format_provider_request_behavior, ensure_upstream_auth_header,
gemini_files_transport_unsupported_reason, header_rules_are_locally_supported,
header_rules_have_enabled_rules, is_gemini_cli_provider_transport,
build_windsurf_cascade_headers, build_windsurf_cascade_request_body,
build_windsurf_cascade_upstream_url, candidate_common_transport_skip_reason,
candidate_transport_pair_skip_reason, classify_same_format_provider_request_behavior,
ensure_upstream_auth_header, gemini_files_transport_unsupported_reason,
header_rules_are_locally_supported, header_rules_have_enabled_rules,
is_gemini_cli_provider_transport, is_windsurf_provider_transport,
local_gemini_transport_unsupported_reason_with_network,
local_openai_chat_transport_unsupported_reason,
local_standard_transport_unsupported_reason_with_network,
local_windsurf_request_transport_unsupported_reason_with_network,
openai_image_transport_unsupported_reason, request_conversion_direct_auth,
request_conversion_enabled_for_transport, request_conversion_transport_supported,
request_conversion_transport_unsupported_reason, request_pair_allowed_for_transport,
@@ -103,5 +110,5 @@ pub(crate) use aether_provider_transport::{
StandardProviderRequestHeaders, StandardProviderRequestHeadersInput,
TransportRequestBodySemanticsError, TransportRequestUrlParams, GEMINI_CLI_USER_AGENT,
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME, GROK_CHAT_PATH, GROK_INTERNAL_HEADER,
GROK_RATE_LIMITS_PATH,
GROK_RATE_LIMITS_PATH, WINDSURF_ENVELOPE_NAME,
};
+13 -1
View File
@@ -8,7 +8,7 @@ use aether_data_contracts::repository::video_tasks::{
use aether_usage_runtime::{build_upsert_usage_record_from_event, settle_usage_if_needed};
use serde_json::{Map, Value};
use tokio::task::JoinHandle;
use tracing::{info, warn};
use tracing::{debug, info, warn};
use crate::log_ids::short_request_id;
use crate::usage::{UsageEvent, UsageEventData, UsageEventType};
@@ -148,8 +148,20 @@ pub(crate) fn spawn_video_task_poller(state: AppState) -> Option<JoinHandle<()>>
let mut interval = tokio::time::interval(config.interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut deferred_since = None;
loop {
interval.tick().await;
if state
.data
.should_defer_maintenance_for_database_pool_pressure(&mut deferred_since)
{
debug!(
event_name = "video_task_poller_deferred",
log_type = "event",
"gateway video task poller deferred because database pool has no idle reserve"
);
continue;
}
if let Err(err) = poll_video_tasks_once(&state, config.batch_size).await {
warn!(
event_name = "video_task_poller_tick_failed",
+173
View File
@@ -0,0 +1,173 @@
use crate::handlers::shared::{
decrypt_catalog_secret_with_fallbacks, system_config_bool, system_config_string,
};
use crate::{AppState, GatewayError};
use serde_json::{json, Value};
pub(crate) const BARK_PUSH_ENABLED_KEY: &str = "module.bark_push.enabled";
pub(crate) const BARK_PUSH_DEVICE_KEY_KEY: &str = "module.bark_push.device_key";
pub(crate) const BARK_PUSH_SERVER_URL_KEY: &str = "module.bark_push.server_url";
pub(crate) const BARK_PUSH_TEMPLATE_KEY: &str = "module.bark_push.template";
const DEFAULT_BARK_API_BASE: &str = "https://api.day.app";
#[derive(Debug, Clone)]
pub(crate) struct BarkPushConfig {
pub(crate) enabled: bool,
pub(crate) device_key: Option<String>,
pub(crate) server_url: String,
pub(crate) template: Option<String>,
}
pub(crate) async fn bark_push_module_enabled(state: &AppState) -> Result<bool, GatewayError> {
let value = state
.read_system_config_json_value(BARK_PUSH_ENABLED_KEY)
.await?;
Ok(system_config_bool(value.as_ref(), false))
}
pub(crate) async fn bark_push_configured(state: &AppState) -> Result<bool, GatewayError> {
let config = read_bark_push_config(state).await?;
Ok(config.device_key.is_some() && !config.server_url.trim().is_empty())
}
pub(crate) async fn read_bark_push_config(
state: &AppState,
) -> Result<BarkPushConfig, GatewayError> {
let enabled = bark_push_module_enabled(state).await?;
let device_key = state
.read_system_config_json_value(BARK_PUSH_DEVICE_KEY_KEY)
.await?
.and_then(|value| system_config_string(Some(&value)))
.map(|value| {
decrypt_catalog_secret_with_fallbacks(state.encryption_key(), &value).unwrap_or(value)
});
let server_url = state
.read_system_config_json_value(BARK_PUSH_SERVER_URL_KEY)
.await?
.and_then(|value| system_config_string(Some(&value)))
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| DEFAULT_BARK_API_BASE.to_string());
let template = state
.read_system_config_json_value(BARK_PUSH_TEMPLATE_KEY)
.await?
.and_then(|value| system_config_string(Some(&value)));
Ok(BarkPushConfig {
enabled,
device_key,
server_url,
template,
})
}
pub(crate) async fn send_bark_push(
state: &AppState,
config: &BarkPushConfig,
title: &str,
markdown_body: &str,
) -> Result<(), GatewayError> {
let Some(device_key) = config.device_key.as_deref() else {
return Err(GatewayError::Internal("未配置 Bark Device Key".to_string()));
};
let device_key = device_key.trim();
if device_key.is_empty() {
return Err(GatewayError::Internal(
"Bark Device Key 不能为空".to_string(),
));
}
let server_url = normalized_bark_server_url(&config.server_url)?;
let body = render_bark_body(config.template.as_deref(), title, markdown_body);
let response = state
.client
.post(format!("{server_url}/push"))
.json(&json!({
"device_key": device_key,
"title": title,
"body": body,
}))
.send()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if !status.is_success() {
return Err(GatewayError::Internal(format!(
"Bark 返回 HTTP {status}: {text}"
)));
}
if let Ok(payload) = serde_json::from_str::<Value>(&text) {
let code_is_ok = payload
.get("code")
.and_then(|value| {
value
.as_i64()
.map(|code| matches!(code, 0 | 200))
.or_else(|| {
value
.as_str()
.map(|code| matches!(code.trim(), "0" | "200"))
})
})
.unwrap_or(true);
if !code_is_ok {
return Err(GatewayError::Internal(format!("Bark 返回失败: {payload}")));
}
}
Ok(())
}
fn normalized_bark_server_url(server_url: &str) -> Result<String, GatewayError> {
let server_url = server_url.trim().trim_end_matches('/');
if server_url.is_empty() {
return Err(GatewayError::Internal(
"Bark 服务器地址不能为空".to_string(),
));
}
if !server_url.starts_with("https://") && !server_url.starts_with("http://") {
return Err(GatewayError::Internal(
"Bark 服务器地址必须以 http:// 或 https:// 开头".to_string(),
));
}
Ok(server_url.to_string())
}
fn render_bark_body(template: Option<&str>, title: &str, markdown_body: &str) -> String {
match template {
Some(template) if !template.trim().is_empty() => template
.replace("{title}", title)
.replace("{body}", markdown_body),
_ => markdown_body.to_string(),
}
}
#[cfg(test)]
mod tests {
use super::{normalized_bark_server_url, render_bark_body};
#[test]
fn bark_body_uses_template_when_provided() {
let rendered = render_bark_body(Some("{title}\n\n{body}"), "告警", "原始正文");
assert_eq!(rendered, "告警\n\n原始正文");
}
#[test]
fn bark_body_falls_back_to_markdown_body_for_empty_template() {
assert_eq!(render_bark_body(None, "告警", "原始正文"), "原始正文");
assert_eq!(
render_bark_body(Some(" "), "告警", "原始正文"),
"原始正文"
);
}
#[test]
fn bark_server_url_trims_trailing_slashes() {
assert_eq!(
normalized_bark_server_url(" https://api.day.app/ ").expect("url should parse"),
"https://api.day.app"
);
}
}
@@ -103,6 +103,16 @@ pub(super) fn classify_admin_system_family_route(
"admin:system",
false,
))
} else if method == http::Method::POST
&& normalized_path == "/api/admin/system/important-notification/test"
{
Some(classified(
"admin_proxy",
"system_manage",
"important_notification_test",
"admin:system",
false,
))
} else if method == http::Method::POST && normalized_path == "/api/admin/system/cleanup" {
Some(classified(
"admin_proxy",
@@ -183,6 +183,10 @@ fn classifies_admin_system_maintenance_write_routes_as_admin_proxy_route() {
("/api/admin/system/users/import", "users_import"),
("/api/admin/system/data/import", "data_import"),
("/api/admin/system/smtp/test", "smtp_test"),
(
"/api/admin/system/important-notification/test",
"important_notification_test",
),
("/api/admin/system/cleanup", "cleanup"),
("/api/admin/system/purge/config", "purge_config"),
("/api/admin/system/purge/users", "purge_users"),
+40 -15
View File
@@ -2000,7 +2000,7 @@ fn resolve_effective_rate_limit_policy(
groups: &[aether_data::repository::users::StoredUserGroup],
) -> Option<i32> {
let group_policy = groups.iter().fold(None, |effective, group| {
intersect_rate_limit_policies(
union_rate_limit_policies(
effective,
rate_limit_restriction_from_mode(&group.rate_limit_mode, group.rate_limit),
)
@@ -2073,6 +2073,22 @@ fn intersect_rate_limit_policies(
}
}
fn union_rate_limit_policies(
left: Option<RateLimitRestriction>,
right: Option<RateLimitRestriction>,
) -> Option<RateLimitRestriction> {
match (left, right) {
(None, None) => None,
(Some(value), None) | (None, Some(value)) => Some(value),
(Some(RateLimitRestriction::Unlimited), _) | (_, Some(RateLimitRestriction::Unlimited)) => {
Some(RateLimitRestriction::Unlimited)
}
(Some(RateLimitRestriction::Limited(left)), Some(RateLimitRestriction::Limited(right))) => {
Some(RateLimitRestriction::Limited(left.max(right)))
}
}
}
fn rate_limit_policy_value(policy: Option<RateLimitRestriction>) -> Option<i32> {
match policy {
None => None,
@@ -2358,35 +2374,44 @@ mod tests {
}
#[test]
fn rate_limit_policy_uses_most_restrictive_custom_limit() {
let groups = vec![sample_group(
"restricted",
10,
None,
"unrestricted",
Some(60),
"custom",
)];
fn rate_limit_policy_uses_highest_group_limit_before_user_restriction() {
let groups = vec![
sample_group("default", 10, None, "unrestricted", Some(30), "custom"),
sample_group("tier-1", 20, None, "unrestricted", Some(100), "custom"),
];
assert_eq!(
resolve_effective_rate_limit_policy(Some(120), "custom", &groups),
Some(60)
Some(100)
);
}
#[test]
fn rate_limit_unlimited_does_not_bypass_limited_group() {
fn rate_limit_unlimited_group_overrides_limited_groups() {
let groups = vec![
sample_group("default", 10, None, "unrestricted", Some(30), "custom"),
sample_group("tier-2", 20, None, "unrestricted", Some(0), "custom"),
];
assert_eq!(
resolve_effective_rate_limit_policy(None, "system", &groups),
Some(0)
);
}
#[test]
fn rate_limit_user_policy_still_restricts_group_grants() {
let groups = vec![sample_group(
"restricted",
"tier-1",
10,
None,
"unrestricted",
Some(60),
Some(100),
"custom",
)];
assert_eq!(
resolve_effective_rate_limit_policy(Some(0), "custom", &groups),
resolve_effective_rate_limit_policy(Some(60), "custom", &groups),
Some(60)
);
}
@@ -42,8 +42,12 @@ use aether_data_contracts::repository::usage::{
};
use aether_runtime_state::RuntimeQueueStore;
use aether_video_tasks_core::read_data_backed_video_task_response;
use std::time::{Duration, Instant};
impl GatewayDataState {
const MAINTENANCE_POOL_IDLE_RESERVE: usize = 1;
const MAINTENANCE_POOL_PRESSURE_MAX_DEFER: Duration = Duration::from_secs(30);
pub(crate) async fn run_database_maintenance(
&self,
table_names: &[&str],
@@ -112,6 +116,47 @@ impl GatewayDataState {
.and_then(|backends| backends.database_pool_summary())
}
pub(crate) fn database_pool_under_maintenance_pressure(&self) -> bool {
self.database_pool_summary()
.as_ref()
.is_some_and(Self::database_pool_summary_under_maintenance_pressure)
}
pub(crate) fn database_pool_summary_under_maintenance_pressure(
summary: &aether_data::DatabasePoolSummary,
) -> bool {
summary.checked_out > 0 && summary.idle <= Self::MAINTENANCE_POOL_IDLE_RESERVE
}
pub(crate) fn should_defer_maintenance_for_database_pool_pressure(
&self,
deferred_since: &mut Option<Instant>,
) -> bool {
Self::should_defer_maintenance_for_pool_pressure_state(
self.database_pool_under_maintenance_pressure(),
deferred_since,
)
}
pub(crate) fn should_defer_maintenance_for_pool_pressure_state(
pool_under_pressure: bool,
deferred_since: &mut Option<Instant>,
) -> bool {
if !pool_under_pressure {
*deferred_since = None;
return false;
}
let now = Instant::now();
let since = deferred_since.get_or_insert(now);
if now.duration_since(*since) >= Self::MAINTENANCE_POOL_PRESSURE_MAX_DEFER {
*deferred_since = None;
return false;
}
true
}
pub(crate) async fn aggregate_wallet_daily_usage(
&self,
input: &WalletDailyUsageAggregationInput,
+60
View File
@@ -1,4 +1,5 @@
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
@@ -51,6 +52,65 @@ fn disabled_gateway_data_state_has_no_backends() {
assert!(!state.has_video_task_reader());
}
#[test]
fn maintenance_pool_pressure_keeps_idle_reserve_for_foreground_work() {
let pressured = aether_data::DatabasePoolSummary {
driver: DatabaseDriver::Postgres,
checked_out: 6,
pool_size: 6,
idle: 0,
max_connections: 20,
usage_rate: 30.0,
};
assert!(GatewayDataState::database_pool_summary_under_maintenance_pressure(&pressured));
let one_idle_left = aether_data::DatabasePoolSummary {
driver: DatabaseDriver::Postgres,
checked_out: 5,
pool_size: 6,
idle: 1,
max_connections: 20,
usage_rate: 25.0,
};
assert!(GatewayDataState::database_pool_summary_under_maintenance_pressure(&one_idle_left));
let idle = aether_data::DatabasePoolSummary {
driver: DatabaseDriver::Postgres,
checked_out: 0,
pool_size: 4,
idle: 4,
max_connections: 20,
usage_rate: 0.0,
};
assert!(!GatewayDataState::database_pool_summary_under_maintenance_pressure(&idle));
}
#[test]
fn maintenance_pool_pressure_deferral_has_timeout() {
let mut deferred_since = None;
assert!(
GatewayDataState::should_defer_maintenance_for_pool_pressure_state(
true,
&mut deferred_since
)
);
assert!(deferred_since.is_some());
assert!(
!GatewayDataState::should_defer_maintenance_for_pool_pressure_state(
false,
&mut deferred_since
)
);
assert!(deferred_since.is_none());
let mut stale_defer = Some(Instant::now() - Duration::from_secs(31));
assert!(
!GatewayDataState::should_defer_maintenance_for_pool_pressure_state(true, &mut stale_defer)
);
assert!(stale_defer.is_none());
}
#[tokio::test]
async fn postgres_gateway_data_state_builds_video_task_reader() {
let state = GatewayDataState::from_config(GatewayDataConfig::from_postgres_url(
+381
View File
@@ -0,0 +1,381 @@
use base64::Engine;
use crate::handlers::shared::{
decrypt_catalog_secret_with_fallbacks, system_config_bool, system_config_string,
};
use crate::{AppState, GatewayError};
const SMTP_TIMEOUT_SECS: u64 = 30;
#[derive(Debug, Clone)]
pub(crate) struct SmtpDeliveryConfig {
pub(crate) host: String,
pub(crate) port: u16,
pub(crate) user: Option<String>,
pub(crate) password: Option<String>,
pub(crate) use_tls: bool,
pub(crate) use_ssl: bool,
pub(crate) from_email: String,
pub(crate) from_name: String,
}
#[derive(Debug, Clone)]
pub(crate) struct ComposedEmail {
pub(crate) to_email: String,
pub(crate) subject: String,
pub(crate) html_body: String,
pub(crate) text_body: String,
}
pub(crate) async fn read_smtp_delivery_config(
state: &AppState,
) -> Result<Option<SmtpDeliveryConfig>, GatewayError> {
let smtp_host = state.read_system_config_json_value("smtp_host").await?;
let smtp_from_email = state
.read_system_config_json_value("smtp_from_email")
.await?;
let Some(host) = system_config_string(smtp_host.as_ref()) else {
return Ok(None);
};
let Some(from_email) = system_config_string(smtp_from_email.as_ref()) else {
return Ok(None);
};
let smtp_port = state.read_system_config_json_value("smtp_port").await?;
let smtp_user = state.read_system_config_json_value("smtp_user").await?;
let smtp_password = state.read_system_config_json_value("smtp_password").await?;
let smtp_use_tls = state.read_system_config_json_value("smtp_use_tls").await?;
let smtp_use_ssl = state.read_system_config_json_value("smtp_use_ssl").await?;
let smtp_from_name = state
.read_system_config_json_value("smtp_from_name")
.await?;
let password = system_config_string(smtp_password.as_ref()).map(|value| {
decrypt_catalog_secret_with_fallbacks(state.encryption_key(), &value).unwrap_or(value)
});
Ok(Some(SmtpDeliveryConfig {
host,
port: system_config_u16(smtp_port.as_ref(), 587),
user: system_config_string(smtp_user.as_ref()),
password,
use_tls: system_config_bool(smtp_use_tls.as_ref(), true),
use_ssl: system_config_bool(smtp_use_ssl.as_ref(), false),
from_email,
from_name: system_config_string(smtp_from_name.as_ref())
.unwrap_or_else(|| "Aether".to_string()),
}))
}
pub(crate) async fn send_smtp_email(
config: SmtpDeliveryConfig,
email: ComposedEmail,
) -> Result<(), GatewayError> {
tokio::task::spawn_blocking(move || send_smtp_email_blocking(config, email))
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
}
pub(crate) async fn probe_smtp_connection(config: SmtpDeliveryConfig) -> Result<(), GatewayError> {
tokio::task::spawn_blocking(move || probe_smtp_connection_blocking(config))
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
}
pub(crate) fn system_config_u16(value: Option<&serde_json::Value>, default: u16) -> u16 {
match value {
Some(serde_json::Value::Number(value)) => value
.as_u64()
.and_then(|value| u16::try_from(value).ok())
.unwrap_or(default),
Some(serde_json::Value::String(value)) => value.trim().parse::<u16>().unwrap_or(default),
_ => default,
}
}
fn encode_mime_header(value: &str) -> String {
if value.is_ascii() {
return value.to_string();
}
format!(
"=?UTF-8?B?{}?=",
base64::engine::general_purpose::STANDARD.encode(value.as_bytes())
)
}
fn wrap_base64(value: &str) -> String {
let mut wrapped = String::new();
for chunk in value.as_bytes().chunks(76) {
wrapped.push_str(std::str::from_utf8(chunk).unwrap_or_default());
wrapped.push_str("\r\n");
}
wrapped
}
fn build_tls_config() -> std::sync::Arc<rustls::ClientConfig> {
let _ = rustls::crypto::ring::default_provider().install_default();
let root_store =
rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
let config = rustls::ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
std::sync::Arc::new(config)
}
fn resolve_server_name(host: &str) -> Result<rustls::pki_types::ServerName<'static>, GatewayError> {
let host = host.trim().trim_start_matches('[').trim_end_matches(']');
if let Ok(ip) = host.parse::<std::net::IpAddr>() {
return Ok(rustls::pki_types::ServerName::from(ip));
}
rustls::pki_types::ServerName::try_from(host.to_string())
.map_err(|err| GatewayError::Internal(err.to_string()))
}
fn connect_tcp_stream(config: &SmtpDeliveryConfig) -> Result<std::net::TcpStream, GatewayError> {
let stream = std::net::TcpStream::connect((config.host.as_str(), config.port))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
stream
.set_read_timeout(Some(std::time::Duration::from_secs(SMTP_TIMEOUT_SECS)))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
stream
.set_write_timeout(Some(std::time::Duration::from_secs(SMTP_TIMEOUT_SECS)))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok(stream)
}
fn wrap_tls_stream(
stream: std::net::TcpStream,
host: &str,
) -> Result<rustls::StreamOwned<rustls::ClientConnection, std::net::TcpStream>, GatewayError> {
let server_name = resolve_server_name(host)?;
let connection = rustls::ClientConnection::new(build_tls_config(), server_name)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok(rustls::StreamOwned::new(connection, stream))
}
fn smtp_read_response<T: std::io::BufRead>(reader: &mut T) -> Result<(u16, String), GatewayError> {
let mut message = String::new();
let code = loop {
let parsed_code;
let continuation;
let trimmed;
{
let mut line = String::new();
let bytes = reader
.read_line(&mut line)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if bytes == 0 {
return Err(GatewayError::Internal(
"smtp connection closed unexpectedly".to_string(),
));
}
trimmed = line.trim_end_matches(['\r', '\n']).to_string();
if trimmed.len() < 3 {
return Err(GatewayError::Internal("invalid smtp response".to_string()));
}
parsed_code = trimmed[..3]
.parse::<u16>()
.map_err(|err| GatewayError::Internal(err.to_string()))?;
continuation = trimmed.as_bytes().get(3).copied() == Some(b'-');
}
if !message.is_empty() {
message.push('\n');
}
message.push_str(&trimmed);
if !continuation {
break parsed_code;
}
};
Ok((code, message))
}
fn smtp_expect<T: std::io::BufRead>(
reader: &mut T,
allowed_codes: &[u16],
) -> Result<String, GatewayError> {
let (code, message) = smtp_read_response(reader)?;
if allowed_codes.contains(&code) {
return Ok(message);
}
Err(GatewayError::Internal(format!(
"unexpected smtp response {code}: {message}"
)))
}
fn smtp_write_line<T: std::io::Write>(writer: &mut T, line: &str) -> Result<(), GatewayError> {
writer
.write_all(line.as_bytes())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
writer
.write_all(b"\r\n")
.map_err(|err| GatewayError::Internal(err.to_string()))?;
writer
.flush()
.map_err(|err| GatewayError::Internal(err.to_string()))
}
fn smtp_send_command<S: std::io::Read + std::io::Write>(
reader: &mut std::io::BufReader<S>,
command: &str,
allowed_codes: &[u16],
) -> Result<String, GatewayError> {
smtp_write_line(reader.get_mut(), command)?;
smtp_expect(reader, allowed_codes)
}
fn build_email_message(config: &SmtpDeliveryConfig, email: &ComposedEmail) -> String {
let boundary = format!("aether-{}", uuid::Uuid::new_v4().simple());
let text_body =
wrap_base64(&base64::engine::general_purpose::STANDARD.encode(email.text_body.as_bytes()));
let html_body =
wrap_base64(&base64::engine::general_purpose::STANDARD.encode(email.html_body.as_bytes()));
let from_header = if config.from_name.trim().is_empty() {
format!("<{}>", config.from_email)
} else {
format!(
"{} <{}>",
encode_mime_header(config.from_name.trim()),
config.from_email
)
};
format!(
"From: {from_header}\r\nTo: <{to_email}>\r\nSubject: {subject}\r\nMIME-Version: 1.0\r\nContent-Type: multipart/alternative; boundary=\"{boundary}\"\r\n\r\n--{boundary}\r\nContent-Type: text/plain; charset=\"utf-8\"\r\nContent-Transfer-Encoding: base64\r\n\r\n{text_body}--{boundary}\r\nContent-Type: text/html; charset=\"utf-8\"\r\nContent-Transfer-Encoding: base64\r\n\r\n{html_body}--{boundary}--\r\n",
to_email = email.to_email,
subject = encode_mime_header(&email.subject),
)
}
fn smtp_authenticate<S: std::io::Read + std::io::Write>(
reader: &mut std::io::BufReader<S>,
config: &SmtpDeliveryConfig,
) -> Result<(), GatewayError> {
let Some(username) = config
.user
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(());
};
let password = config.password.as_deref().unwrap_or("");
smtp_send_command(reader, "AUTH LOGIN", &[334])?;
smtp_send_command(
reader,
&base64::engine::general_purpose::STANDARD.encode(username.as_bytes()),
&[334],
)?;
smtp_send_command(
reader,
&base64::engine::general_purpose::STANDARD.encode(password.as_bytes()),
&[235],
)?;
Ok(())
}
fn smtp_deliver_message<S: std::io::Read + std::io::Write>(
reader: &mut std::io::BufReader<S>,
config: &SmtpDeliveryConfig,
email: &ComposedEmail,
) -> Result<(), GatewayError> {
smtp_send_command(
reader,
&format!("MAIL FROM:<{}>", config.from_email),
&[250],
)?;
smtp_send_command(
reader,
&format!("RCPT TO:<{}>", email.to_email),
&[250, 251],
)?;
smtp_send_command(reader, "DATA", &[354])?;
let message = build_email_message(config, email);
reader
.get_mut()
.write_all(message.as_bytes())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
reader
.get_mut()
.write_all(b"\r\n.\r\n")
.map_err(|err| GatewayError::Internal(err.to_string()))?;
reader
.get_mut()
.flush()
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let _ = smtp_expect(reader, &[250])?;
let _ = smtp_send_command(reader, "QUIT", &[221]);
Ok(())
}
fn smtp_send_message<S: std::io::Read + std::io::Write>(
reader: &mut std::io::BufReader<S>,
config: &SmtpDeliveryConfig,
email: &ComposedEmail,
) -> Result<(), GatewayError> {
smtp_send_command(reader, "EHLO aether.local", &[250])?;
smtp_authenticate(reader, config)?;
smtp_deliver_message(reader, config, email)
}
fn smtp_probe_connection<S: std::io::Read + std::io::Write>(
reader: &mut std::io::BufReader<S>,
config: &SmtpDeliveryConfig,
) -> Result<(), GatewayError> {
smtp_send_command(reader, "EHLO aether.local", &[250])?;
smtp_authenticate(reader, config)?;
let _ = smtp_send_command(reader, "QUIT", &[221]);
Ok(())
}
fn send_smtp_email_blocking(
config: SmtpDeliveryConfig,
email: ComposedEmail,
) -> Result<(), GatewayError> {
if config.use_ssl {
let stream = connect_tcp_stream(&config)?;
let tls_stream = wrap_tls_stream(stream, &config.host)?;
let mut reader = std::io::BufReader::new(tls_stream);
let _ = smtp_expect(&mut reader, &[220])?;
return smtp_send_message(&mut reader, &config, &email);
}
let stream = connect_tcp_stream(&config)?;
let mut reader = std::io::BufReader::new(stream);
let _ = smtp_expect(&mut reader, &[220])?;
let _ = smtp_send_command(&mut reader, "EHLO aether.local", &[250])?;
if config.use_tls {
let _ = smtp_send_command(&mut reader, "STARTTLS", &[220])?;
let stream = reader.into_inner();
let tls_stream = wrap_tls_stream(stream, &config.host)?;
let mut reader = std::io::BufReader::new(tls_stream);
return smtp_send_message(&mut reader, &config, &email);
}
smtp_authenticate(&mut reader, &config)?;
smtp_deliver_message(&mut reader, &config, &email)
}
fn probe_smtp_connection_blocking(config: SmtpDeliveryConfig) -> Result<(), GatewayError> {
if config.use_ssl {
let stream = connect_tcp_stream(&config)?;
let tls_stream = wrap_tls_stream(stream, &config.host)?;
let mut reader = std::io::BufReader::new(tls_stream);
let _ = smtp_expect(&mut reader, &[220])?;
return smtp_probe_connection(&mut reader, &config);
}
let stream = connect_tcp_stream(&config)?;
let mut reader = std::io::BufReader::new(stream);
let _ = smtp_expect(&mut reader, &[220])?;
let _ = smtp_send_command(&mut reader, "EHLO aether.local", &[250])?;
if config.use_tls {
let _ = smtp_send_command(&mut reader, "STARTTLS", &[220])?;
let stream = reader.into_inner();
let tls_stream = wrap_tls_stream(stream, &config.host)?;
let mut reader = std::io::BufReader::new(tls_stream);
return smtp_probe_connection(&mut reader, &config);
}
smtp_authenticate(&mut reader, &config)?;
let _ = smtp_send_command(&mut reader, "QUIT", &[221]);
Ok(())
}
@@ -20,6 +20,7 @@ mod stream_pump;
pub(crate) mod submission;
pub(crate) mod sync;
pub(crate) mod transport;
mod windsurf;
pub(crate) use self::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync;
pub(crate) use self::constants::{
File diff suppressed because it is too large Load Diff
@@ -126,6 +126,54 @@ pub(super) fn build_stream_failure_from_execution_error(
}
}
pub(super) fn build_stream_failure_from_provider_error_body(
status_code: u16,
body_json: &Value,
) -> StreamFailureReport {
let body_object = body_json.as_object();
let error_object = body_object
.and_then(|object| object.get("error"))
.and_then(Value::as_object);
let error_type =
first_non_empty_error_text(error_object, body_object, &["type", "code", "status"])
.unwrap_or_else(|| "upstream_error".to_string());
let error_message = first_non_empty_error_text(
error_object,
body_object,
&["message", "detail", "reason", "status", "type", "code"],
)
.unwrap_or_else(|| format!("upstream stream returned error status {status_code}"));
StreamFailureReport {
status_code,
error_type,
error_message,
extra_error_fields: Map::new(),
}
}
fn first_non_empty_error_text(
error_object: Option<&Map<String, Value>>,
body_object: Option<&Map<String, Value>>,
keys: &[&str],
) -> Option<String> {
for object in [error_object, body_object].into_iter().flatten() {
for key in keys {
let Some(value) = object.get(*key) else {
continue;
};
match value {
Value::String(text) if !text.trim().is_empty() => {
return Some(text.trim().to_string());
}
Value::Number(number) => return Some(number.to_string()),
_ => {}
}
}
}
None
}
fn build_stream_failure_sync_payload(
trace_id: &str,
report_kind: String,
@@ -296,6 +344,49 @@ async fn record_stream_sync_failure(
.await;
}
#[allow(clippy::too_many_arguments)] // internal helper for prefetch error handling
pub(super) async fn handle_prefetch_provider_private_stream_error(
state: &AppState,
trace_id: &str,
decision: &GatewayControlDecision,
plan: &ExecutionPlan,
report_context: Option<Value>,
request_id: &str,
candidate_id: Option<&str>,
report_kind: &str,
mut headers: std::collections::BTreeMap<String, String>,
telemetry: Option<ExecutionTelemetry>,
buffered_body: &[u8],
status_code: u16,
body_json: Value,
) -> Result<Option<Response<Body>>, GatewayError> {
headers.remove("content-encoding");
headers.remove("content-length");
headers.insert("content-type".to_string(), "application/json".to_string());
let payload = GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind: report_kind.to_string(),
report_context,
status_code,
headers,
body_json: Some(body_json),
client_body_json: None,
body_base64: (!buffered_body.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(buffered_body)),
telemetry,
};
record_stream_sync_failure(state, plan, payload.report_context.as_ref(), &payload, None).await;
let response =
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
Ok(Some(attach_control_metadata_headers(
response,
Some(request_id),
candidate_id,
)?))
}
#[allow(clippy::too_many_arguments)] // internal helper for prefetch error handling
pub(super) async fn handle_prefetch_stream_failure(
state: &AppState,
@@ -372,7 +372,10 @@ pub(crate) fn resolve_core_success_background_report_kind(report_kind: &str) ->
core_success_background_report_kind(report_kind).map(ToOwned::to_owned)
}
fn resolve_local_sync_error_status_code(status_code: u16, body_json: &serde_json::Value) -> u16 {
pub(crate) fn resolve_local_sync_error_status_code(
status_code: u16,
body_json: &serde_json::Value,
) -> u16 {
if (400..600).contains(&status_code) {
return status_code;
}
@@ -3,6 +3,7 @@ use std::io::Error as IoError;
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_ai_serving::UPSTREAM_IS_STREAM_KEY;
use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult,
ExecutionTelemetry,
@@ -20,7 +21,6 @@ use async_stream::stream;
use axum::body::{to_bytes, Body, Bytes};
use axum::http::header::{CACHE_CONTROL, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE};
use axum::http::{HeaderName, HeaderValue, Response, StatusCode};
use base64::Engine as _;
use futures_util::StreamExt;
use serde_json::{json, Value};
use tokio::sync::mpsc;
@@ -29,8 +29,9 @@ use tokio::time::MissedTickBehavior;
use tracing::{debug, warn};
use crate::ai_serving::api::{
build_core_error_body_for_client_format, implicit_sync_finalize_report_kind,
maybe_build_sync_finalize_outcome, LocalCoreSyncErrorKind, LocalCoreSyncFinalizeOutcome,
build_core_error_body_for_client_format, extract_provider_private_stream_error_body,
implicit_sync_finalize_report_kind, maybe_build_sync_finalize_outcome, LocalCoreSyncErrorKind,
LocalCoreSyncFinalizeOutcome,
};
use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
@@ -43,12 +44,16 @@ use crate::execution_runtime::grok::maybe_execute_grok_sync;
use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry;
#[cfg(test)]
use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_runtime;
use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
use crate::execution_runtime::transport::{
build_request_body, collect_response_headers, decode_response_body_bytes,
format_upstream_request_error, format_wreq_upstream_request_error, response_body_is_json,
send_request, DirectHttpResponse, DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
use crate::execution_runtime::submission::{
resolve_local_sync_error_status_code, submit_local_core_error_or_sync_finalize,
};
use crate::execution_runtime::transport::{
build_execution_response_body, build_request_body, collect_response_headers,
decode_response_body_bytes, format_upstream_request_error, format_wreq_upstream_request_error,
response_body_is_json, send_request, DirectHttpResponse, DirectSyncExecutionRuntime,
ExecutionRuntimeTransportError,
};
use crate::execution_runtime::windsurf::maybe_execute_windsurf_sync;
use crate::execution_runtime::{
analyze_local_candidate_failover_sync, apply_endpoint_response_header_rules,
attach_provider_response_headers_to_report_context, local_failover_response_text,
@@ -409,6 +414,43 @@ fn build_invalid_provider_success_body(
)
}
fn provider_private_error_details(body_json: &Value) -> (Option<String>, Option<String>) {
let body_object = body_json.as_object();
let error_object = body_object
.and_then(|object| object.get("error"))
.and_then(Value::as_object);
let error_type =
first_non_empty_error_text(error_object, body_object, &["type", "code", "status"]);
let error_message = first_non_empty_error_text(
error_object,
body_object,
&["message", "detail", "reason", "status", "type", "code"],
);
(error_type, error_message)
}
fn first_non_empty_error_text(
error_object: Option<&serde_json::Map<String, Value>>,
body_object: Option<&serde_json::Map<String, Value>>,
keys: &[&str],
) -> Option<String> {
for object in [error_object, body_object].into_iter().flatten() {
for key in keys {
let Some(value) = object.get(*key) else {
continue;
};
match value {
Value::String(text) if !text.trim().is_empty() => {
return Some(text.trim().to_string());
}
Value::Number(number) => return Some(number.to_string()),
_ => {}
}
}
}
None
}
#[derive(Debug, Clone)]
struct OpenAiImageSyncProgressSnapshot {
phase: &'static str,
@@ -799,6 +841,12 @@ async fn execute_direct_sync_runtime_candidate(
candidate_index: &str,
progress_snapshot: Option<Arc<Mutex<OpenAiImageSyncProgressSnapshot>>>,
) -> Result<ExecutionResult, SyncExecutionFailure> {
if let Some(result) = maybe_execute_windsurf_sync(state, plan, report_context)
.await
.map_err(SyncExecutionFailure::from_transport)?
{
return Ok(result);
}
if !should_track_openai_image_sync_upstream_sse(plan_kind, plan, report_context) {
return DirectSyncExecutionRuntime::new()
.execute_sync(plan)
@@ -956,27 +1004,9 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
let upstream_bytes = body_bytes.len() as u64;
progress.finish(status_code, elapsed_ms).await;
let body = if body_bytes.is_empty() {
None
} else if plan.stream {
Some(aether_contracts::ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(&body_bytes)),
})
} else if response_body_is_json(&headers, &decoded_body_bytes) {
let body_json: Value = serde_json::from_slice(&decoded_body_bytes)
.map_err(ExecutionRuntimeTransportError::InvalidJson)
let body =
build_execution_response_body(&headers, &body_bytes, &decoded_body_bytes, plan.stream)
.map_err(SyncExecutionFailure::from_transport)?;
Some(aether_contracts::ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
})
} else {
Some(aether_contracts::ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(&body_bytes)),
})
};
Ok(ExecutionResult {
request_id: plan.request_id.clone(),
@@ -1016,7 +1046,7 @@ fn resolve_openai_image_sync_total_timeout_ms(plan: &ExecutionPlan) -> u64 {
fn report_context_upstream_is_stream(report_context: Option<&Value>) -> bool {
report_context
.and_then(|value| value.get("upstream_is_stream"))
.and_then(|value| value.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool)
.unwrap_or(false)
}
@@ -1707,8 +1737,21 @@ async fn execute_execution_runtime_sync_impl(
headers.insert("content-type".to_string(), "application/json".to_string());
}
}
let (result_error_type, result_error_message) =
let (mut result_error_type, mut result_error_message) =
execution_error_details(result.error.as_ref(), body_json.as_ref());
if result.status_code < 400 && body_json.is_none() {
if let Some(error_body_json) =
extract_provider_private_stream_error_body(report_context.as_ref(), &body_bytes)
{
result.status_code =
resolve_local_sync_error_status_code(result.status_code, &error_body_json);
let (private_error_type, private_error_message) =
provider_private_error_details(&error_body_json);
result_error_type = private_error_type.or(result_error_type);
result_error_message = private_error_message.or(result_error_message);
body_json = Some(error_body_json);
}
}
let local_failover_response_text = local_failover_response_text(
body_json.as_ref(),
&body_bytes,
@@ -25,8 +25,10 @@ use serde_json::json;
use serde_json::Value;
use thiserror::Error;
use crate::ai_serving::api::extract_provider_private_stream_error_body;
#[cfg(test)]
use crate::execution_runtime::remote_compat::execute_sync_plan_via_remote_execution_runtime;
use crate::execution_runtime::windsurf::maybe_execute_windsurf_sync;
use crate::frontdoor_loop_guard::{
configured_gateway_frontdoor_base_url, gateway_frontdoor_self_loop_guard_error,
};
@@ -232,26 +234,8 @@ impl DirectSyncExecutionRuntime {
let elapsed_ms = started_at.elapsed().as_millis() as u64;
let upstream_bytes = body_bytes.len() as u64;
let body = if body_bytes.is_empty() {
None
} else if plan.stream {
Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(&body_bytes)),
})
} else if response_body_is_json(&headers, &decoded_body_bytes) {
let body_json: Value = serde_json::from_slice(&decoded_body_bytes)
.map_err(ExecutionRuntimeTransportError::InvalidJson)?;
Some(ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
})
} else {
Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(&body_bytes)),
})
};
let body =
build_execution_response_body(&headers, &body_bytes, &decoded_body_bytes, plan.stream)?;
Ok(ExecutionResult {
request_id: plan.request_id.clone(),
@@ -347,6 +331,11 @@ pub(crate) async fn execute_sync_plan_with_report_context(
}
let _ = trace_id;
match maybe_execute_windsurf_sync(state, plan, None).await {
Ok(Some(result)) => return Ok(result),
Ok(None) => {}
Err(err) => return Err(GatewayError::Internal(err.to_string())),
}
match DirectSyncExecutionRuntime::new().execute_sync(plan).await {
Ok(result) => {
record_manual_proxy_request_outcome(state, plan, result.status_code).await;
@@ -567,26 +556,8 @@ async fn execute_sync_plan_via_local_tunnel(
);
}
let body = if body_bytes.is_empty() {
None
} else if plan.stream {
Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(&body_bytes)),
})
} else if response_body_is_json(&headers, &decoded_body_bytes) {
let body_json: Value = serde_json::from_slice(&decoded_body_bytes)
.map_err(ExecutionRuntimeTransportError::InvalidJson)?;
Some(ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
})
} else {
Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(&body_bytes)),
})
};
let body =
build_execution_response_body(&headers, &body_bytes, &decoded_body_bytes, plan.stream)?;
Ok(ExecutionResult {
request_id: plan.request_id.clone(),
@@ -1470,17 +1441,63 @@ pub(crate) fn decode_response_body_bytes(
}
pub(crate) fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8]) -> bool {
if headers
let content_type = headers
.get("content-type")
.map(|value| value.to_ascii_lowercase())
.is_some_and(|value| value.contains("json"))
.unwrap_or_default();
if content_type.contains("application/connect+json")
|| content_type.contains("application/connect+proto")
{
return false;
}
if content_type.contains("json") {
return true;
}
serde_json::from_slice::<Value>(body_bytes).is_ok()
}
pub(crate) fn build_execution_response_body(
headers: &BTreeMap<String, String>,
body_bytes: &[u8],
decoded_body_bytes: &[u8],
stream: bool,
) -> Result<Option<ResponseBody>, ExecutionRuntimeTransportError> {
if body_bytes.is_empty() {
return Ok(None);
}
if let Some(body_json) = extract_provider_private_stream_error_body(None, decoded_body_bytes)
.or_else(|| extract_provider_private_stream_error_body(None, body_bytes))
{
return Ok(Some(ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
}));
}
if stream {
return Ok(Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body_bytes)),
}));
}
if response_body_is_json(headers, decoded_body_bytes) {
let body_json: Value = serde_json::from_slice(decoded_body_bytes)
.map_err(ExecutionRuntimeTransportError::InvalidJson)?;
return Ok(Some(ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
}));
}
Ok(Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body_bytes)),
}))
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
@@ -1505,11 +1522,12 @@ mod tests {
use tokio::sync::watch;
use super::{
build_browser_wreq_client, build_client, build_request_headers, execute_sync_plan,
record_manual_proxy_request_failure, record_manual_proxy_request_outcome,
record_manual_proxy_request_success, record_manual_proxy_stream_error,
resolve_execution_transport_controls, DirectSyncExecutionRuntime,
ExecutionRuntimeTransportError, ExecutionTransportControls,
build_browser_wreq_client, build_client, build_execution_response_body,
build_request_headers, execute_sync_plan, record_manual_proxy_request_failure,
record_manual_proxy_request_outcome, record_manual_proxy_request_success,
record_manual_proxy_stream_error, resolve_execution_transport_controls,
response_body_is_json, DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
ExecutionTransportControls,
};
use crate::constants::{
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VIA_TOKEN,
@@ -2764,6 +2782,41 @@ mod tests {
));
}
#[test]
fn connect_json_response_is_not_treated_as_plain_json() {
let headers = BTreeMap::from([(
"content-type".to_string(),
"application/connect+json".to_string(),
)]);
let body = [2, 0, 0, 0, 2, b'{', b'}'];
assert!(!response_body_is_json(&headers, &body));
}
#[test]
fn connect_json_error_response_is_decoded_for_stream_sync_body() {
let headers = BTreeMap::from([(
"content-type".to_string(),
"application/connect+json".to_string(),
)]);
let payload = br#"{"error":{"code":"resource_exhausted","message":"quota exhausted"}}"#;
let mut body_bytes = vec![2];
body_bytes.extend_from_slice(&(payload.len() as u32).to_be_bytes());
body_bytes.extend_from_slice(payload);
let body = build_execution_response_body(&headers, &body_bytes, &body_bytes, true)
.expect("body should build")
.expect("body should be present");
assert_eq!(
body.json_body
.as_ref()
.and_then(|value| value.pointer("/error/code")),
Some(&json!("resource_exhausted"))
);
assert!(body.body_bytes_b64.is_none());
}
#[tokio::test]
async fn direct_sync_execution_runtime_compresses_json_body_when_requested() {
let listener = crate::test_support::bind_loopback_listener()
File diff suppressed because it is too large Load Diff
@@ -1,5 +1,6 @@
use aether_ai_serving::{
run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt,
UPSTREAM_IS_STREAM_KEY,
};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::{
@@ -471,11 +472,24 @@ fn should_skip_unused_persistence_from_metadata(
metadata.candidate_group_id.is_some() && metadata.pool_key_index.is_some()
}
fn resolve_stream_candidate_watchdog_timeout(plan: &aether_contracts::ExecutionPlan) -> Duration {
fn resolve_stream_candidate_watchdog_timeout(
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Duration {
let upstream_is_stream = report_context
.and_then(|context| context.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.unwrap_or(true);
let timeout_ms = plan
.timeouts
.as_ref()
.and_then(|timeouts| timeouts.first_byte_ms.or(timeouts.total_ms))
.and_then(|timeouts| {
if upstream_is_stream {
timeouts.first_byte_ms.or(timeouts.total_ms)
} else {
timeouts.total_ms.or(timeouts.first_byte_ms)
}
})
.unwrap_or(DEFAULT_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS)
.max(1);
Duration::from_millis(timeout_ms)
@@ -493,7 +507,7 @@ where
Fut:
std::future::Future<Output = Result<Option<Response<Body>>, GatewayError>> + Send + 'static,
{
let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan);
let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan, report_context);
let candidate_started_unix_ms = current_unix_ms();
let mut join_handle = tokio::spawn(execute());
match timeout(timeout_duration, &mut join_handle).await {
@@ -655,19 +669,22 @@ mod tests {
#[test]
fn stream_candidate_watchdog_prefers_first_byte_timeout() {
let timeout =
resolve_stream_candidate_watchdog_timeout(&test_plan(Some(ExecutionTimeouts {
let report_context = json!({"upstream_is_stream": true});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(12_345),
total_ms: Some(90_000),
..ExecutionTimeouts::default()
})));
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(12_345));
}
#[test]
fn stream_candidate_watchdog_uses_default_when_timeouts_missing() {
let timeout = resolve_stream_candidate_watchdog_timeout(&test_plan(None));
let timeout = resolve_stream_candidate_watchdog_timeout(&test_plan(None), None);
assert_eq!(
timeout,
@@ -675,6 +692,50 @@ mod tests {
);
}
#[test]
fn stream_candidate_watchdog_prefers_total_timeout_when_upstream_non_stream() {
let report_context = json!({"upstream_is_stream": false});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(300_000),
total_ms: Some(599_000),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(599_000));
}
#[test]
fn stream_candidate_watchdog_falls_back_to_first_byte_when_upstream_non_stream_lacks_total() {
let report_context = json!({"upstream_is_stream": false});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(300_000),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(300_000));
}
#[test]
fn stream_candidate_watchdog_defaults_to_streaming_when_flag_missing() {
let report_context = json!({});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(12_345),
total_ms: Some(90_000),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(12_345));
}
#[test]
fn unused_persistence_skips_pool_internal_candidates() {
assert!(should_skip_unused_persistence(Some(&json!({
@@ -37,6 +37,7 @@ pub(crate) use self::provider::oauth::runtime::{
refresh_provider_oauth_account_state_after_update,
};
pub(crate) use self::provider::ops::providers::actions::admin_provider_ops_local_action_response;
pub(crate) use self::provider::ops::providers::store_admin_provider_ops_balance_cache;
pub(crate) use self::provider::pool::config::admin_provider_pool_config;
pub(crate) use self::provider::pool_admin::maybe_build_local_admin_pool_response;
pub(crate) use self::provider::shared::payloads::{
@@ -59,6 +59,11 @@ pub(super) async fn build_admin_monitoring_redis_cache_categories_response(
) -> Result<Response<Body>, GatewayError> {
let mut categories = Vec::with_capacity(ADMIN_MONITORING_REDIS_CACHE_CATEGORIES.len());
let mut total_keys = 0usize;
let diagnostics = state
.runtime_state()
.redis_diagnostics()
.await
.map_err(|err| GatewayError::Internal(format!("redis diagnostics failed: {err}")))?;
for (key, name, pattern, description) in ADMIN_MONITORING_REDIS_CACHE_CATEGORIES {
let count = list_admin_monitoring_namespaced_keys(state, pattern)
@@ -81,6 +86,7 @@ pub(super) async fn build_admin_monitoring_redis_cache_categories_response(
"backend": state.runtime_state().backend_kind().as_str(),
"categories": categories,
"total_keys": total_keys,
"diagnostics": diagnostics,
}
}))
.into_response())
@@ -1194,6 +1194,7 @@ async fn admin_monitoring_redis_keys_returns_local_payload_without_redis() {
assert_eq!(payload["data"]["available"], json!(true));
assert_eq!(payload["data"]["backend"], json!("memory"));
assert_eq!(payload["data"]["total_keys"], json!(0));
assert_eq!(payload["data"]["diagnostics"], serde_json::Value::Null);
}
#[tokio::test]
@@ -5,8 +5,11 @@ use super::local_monitoring_response;
use crate::data::GatewayDataState;
use crate::AppState;
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_data_contracts::repository::{
candidates::RequestCandidateStatus, usage::UsageBodyCaptureState,
};
use axum::body::to_bytes;
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
use serde_json::json;
use std::sync::Arc;
@@ -530,6 +533,89 @@ async fn admin_monitoring_trace_request_exposes_failed_candidate_upstream_respon
assert!(extra.get("provider_response").is_none());
}
#[tokio::test]
async fn admin_monitoring_trace_request_decodes_connect_json_response_body_refs() {
let mut candidate = sample_candidate(
"cand-used",
"request-connect",
0,
RequestCandidateStatus::Failed,
Some(101),
Some(33),
Some(429),
);
candidate.extra_data = Some(json!({"cache_1h": true}));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![candidate]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let mut usage = sample_usage(
"request-connect",
"provider-1",
"Windsurf",
0,
0.0,
"failed",
Some(429),
100,
);
usage.candidate_id = Some("cand-used".to_string());
usage.response_headers = Some(json!({
"content-type": "application/connect+json"
}));
let mut framed = Vec::new();
framed.push(2);
let payload = br#"{"error":{"code":"resource_exhausted","message":"quota exhausted"}}"#;
framed.extend_from_slice(&(payload.len() as u32).to_be_bytes());
framed.extend_from_slice(payload);
usage.response_body = Some(json!(BASE64_STANDARD.encode(framed)));
usage.response_body_ref = Some("usage://request/request-connect/response_body".to_string());
usage.response_body_state = Some(UsageBodyCaptureState::Inline);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
let data_state =
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
request_candidates,
usage_repository,
)
.with_provider_catalog_reader(provider_catalog);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let context = request_context(
http::Method::GET,
"/api/admin/monitoring/trace/request-connect",
);
let response = local_monitoring_response(&state, &context)
.await
.expect("handler should not error")
.expect("route should be handled locally");
assert_eq!(response.status(), http::StatusCode::OK);
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");
let upstream_response = &payload["candidates"][0]["extra_data"]["upstream_response"];
assert_eq!(upstream_response["status_code"], json!(429));
assert_eq!(
upstream_response["body"]["error"]["code"],
json!("resource_exhausted")
);
assert_eq!(
upstream_response["body"]["error"]["message"],
json!("quota exhausted")
);
assert_eq!(
upstream_response["body_ref"],
json!("usage://request/request-connect/response_body")
);
assert_eq!(upstream_response["body_state"], json!("inline"));
}
#[tokio::test]
async fn admin_monitoring_trace_request_exposes_structured_ranking_metadata() {
let mut candidate = sample_candidate(
@@ -1,6 +1,6 @@
use crate::handlers::admin::request::AdminAppState;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use std::time::{SystemTime, UNIX_EPOCH};
@@ -25,10 +25,15 @@ use crate::handlers::admin::provider::oauth::runtime::{
use crate::handlers::admin::provider::oauth::state::{
admin_provider_oauth_template, exchange_admin_provider_oauth_refresh_token,
};
use crate::handlers::admin::provider::shared::support::ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL;
use crate::handlers::admin::request::{AdminAppState, AdminProviderOAuthTemplate};
use crate::GatewayError;
use aether_admin::provider::oauth::parse_admin_provider_oauth_kiro_batch_import_entries;
use aether_contracts::ProxySnapshot;
use aether_oauth::core::OAuthError;
use aether_oauth::provider::{
ProviderOAuthImportInput, ProviderOAuthService, ProviderOAuthTransportContext,
};
use serde_json::{json, Map, Value};
struct AdminProviderOAuthResolvedBatchImport {
@@ -37,6 +42,16 @@ struct AdminProviderOAuthResolvedBatchImport {
expires_at: Option<u64>,
}
fn sanitize_windsurf_batch_import_error(error: &OAuthError) -> String {
match error {
OAuthError::InvalidRequest(_) => "Windsurf 凭据验证失败: 请求参数无效".to_string(),
OAuthError::HttpStatus { status_code, .. } => {
format!("Windsurf 凭据验证失败: HTTP {status_code}")
}
_ => "Windsurf 凭据验证失败".to_string(),
}
}
pub(super) fn estimate_admin_provider_oauth_batch_import_total(
provider_type: &str,
raw_credentials: &str,
@@ -98,6 +113,61 @@ async fn resolve_admin_provider_oauth_batch_import_tokens(
.map(str::trim)
.filter(|value| !value.is_empty());
if provider_type.eq_ignore_ascii_case("windsurf") {
let token_for_import = refresh_token.or(access_token);
let ctx = ProviderOAuthTransportContext {
provider_id: String::new(),
provider_type: provider_type.to_string(),
endpoint_id: None,
key_id: None,
auth_type: Some("oauth".to_string()),
decrypted_api_key: None,
decrypted_auth_config: None,
provider_config: None,
endpoint_config: None,
key_config: None,
network: aether_oauth::network::OAuthNetworkContext::provider_operation(
request_proxy.clone(),
),
};
let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state);
let result = ProviderOAuthService::with_builtin_adapters()
.import_credentials(
&executor,
&ctx,
ProviderOAuthImportInput {
provider_type: provider_type.to_string(),
name: entry
.raw_credentials
.as_ref()
.and_then(|raw| raw.get("name"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
refresh_token: token_for_import.map(ToOwned::to_owned),
raw_credentials: entry.raw_credentials.clone(),
network: ctx.network.clone(),
},
)
.await
.map_err(|error| sanitize_windsurf_batch_import_error(&error))?;
let access_token = result.token_set.access_token.trim().to_string();
if access_token.is_empty() {
return Err("Windsurf 凭据验证返回缺少 apiKey/sessionToken".to_string());
}
let auth_config = result
.auth_config
.as_object()
.cloned()
.ok_or_else(|| "Windsurf 凭据验证返回缺少 auth_config".to_string())?;
return Ok(AdminProviderOAuthResolvedBatchImport {
access_token,
auth_config,
expires_at: result.token_set.expires_at_unix_secs,
});
}
if let Some(refresh_token) = refresh_token {
let Some(template) = template else {
if provider_type_supports_access_token_import(provider_type) {
@@ -228,6 +298,28 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
};
let template = admin_provider_oauth_template(provider_type);
if template.is_none()
&& !provider_type.eq_ignore_ascii_case("windsurf")
&& !provider_type_supports_access_token_import(provider_type)
{
return Ok(AdminProviderOAuthBatchImportOutcome {
total: entries.len(),
success: 0,
failed: entries.len(),
results: entries
.iter()
.enumerate()
.map(|(index, _)| {
json!({
"index": index,
"status": "error",
"error": ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL,
"replaced": false,
})
})
.collect(),
});
}
let endpoint_resolution =
resolve_provider_oauth_runtime_endpoints(state, &provider, provider_type).await?;
@@ -251,6 +343,25 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
let mut failed = 0usize;
for (index, entry) in entries.iter().enumerate() {
if let Some(error) = entry.parse_error.as_ref() {
failed += 1;
results.push(json!({
"index": index,
"status": "error",
"error": error,
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
}
let resolved_import = match resolve_admin_provider_oauth_batch_import_tokens(
state,
template,
@@ -418,3 +529,33 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
results,
})
}
#[cfg(test)]
mod tests {
use super::sanitize_windsurf_batch_import_error;
use aether_oauth::core::OAuthError;
#[test]
fn windsurf_batch_import_error_redacts_http_body() {
let error = OAuthError::HttpStatus {
status_code: 401,
body_excerpt: "sessionToken=devin-session-token$secret".to_string(),
};
let detail = sanitize_windsurf_batch_import_error(&error);
assert_eq!(detail, "Windsurf 凭据验证失败: HTTP 401");
assert!(!detail.contains("devin-session-token$secret"));
}
#[test]
fn windsurf_batch_import_error_redacts_provider_detail() {
let error = OAuthError::invalid_response("apiKey=sk-secret token=secret-token");
let detail = sanitize_windsurf_batch_import_error(&error);
assert_eq!(detail, "Windsurf 凭据验证失败");
assert!(!detail.contains("sk-secret"));
assert!(!detail.contains("secret-token"));
}
}
@@ -8,7 +8,7 @@ use super::parse::{
};
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::state::{
build_admin_provider_oauth_backend_unavailable_response,
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
is_fixed_provider_type_for_provider_oauth,
};
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_provider_id;
@@ -60,6 +60,13 @@ pub(in super::super) async fn handle_admin_provider_oauth_batch_import(
"该 Provider 不是固定类型,无法使用 provider-oauth",
));
}
if provider_type != "kiro"
&& provider_type != "windsurf"
&& admin_provider_oauth_template(&provider_type).is_none()
{
return Ok(build_admin_provider_oauth_backend_unavailable_response());
}
let total = estimate_admin_provider_oauth_batch_import_total(
&provider_type,
payload.credentials.as_str(),
@@ -19,8 +19,10 @@ pub(super) struct AdminProviderOAuthBatchImportRequest {
#[derive(Debug, Clone)]
pub(super) struct AdminProviderOAuthBatchImportEntry {
pub parse_error: Option<String>,
pub refresh_token: Option<String>,
pub access_token: Option<String>,
pub raw_credentials: Option<serde_json::Value>,
pub expires_at: Option<u64>,
pub account_id: Option<String>,
pub account_user_id: Option<String>,
@@ -156,8 +158,10 @@ fn extract_admin_provider_oauth_batch_import_entry(
access_token.as_deref(),
);
Some(AdminProviderOAuthBatchImportEntry {
parse_error: None,
refresh_token,
access_token,
raw_credentials: None,
expires_at: None,
account_id: None,
account_user_id: None,
@@ -176,6 +180,8 @@ fn extract_admin_provider_oauth_batch_import_entry(
}
}
serde_json::Value::Object(object) => {
let is_grok = provider_type.trim().eq_ignore_ascii_case("grok");
let is_windsurf = provider_type.trim().eq_ignore_ascii_case("windsurf");
let refresh_token = coerce_admin_provider_oauth_import_str(
object
.get("refresh_token")
@@ -186,12 +192,8 @@ fn extract_admin_provider_oauth_batch_import_entry(
.get("access_token")
.or_else(|| object.get("accessToken")),
);
let grok_token_alias = if provider_type.trim().eq_ignore_ascii_case("grok") {
object.get("token")
} else {
None
};
let grok_cookie = if provider_type.trim().eq_ignore_ascii_case("grok") {
let grok_token_alias = if is_grok { object.get("token") } else { None };
let grok_cookie = if is_grok {
coerce_admin_provider_oauth_import_str(
object.get("cookie").or_else(|| object.get("cookieHeader")),
)
@@ -214,9 +216,43 @@ fn extract_admin_provider_oauth_batch_import_entry(
refresh_token.as_deref(),
access_token.as_deref().or(session_token.as_deref()),
);
if refresh_token.is_none() && access_token.is_none() {
let windsurf_api_key = is_windsurf
.then(|| {
coerce_admin_provider_oauth_import_str(
object.get("api_key").or_else(|| object.get("apiKey")),
)
})
.flatten();
let windsurf_token = is_windsurf
.then(|| {
coerce_admin_provider_oauth_import_str(
object
.get("token")
.or_else(|| object.get("auth_token"))
.or_else(|| object.get("authToken")),
)
})
.flatten();
let windsurf_password = is_windsurf
.then(|| coerce_admin_provider_oauth_import_str(object.get("password")))
.flatten();
let raw_credentials = if is_windsurf
&& (windsurf_api_key.is_some()
|| windsurf_token.is_some()
|| windsurf_password.is_some())
{
Some(item.clone())
} else {
None
};
if refresh_token.is_none() && access_token.is_none() && raw_credentials.is_none() {
return None;
}
let refresh_token = if is_windsurf {
refresh_token.or(windsurf_api_key).or(windsurf_token)
} else {
refresh_token
};
let expires_at =
json_u64_value(object.get("expires_at").or_else(|| object.get("expiresAt")));
let account_id = coerce_admin_provider_oauth_import_str(
@@ -308,8 +344,10 @@ fn extract_admin_provider_oauth_batch_import_entry(
.or_else(|| object.get("impersonate")),
);
Some(AdminProviderOAuthBatchImportEntry {
parse_error: None,
refresh_token,
access_token,
raw_credentials,
expires_at,
account_id,
account_user_id,
@@ -340,14 +378,17 @@ pub(super) fn parse_admin_provider_oauth_batch_import_entries(
}
if raw.starts_with('[') {
if let Ok(serde_json::Value::Array(items)) = serde_json::from_str::<serde_json::Value>(raw)
{
return items
.iter()
.filter_map(|item| {
extract_admin_provider_oauth_batch_import_entry(provider_type, item)
})
.collect();
match serde_json::from_str::<serde_json::Value>(raw) {
Ok(serde_json::Value::Array(items)) => {
return items
.iter()
.filter_map(|item| {
extract_admin_provider_oauth_batch_import_entry(provider_type, item)
})
.collect();
}
Ok(_) => {}
Err(error) => return vec![parse_error_entry(format!("JSON 数组解析失败: {error}"))],
}
}
@@ -364,23 +405,61 @@ pub(super) fn parse_admin_provider_oauth_batch_import_entries(
raw.lines()
.map(str::trim)
.filter(|line| !line.is_empty() && !line.starts_with('#'))
.filter_map(|line| {
if line.starts_with('{') {
return serde_json::from_str::<serde_json::Value>(line)
.ok()
.and_then(|value| {
extract_admin_provider_oauth_batch_import_entry(provider_type, &value)
});
.filter_map(|token| {
if is_json_like_batch_line(token) {
match serde_json::from_str::<serde_json::Value>(token) {
Ok(value @ serde_json::Value::Object(_)) => {
return extract_admin_provider_oauth_batch_import_entry(
provider_type,
&value,
);
}
Ok(_) => {
return Some(parse_error_entry(
"JSON 行必须是账号对象,不能作为 raw token 导入".to_string(),
));
}
Err(error) => {
return Some(parse_error_entry(format!("JSON 行解析失败: {error}")));
}
}
}
extract_admin_provider_oauth_batch_import_entry(
provider_type,
&serde_json::Value::String(line.to_string()),
&serde_json::Value::String(token.to_string()),
)
})
.collect()
}
fn parse_error_entry(error: String) -> AdminProviderOAuthBatchImportEntry {
AdminProviderOAuthBatchImportEntry {
parse_error: Some(error),
refresh_token: None,
access_token: None,
raw_credentials: None,
expires_at: None,
account_id: None,
account_user_id: None,
plan_type: None,
pool_tier: None,
user_id: None,
email: None,
account_name: None,
sso_rw_token: None,
cf_cookies: None,
cf_clearance: None,
user_agent: None,
browser_profile: None,
}
}
fn is_json_like_batch_line(line: &str) -> bool {
let line = line.trim_start();
line.starts_with('{') || line.starts_with('[')
}
pub(super) fn apply_admin_provider_oauth_batch_import_hints(
provider_type: &str,
entry: &AdminProviderOAuthBatchImportEntry,
@@ -705,4 +784,115 @@ mod tests {
Some(&json!("project-gemini-cli-2"))
);
}
#[test]
fn parses_windsurf_json_credentials_for_native_import() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"windsurf",
r#"[
{"api_key":"devin-session-token$abc","email":"a@example.com"},
{"token":"firebase-id-token","name":"Browser Login"},
{"email":"b@example.com","password":"secret"},
{"access_token":"devin-session-token$alias","email":"c@example.com"}
]"#,
);
assert_eq!(entries.len(), 4);
assert_eq!(
entries[0].refresh_token.as_deref(),
Some("devin-session-token$abc")
);
assert_eq!(entries[0].email.as_deref(), Some("a@example.com"));
assert_eq!(
entries[0]
.raw_credentials
.as_ref()
.and_then(|value| value.get("api_key")),
Some(&json!("devin-session-token$abc"))
);
assert_eq!(
entries[1]
.raw_credentials
.as_ref()
.and_then(|value| value.get("token")),
Some(&json!("firebase-id-token"))
);
assert_eq!(
entries[2]
.raw_credentials
.as_ref()
.and_then(|value| value.get("password")),
Some(&json!("secret"))
);
assert_eq!(
entries[3].access_token.as_deref(),
Some("devin-session-token$alias")
);
}
#[test]
fn parses_windsurf_json_lines_credentials_for_native_import() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"windsurf",
r#"{"api_key":"devin-session-token$abc","email":"a@example.com"}
{"token":"firebase-id-token","name":"Browser Login"}
{"email":"b@example.com","password":"secret"}"#,
);
assert_eq!(entries.len(), 3);
assert_eq!(
entries[0]
.raw_credentials
.as_ref()
.and_then(|value| value.get("api_key")),
Some(&json!("devin-session-token$abc"))
);
assert_eq!(
entries[1]
.raw_credentials
.as_ref()
.and_then(|value| value.get("token")),
Some(&json!("firebase-id-token"))
);
assert_eq!(
entries[2]
.raw_credentials
.as_ref()
.and_then(|value| value.get("password")),
Some(&json!("secret"))
);
}
#[test]
fn invalid_json_line_is_parse_error_not_token() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"windsurf",
r#"{"email":"b@example.com","password":"secret""#,
);
assert_eq!(entries.len(), 1);
assert!(entries[0].parse_error.is_some());
assert!(entries[0].refresh_token.is_none());
assert!(entries[0].access_token.is_none());
assert!(entries[0].raw_credentials.is_none());
}
#[test]
fn json_like_line_after_token_is_parse_error_not_token() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"windsurf",
"devin-session-token$abc\n[not-json",
);
assert_eq!(entries.len(), 2);
assert!(entries[0].parse_error.is_none());
assert_eq!(
entries[0].refresh_token.as_deref(),
Some("devin-session-token$abc")
);
assert!(entries[1].parse_error.is_some());
assert!(entries[1].refresh_token.is_none());
assert!(entries[1].access_token.is_none());
assert!(entries[1].raw_credentials.is_none());
}
}
@@ -10,7 +10,7 @@ use super::progress::{
};
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::state::{
build_admin_provider_oauth_backend_unavailable_response,
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
is_fixed_provider_type_for_provider_oauth,
};
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_task_provider_id;
@@ -124,6 +124,13 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
"该 Provider 不是固定类型,无法使用 provider-oauth",
));
}
if provider_type != "kiro"
&& provider_type != "windsurf"
&& admin_provider_oauth_template(&provider_type).is_none()
{
return Ok(build_admin_provider_oauth_backend_unavailable_response());
}
let total = estimate_admin_provider_oauth_batch_import_total(
&provider_type,
payload.credentials.as_str(),
@@ -12,6 +12,7 @@ use crate::GatewayError;
use aether_data::repository::provider_oauth::{
StoredAdminProviderOAuthDeviceSession, KIRO_DEVICE_AUTH_SESSION_TTL_BUFFER_SECS,
};
use aether_oauth::provider::{ProviderOAuthService, ProviderOAuthTransportContext};
use axum::{
body::{Body, Bytes},
http,
@@ -28,6 +29,8 @@ const KIRO_SOCIAL_MANUAL_CALLBACK_PORT: u16 = 49153;
const KIRO_SOCIAL_ALLOWED_CALLBACK_PORTS: &[u16] = &[
3128, 4649, 6588, 8008, 9091, 49153, 50153, 51153, 52153, 53153,
];
const WINDSURF_BROWSER_AUTH_EXPIRES_IN_SECS: u64 = 600;
const WINDSURF_BROWSER_AUTH_POLL_INTERVAL_SECS: u64 = 5;
fn normalize_kiro_device_auth_type(raw: Option<&str>) -> String {
match raw
@@ -119,6 +122,26 @@ fn build_kiro_social_authorization_url(
)
}
fn build_windsurf_authorization_url(authorize_url: &str, login_option: &str) -> String {
let login_option = login_option.trim();
if login_option.is_empty() {
return authorize_url.to_string();
}
if let Ok(mut url) = Url::parse(authorize_url) {
url.query_pairs_mut()
.append_pair("login_option", login_option);
return url.to_string();
}
let separator = if authorize_url.contains('?') {
'&'
} else {
'?'
};
let mut serializer = form_urlencoded::Serializer::new(String::new());
serializer.append_pair("login_option", login_option);
format!("{authorize_url}{separator}{}", serializer.finish())
}
pub(super) async fn handle_admin_provider_oauth_device_authorize(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
@@ -163,14 +186,14 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
));
};
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
if provider_type != "kiro" {
if provider_type != "kiro" && provider_type != "windsurf" {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"设备授权仅支持 Kiro provider",
"设备授权仅支持 Kiro / Windsurf provider",
));
}
let endpoint_resolution =
resolve_provider_oauth_runtime_endpoints(state, &provider, "kiro").await?;
resolve_provider_oauth_runtime_endpoints(state, &provider, &provider_type).await?;
let runtime_endpoint = endpoint_resolution.runtime_endpoint;
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
@@ -184,6 +207,103 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
)
.await;
if provider_type == "windsurf" {
let session_id = generate_provider_oauth_nonce();
let login_option = payload
.login_option
.as_deref()
.or(payload.auth_type.as_deref())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("default")
.to_ascii_lowercase();
let ctx = ProviderOAuthTransportContext {
provider_id: provider_id.clone(),
provider_type: provider_type.clone(),
endpoint_id: runtime_endpoint
.as_ref()
.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
.as_ref()
.and_then(|endpoint| endpoint.config.clone()),
key_config: None,
network: aether_oauth::network::OAuthNetworkContext::provider_operation(
request_proxy.clone(),
),
};
let mut authorization = match ProviderOAuthService::with_builtin_adapters()
.build_authorize_url(&ctx, &session_id, None)
{
Ok(authorization) => authorization,
Err(error) => {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
format!("Windsurf 授权 URL 构建失败: {error}"),
));
}
};
authorization.authorize_url =
build_windsurf_authorization_url(&authorization.authorize_url, &login_option);
let now_unix_secs = current_unix_secs();
let session = StoredAdminProviderOAuthDeviceSession {
provider_id: provider_id.clone(),
region: String::new(),
client_id: String::new(),
client_secret: String::new(),
device_code: String::new(),
auth_type: Some("browser".to_string()),
social_provider: Some(login_option.clone()),
code_verifier: None,
redirect_uri: Some("show-auth-token".to_string()),
machine_id: Some(uuid::Uuid::new_v4().to_string().to_ascii_lowercase()),
interval: WINDSURF_BROWSER_AUTH_POLL_INTERVAL_SECS,
expires_at_unix_secs: now_unix_secs
.saturating_add(WINDSURF_BROWSER_AUTH_EXPIRES_IN_SECS),
status: "pending".to_string(),
proxy_node_id: payload
.proxy_node_id
.as_deref()
.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,
WINDSURF_BROWSER_AUTH_EXPIRES_IN_SECS
.saturating_add(KIRO_DEVICE_AUTH_SESSION_TTL_BUFFER_SECS),
)
.await
{
return Ok(response);
}
return Ok(Json(json!({
"session_id": session_id,
"user_code": "",
"verification_uri": "https://windsurf.com/windsurf/signin",
"verification_uri_complete": authorization.authorize_url,
"expires_in": WINDSURF_BROWSER_AUTH_EXPIRES_IN_SECS,
"interval": WINDSURF_BROWSER_AUTH_POLL_INTERVAL_SECS,
"auth_type": "browser",
"login_option": login_option,
"redirect_uri": "show-auth-token",
"callback_required": true,
}))
.into_response());
}
let auth_type = normalize_kiro_device_auth_type(payload.auth_type.as_deref());
if let Some(social_provider) = kiro_social_provider_id(&auth_type) {
let redirect_uri = match normalize_kiro_social_redirect_uri(payload.redirect_uri.as_deref())
@@ -394,3 +514,27 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
}))
.into_response())
}
#[cfg(test)]
mod tests {
use super::build_windsurf_authorization_url;
#[test]
fn windsurf_authorization_url_includes_login_option() {
let url = build_windsurf_authorization_url(
"https://windsurf.com/windsurf/signin?state=session-1",
"github",
);
let parsed = url::Url::parse(&url).expect("url should parse");
let params = parsed
.query_pairs()
.map(|(key, value)| (key.to_string(), value.to_string()))
.collect::<std::collections::BTreeMap<_, _>>();
assert_eq!(params.get("state").map(String::as_str), Some("session-1"));
assert_eq!(
params.get("login_option").map(String::as_str),
Some("github")
);
}
}
@@ -26,6 +26,9 @@ use aether_data::repository::provider_oauth::StoredAdminProviderOAuthDeviceSessi
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
use aether_oauth::provider::{
ProviderOAuthImportInput, ProviderOAuthService, ProviderOAuthTransportContext,
};
use axum::{
body::{Body, Bytes},
http,
@@ -86,6 +89,99 @@ fn kiro_social_poll_error_response(error: impl Into<String>) -> Response<Body> {
.into_response()
}
fn windsurf_browser_poll_error_response(error: impl Into<String>) -> Response<Body> {
Json(json!({
"status": "error",
"error": error.into(),
"replaced": false,
}))
.into_response()
}
fn sanitize_windsurf_browser_poll_detail(detail: impl AsRef<str>) -> String {
let detail = detail.as_ref().trim();
if detail.is_empty() {
return "-".to_string();
}
if contains_windsurf_sensitive_marker(detail) {
"[REDACTED upstream error body]".to_string()
} else {
detail.chars().take(500).collect()
}
}
fn sanitize_windsurf_browser_poll_callback_error(error: &str, description: &str) -> String {
let error = sanitize_windsurf_browser_poll_error_code(error);
let description = sanitize_windsurf_browser_poll_detail(description);
format!("{error}: {description}")
}
fn sanitize_windsurf_browser_poll_error_code(error: &str) -> String {
let error = error.trim();
if !error.is_empty()
&& error.len() <= 80
&& error
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-' | '.'))
{
return error.to_string();
}
sanitize_windsurf_browser_poll_detail(error)
}
fn sanitize_windsurf_browser_poll_oauth_error(error: &aether_oauth::core::OAuthError) -> String {
match error {
aether_oauth::core::OAuthError::InvalidRequest(_) => {
"Windsurf token 验证失败: 请求参数无效".to_string()
}
aether_oauth::core::OAuthError::HttpStatus { status_code, .. } => {
format!("Windsurf token 验证失败: HTTP {status_code}")
}
_ => "Windsurf token 验证失败".to_string(),
}
}
fn contains_windsurf_sensitive_marker(value: &str) -> bool {
let lowered = value.to_ascii_lowercase();
[
"token",
"api_key",
"apikey",
"sessiontoken",
"firebase_id_token",
"idtoken",
"authorization",
"password",
"secret",
"devin-session-token$",
]
.iter()
.any(|marker| lowered.contains(marker))
|| value.contains("sk-")
}
fn secret_fingerprint(value: &str) -> Option<String> {
let value = value.trim();
if value.is_empty() {
return None;
}
use sha2::{Digest, Sha256};
let digest = Sha256::digest(value.as_bytes());
Some(
digest[..8]
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>(),
)
}
fn insert_secret_fingerprint(target: &mut serde_json::Map<String, Value>, key: &str, secret: &str) {
if let Some(fingerprint) = secret_fingerprint(secret) {
target.insert(key.to_string(), json!(fingerprint));
}
}
fn kiro_social_provider_from_login_option(login_option: Option<&str>) -> Option<&'static str> {
match login_option
.map(str::trim)
@@ -322,8 +418,9 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
"Provider 不存在",
));
};
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
let endpoint_resolution =
resolve_provider_oauth_runtime_endpoints(state, &provider, "kiro").await?;
resolve_provider_oauth_runtime_endpoints(state, &provider, &provider_type).await?;
let endpoints = endpoint_resolution.endpoints;
let runtime_endpoint = endpoint_resolution.runtime_endpoint;
let request_proxy = state
@@ -338,6 +435,20 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
)
.await;
if provider_type == "windsurf" {
return handle_admin_provider_oauth_windsurf_browser_device_poll(
state,
&provider,
&endpoints,
request_proxy,
session_id,
session,
payload.callback_url.as_deref(),
payload.token.as_deref(),
)
.await;
}
if kiro_device_session_is_social(&session) {
return handle_admin_provider_oauth_kiro_social_device_poll(
state,
@@ -630,6 +741,276 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
))
}
fn windsurf_raw_api_key(value: &str) -> Option<&str> {
let value = value.trim();
if value.starts_with("devin-session-token$") || value.starts_with("sk-") {
Some(value)
} else {
None
}
}
async fn handle_admin_provider_oauth_windsurf_browser_device_poll(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoints: &[StoredProviderCatalogEndpoint],
request_proxy: Option<ProxySnapshot>,
session_id: &str,
mut session: StoredAdminProviderOAuthDeviceSession,
callback_url: Option<&str>,
token: Option<&str>,
) -> Result<Response<Body>, GatewayError> {
let callback_url = callback_url
.map(str::trim)
.filter(|value| !value.is_empty());
let token = token.map(str::trim).filter(|value| !value.is_empty());
if callback_url.is_none() && token.is_none() {
return Ok(Json(json!({"status": "pending", "replaced": false})).into_response());
}
let mut social_provider = session
.social_provider
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let imported_token = if let Some(callback_url) = callback_url {
let callback_params = parse_provider_oauth_callback_params(callback_url);
if let Some(error) = callback_params.get("error").map(String::as_str) {
let error_description = callback_params
.get("error_description")
.map(String::as_str)
.unwrap_or("用户拒绝授权");
let sanitized_error =
sanitize_windsurf_browser_poll_callback_error(error, error_description);
session.status = "error".to_string();
session.error_msg = Some(sanitized_error.clone());
let _ = state
.save_provider_oauth_device_session(session_id, &session, 30)
.await;
return Ok(attach_admin_provider_oauth_device_poll_terminal_response(
session_id,
"error",
windsurf_browser_poll_error_response(sanitized_error),
));
}
let Some(callback_state) = callback_params
.get("state")
.map(String::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(windsurf_browser_poll_error_response("回调 URL 缺少 state"));
};
if callback_state != session_id {
return Ok(windsurf_browser_poll_error_response(
"回调 state 与会话不匹配",
));
}
if let Some(provider) = callback_params
.get("provider")
.or_else(|| callback_params.get("login_option"))
.map(String::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
social_provider = Some(provider.to_string());
}
let Some(callback_token) = callback_params
.get("token")
.or_else(|| callback_params.get("auth_token"))
.or_else(|| callback_params.get("access_token"))
.map(String::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(windsurf_browser_poll_error_response("回调 URL 缺少 token"));
};
callback_token.to_string()
} else {
let token = token.unwrap_or_default();
if windsurf_raw_api_key(token).is_none() {
return Ok(windsurf_browser_poll_error_response(
"浏览器授权请提交包含 state 的回调 URL;纯 token 请使用导入授权",
));
}
token.to_string()
};
let mut raw_credentials = serde_json::Map::new();
if windsurf_raw_api_key(&imported_token).is_some() {
raw_credentials.insert("api_key".to_string(), json!(imported_token));
} else {
raw_credentials.insert("token".to_string(), json!(imported_token));
}
if let Some(social_provider) = social_provider.as_ref() {
raw_credentials.insert("social_provider".to_string(), json!(social_provider));
}
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 result = match ProviderOAuthService::with_builtin_adapters()
.import_credentials(
&executor,
&ctx,
ProviderOAuthImportInput {
provider_type: provider.provider_type.clone(),
name: None,
refresh_token: None,
raw_credentials: Some(Value::Object(raw_credentials)),
network: ctx.network.clone(),
},
)
.await
{
Ok(result) => result,
Err(error) => {
let sanitized_error = sanitize_windsurf_browser_poll_oauth_error(&error);
session.status = "error".to_string();
session.error_msg = Some(sanitized_error.clone());
let _ = state
.save_provider_oauth_device_session(session_id, &session, 30)
.await;
return Ok(attach_admin_provider_oauth_device_poll_terminal_response(
session_id,
"error",
windsurf_browser_poll_error_response(sanitized_error),
));
}
};
let access_token = result.token_set.access_token.trim().to_string();
if access_token.is_empty() {
return Ok(windsurf_browser_poll_error_response(
"Windsurf token 验证返回缺少 apiKey/sessionToken",
));
}
let mut auth_config = result.auth_config.as_object().cloned().unwrap_or_default();
auth_config.insert("provider_type".to_string(), json!("windsurf"));
auth_config.insert("auth_method".to_string(), json!("browser"));
if let Some(social_provider) = social_provider.as_ref() {
auth_config
.entry("social_provider".to_string())
.or_insert_with(|| json!(social_provider));
}
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 = email
.as_deref()
.map(|email| format!("windsurf_{email}"))
.unwrap_or_else(|| format!("windsurf_{}", current_unix_secs()));
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 handle_admin_provider_oauth_kiro_social_device_poll(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
@@ -814,7 +1195,7 @@ async fn handle_admin_provider_oauth_kiro_social_device_poll(
.get("idToken")
.or_else(|| token_result.get("id_token")),
) {
auth_config_object.insert("id_token".to_string(), json!(id_token));
insert_secret_fingerprint(&mut auth_config_object, "id_token_fingerprint", &id_token);
}
if let Some(token_type) = json_non_empty_string(
token_result
@@ -921,3 +1302,18 @@ async fn handle_admin_provider_oauth_kiro_social_device_poll(
.into_response(),
))
}
#[cfg(test)]
mod tests {
#[test]
fn windsurf_browser_poll_callback_error_redacts_sensitive_values() {
let detail = super::sanitize_windsurf_browser_poll_callback_error(
"access_denied",
"bad token devin-session-token$secret and apiKey sk-secret",
);
assert_eq!(detail, "access_denied: [REDACTED upstream error body]");
assert!(!detail.contains("devin-session-token$secret"));
assert!(!detail.contains("sk-secret"));
}
}
@@ -12,6 +12,7 @@ pub(super) struct AdminProviderOAuthDeviceAuthorizePayload {
#[serde(default = "default_kiro_device_region")]
pub(super) region: String,
pub(super) auth_type: Option<String>,
pub(super) login_option: Option<String>,
pub(super) redirect_uri: Option<String>,
pub(super) proxy_node_id: Option<String>,
}
@@ -20,6 +21,7 @@ pub(super) struct AdminProviderOAuthDeviceAuthorizePayload {
pub(super) struct AdminProviderOAuthDevicePollPayload {
pub(super) session_id: String,
pub(super) callback_url: Option<String>,
pub(super) token: Option<String>,
}
pub(super) fn attach_admin_provider_oauth_device_poll_terminal_response(
@@ -25,6 +25,10 @@ use crate::handlers::admin::request::{
};
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
use aether_oauth::core::OAuthError;
use aether_oauth::provider::{
ProviderOAuthImportInput, ProviderOAuthService, ProviderOAuthTransportContext,
};
use axum::{
body::Body,
http,
@@ -39,6 +43,49 @@ struct AdminProviderOAuthSingleImportTokens {
expires_at: Option<u64>,
}
fn sanitize_windsurf_import_error(error: &OAuthError) -> String {
match error {
OAuthError::InvalidRequest(_) => "Windsurf 凭据验证失败: 请求参数无效".to_string(),
OAuthError::HttpStatus { status_code, .. } => {
format!("Windsurf 凭据验证失败: HTTP {status_code}")
}
OAuthError::InvalidResponse(detail) => sanitize_windsurf_invalid_response_detail(detail)
.unwrap_or_else(|| "Windsurf 凭据验证失败".to_string()),
_ => "Windsurf 凭据验证失败".to_string(),
}
}
fn sanitize_windsurf_invalid_response_detail(detail: &str) -> Option<String> {
let detail = detail.trim();
if detail.eq_ignore_ascii_case("Auth1 response is not json") {
return Some("Windsurf 凭据验证失败: Auth1 响应无法解析".to_string());
}
if detail.eq_ignore_ascii_case("Auth1 response missing token") {
return Some("Windsurf 凭据验证失败: Auth1 响应缺少 token".to_string());
}
if detail.contains("WindsurfPostAuth response missing sessionToken")
|| detail.contains("WindsurfPostAuth response is not json")
|| (detail.contains("WindsurfPostAuth failed") && detail.contains("missing sessionToken"))
{
return Some("Windsurf 凭据验证失败: PostAuth 未返回 sessionToken".to_string());
}
if detail.contains("WindsurfPostAuth failed") {
return Some("Windsurf 凭据验证失败: PostAuth 失败".to_string());
}
None
}
fn import_payload_has_windsurf_credentials(
payload: &serde_json::Map<String, serde_json::Value>,
) -> bool {
import_payload_string(payload, "api_key", "apiKey").is_some()
|| import_payload_string_any(payload, &["token", "auth_token", "authToken"]).is_some()
|| import_payload_string(payload, "refresh_token", "refreshToken").is_some()
|| import_payload_string(payload, "access_token", "accessToken").is_some()
|| (import_payload_string_any(payload, &["email"]).is_some()
&& import_payload_string_any(payload, &["password"]).is_some())
}
fn import_payload_string(
payload: &serde_json::Map<String, serde_json::Value>,
snake_case: &str,
@@ -266,6 +313,70 @@ async fn resolve_admin_provider_oauth_single_import_tokens(
})
}
async fn resolve_admin_provider_oauth_windsurf_single_import_tokens(
state: &AdminAppState<'_>,
provider_type: &str,
name: Option<String>,
raw_payload: &serde_json::Map<String, serde_json::Value>,
refresh_token: Option<&str>,
request_proxy: Option<ProxySnapshot>,
) -> Result<AdminProviderOAuthSingleImportTokens, Response<Body>> {
let ctx = ProviderOAuthTransportContext {
provider_id: String::new(),
provider_type: provider_type.to_string(),
endpoint_id: None,
key_id: None,
auth_type: Some("oauth".to_string()),
decrypted_api_key: None,
decrypted_auth_config: None,
provider_config: None,
endpoint_config: None,
key_config: None,
network: aether_oauth::network::OAuthNetworkContext::provider_operation(
request_proxy.clone(),
),
};
let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state);
let service = ProviderOAuthService::with_builtin_adapters();
let result = service
.import_credentials(
&executor,
&ctx,
ProviderOAuthImportInput {
provider_type: provider_type.to_string(),
name,
refresh_token: refresh_token.map(ToOwned::to_owned),
raw_credentials: Some(serde_json::Value::Object(raw_payload.clone())),
network: ctx.network.clone(),
},
)
.await
.map_err(|error| {
build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
sanitize_windsurf_import_error(&error),
)
})?;
let access_token = result.token_set.access_token.trim().to_string();
if access_token.is_empty() {
return Err(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"Windsurf 凭据验证返回缺少 apiKey/sessionToken",
));
}
let auth_config = result.auth_config.as_object().cloned().ok_or_else(|| {
build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"Windsurf 凭据验证返回缺少 auth_config",
)
})?;
Ok(AdminProviderOAuthSingleImportTokens {
access_token,
auth_config,
expires_at: result.token_set.expires_at_unix_secs,
})
}
pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
@@ -370,19 +481,47 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
.await;
let key_proxy = provider_oauth_key_proxy_value(proxy_node_id.as_deref());
let resolved_import = match resolve_admin_provider_oauth_single_import_tokens(
state,
template,
&provider_type,
refresh_token_input.as_deref(),
access_token_input.as_deref(),
imported_expires_at,
request_proxy.clone(),
)
.await
{
Ok(value) => value,
Err(response) => return Ok(response),
let resolved_import = if provider_type == "windsurf" {
if !import_payload_has_windsurf_credentials(&raw_payload) {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"Windsurf 凭据不能为空",
));
}
match resolve_admin_provider_oauth_windsurf_single_import_tokens(
state,
&provider_type,
name.clone(),
&raw_payload,
refresh_token_input.as_deref(),
request_proxy.clone(),
)
.await
{
Ok(value) => value,
Err(response) => return Ok(response),
}
} else {
if 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 不能为空",
));
}
match resolve_admin_provider_oauth_single_import_tokens(
state,
template,
&provider_type,
refresh_token_input.as_deref(),
access_token_input.as_deref(),
imported_expires_at,
request_proxy.clone(),
)
.await
{
Ok(value) => value,
Err(response) => return Ok(response),
}
};
let AdminProviderOAuthSingleImportTokens {
access_token,
@@ -480,3 +619,47 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
}))
.into_response())
}
#[cfg(test)]
mod tests {
use super::sanitize_windsurf_import_error;
use aether_oauth::core::OAuthError;
#[test]
fn windsurf_import_error_redacts_http_body() {
let error = OAuthError::HttpStatus {
status_code: 400,
body_excerpt: "token=secret-token password=secret-password".to_string(),
};
let detail = sanitize_windsurf_import_error(&error);
assert_eq!(detail, "Windsurf 凭据验证失败: HTTP 400");
assert!(!detail.contains("secret-token"));
assert!(!detail.contains("secret-password"));
}
#[test]
fn windsurf_import_error_redacts_invalid_response_detail() {
let error =
OAuthError::invalid_response("RegisterUser failed with firebase_id_token=secret-token");
let detail = sanitize_windsurf_import_error(&error);
assert_eq!(detail, "Windsurf 凭据验证失败");
assert!(!detail.contains("secret-token"));
assert!(!detail.contains("firebase_id_token"));
}
#[test]
fn windsurf_import_error_keeps_safe_post_auth_stage() {
let error = OAuthError::invalid_response("WindsurfPostAuth response missing sessionToken");
let detail = sanitize_windsurf_import_error(&error);
assert_eq!(
detail,
"Windsurf 凭据验证失败: PostAuth 未返回 sessionToken"
);
}
}
@@ -63,6 +63,12 @@ pub(super) async fn handle_admin_provider_oauth_start_key(
"该 Provider 不是固定类型,无法使用 provider-oauth",
));
}
if provider_type == "windsurf" {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"Windsurf 请使用浏览器登录或导入凭据。",
));
}
let Some(template) = admin_provider_oauth_template(&provider_type) else {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
@@ -134,6 +140,12 @@ pub(super) async fn handle_admin_provider_oauth_start_provider(
"Kiro 不支持 OAuth 授权,请使用导入授权。",
));
}
if provider_type == "windsurf" {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"Windsurf 请使用浏览器登录或导入凭据。",
));
}
let Some(template) = admin_provider_oauth_template(&provider_type) else {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
@@ -36,6 +36,13 @@ fn is_openai_provider_oauth_provider_type(value: Option<&serde_json::Value>) ->
})
}
fn is_windsurf_provider_oauth_provider_type(value: Option<&serde_json::Value>) -> bool {
value
.and_then(serde_json::Value::as_str)
.map(str::trim)
.is_some_and(|provider_type| provider_type.eq_ignore_ascii_case("windsurf"))
}
fn match_codex_provider_oauth_identity(
new_auth_config: &serde_json::Map<String, serde_json::Value>,
existing_auth_config: &serde_json::Map<String, serde_json::Value>,
@@ -106,6 +113,41 @@ fn match_codex_provider_oauth_identity(
None
}
fn match_windsurf_provider_oauth_identity(
new_auth_config: &serde_json::Map<String, serde_json::Value>,
existing_auth_config: &serde_json::Map<String, serde_json::Value>,
) -> Option<bool> {
let new_provider_type = new_auth_config.get("provider_type");
let existing_provider_type = existing_auth_config.get("provider_type");
if !is_windsurf_provider_oauth_provider_type(new_provider_type)
&& !is_windsurf_provider_oauth_provider_type(existing_provider_type)
{
return None;
}
let new_account_id = normalize_provider_oauth_identity_value(new_auth_config.get("account_id"));
let existing_account_id =
normalize_provider_oauth_identity_value(existing_auth_config.get("account_id"));
if let (Some(new_account_id), Some(existing_account_id)) =
(new_account_id.as_deref(), existing_account_id.as_deref())
{
return Some(new_account_id == existing_account_id);
}
let new_credential_fingerprint =
normalize_provider_oauth_identity_value(new_auth_config.get("credential_fingerprint"));
let existing_credential_fingerprint =
normalize_provider_oauth_identity_value(existing_auth_config.get("credential_fingerprint"));
if let (Some(new_fingerprint), Some(existing_fingerprint)) = (
new_credential_fingerprint.as_deref(),
existing_credential_fingerprint.as_deref(),
) {
return Some(new_fingerprint == existing_fingerprint);
}
None
}
fn is_codex_cross_plan_group_non_duplicate(
new_auth_config: &serde_json::Map<String, serde_json::Value>,
existing_auth_config: &serde_json::Map<String, serde_json::Value>,
@@ -169,10 +211,17 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
) -> Result<Option<StoredProviderCatalogKey>, String> {
let new_email = normalize_provider_oauth_identity_value(auth_config.get("email"));
let new_user_id = normalize_provider_oauth_identity_value(auth_config.get("user_id"));
let new_account_id = normalize_provider_oauth_identity_value(auth_config.get("account_id"));
let new_credential_fingerprint =
normalize_provider_oauth_identity_value(auth_config.get("credential_fingerprint"));
let new_auth_method = normalize_provider_oauth_identity_value(auth_config.get("auth_method"));
let new_kiro_provider = normalize_provider_oauth_identity_value(auth_config.get("provider"));
if new_email.is_none() && new_user_id.is_none() {
if new_email.is_none()
&& new_user_id.is_none()
&& new_account_id.is_none()
&& new_credential_fingerprint.is_none()
{
return Ok(None);
}
@@ -201,15 +250,28 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
normalize_provider_oauth_identity_value(existing_auth_config.get("auth_method"));
let existing_kiro_provider =
normalize_provider_oauth_identity_value(existing_auth_config.get("provider"));
let is_windsurf = auth_config
.get("provider_type")
.and_then(serde_json::Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("windsurf"))
|| existing_auth_config
.get("provider_type")
.and_then(serde_json::Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("windsurf"));
let mut is_duplicate = false;
let codex_identity_match =
match_codex_provider_oauth_identity(auth_config, &existing_auth_config);
let windsurf_identity_match =
match_windsurf_provider_oauth_identity(auth_config, &existing_auth_config);
if let Some(codex_identity_match) = codex_identity_match {
is_duplicate = codex_identity_match;
} else if let Some(windsurf_identity_match) = windsurf_identity_match {
is_duplicate = windsurf_identity_match;
}
if codex_identity_match.is_none()
&& windsurf_identity_match.is_none()
&& !is_duplicate
&& new_user_id.is_some()
&& existing_user_id.is_some()
@@ -220,7 +282,9 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
}
if codex_identity_match.is_none()
&& windsurf_identity_match.is_none()
&& !is_duplicate
&& !is_windsurf
&& new_email.is_some()
&& existing_email.is_some()
&& new_email == existing_email
@@ -261,6 +325,12 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
let identifier =
normalize_provider_oauth_identity_value(auth_config.get("account_user_id"))
.or_else(|| normalize_provider_oauth_identity_value(auth_config.get("account_id")))
.or_else(|| {
normalize_provider_oauth_identity_value(
auth_config.get("credential_fingerprint"),
)
.map(|value| format!("fingerprint:{value}"))
})
.or_else(|| new_email.clone())
.or_else(|| new_user_id.clone())
.unwrap_or_default();
@@ -272,3 +342,91 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
Ok(None)
}
#[cfg(test)]
mod tests {
use super::match_windsurf_provider_oauth_identity;
use serde_json::{json, Map, Value};
fn auth_config(value: Value) -> Map<String, Value> {
value.as_object().cloned().expect("auth config object")
}
#[test]
fn windsurf_identity_matches_account_id_without_email() {
let new_auth_config = auth_config(json!({
"provider_type": "windsurf",
"auth_method": "api_key",
"account_id": "acct-ws-1"
}));
let existing_auth_config = auth_config(json!({
"provider_type": "windsurf",
"auth_method": "browser",
"account_id": "acct-ws-1"
}));
assert_eq!(
match_windsurf_provider_oauth_identity(&new_auth_config, &existing_auth_config),
Some(true)
);
}
#[test]
fn windsurf_identity_rejects_different_account_id() {
let new_auth_config = auth_config(json!({
"provider_type": "windsurf",
"account_id": "acct-ws-1",
"email": "same@example.com"
}));
let existing_auth_config = auth_config(json!({
"provider_type": "windsurf",
"account_id": "acct-ws-2",
"email": "same@example.com"
}));
assert_eq!(
match_windsurf_provider_oauth_identity(&new_auth_config, &existing_auth_config),
Some(false)
);
}
#[test]
fn windsurf_identity_matches_credential_fingerprint_without_profile() {
let new_auth_config = auth_config(json!({
"provider_type": "windsurf",
"auth_method": "api_key",
"credential_fingerprint": "abcdef0123456789"
}));
let existing_auth_config = auth_config(json!({
"provider_type": "windsurf",
"auth_method": "browser",
"credential_fingerprint": "abcdef0123456789"
}));
assert_eq!(
match_windsurf_provider_oauth_identity(&new_auth_config, &existing_auth_config),
Some(true)
);
}
#[test]
fn windsurf_identity_does_not_match_user_supplied_email_only() {
let new_auth_config = auth_config(json!({
"provider_type": "windsurf",
"auth_method": "api_key",
"email": "same@example.com",
"email_verified": false
}));
let existing_auth_config = auth_config(json!({
"provider_type": "windsurf",
"auth_method": "api_key",
"email": "same@example.com",
"email_verified": false
}));
assert_eq!(
match_windsurf_provider_oauth_identity(&new_auth_config, &existing_auth_config),
None
);
}
}
@@ -7,6 +7,7 @@ use super::codex::refresh_codex_provider_quota_locally;
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 crate::handlers::admin::request::AdminAppState;
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
@@ -41,6 +42,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),
];
pub(crate) async fn refresh_provider_pool_quota_locally(
@@ -156,3 +158,19 @@ fn refresh_grok_provider_quota_locally_boxed<'a>(
proxy_override,
))
}
fn refresh_windsurf_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_windsurf_provider_quota_locally(
state,
provider,
endpoint,
keys,
proxy_override,
))
}
@@ -6,3 +6,4 @@ pub(crate) mod gemini_cli;
pub(crate) mod grok;
pub(crate) mod kiro;
pub(crate) mod shared;
pub(crate) mod windsurf;
@@ -0,0 +1,632 @@
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, persist_provider_quota_refresh_state,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_provider_pool::{
build_windsurf_pool_model_configs_request_with_base_url,
build_windsurf_pool_quota_request_with_base_url,
build_windsurf_pool_rate_limit_request_with_base_url, ProviderPoolQuotaRequestSpec,
};
use serde_json::json;
use std::time::{SystemTime, UNIX_EPOCH};
async fn execute_windsurf_probe_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
spec: ProviderPoolQuotaRequestSpec,
proxy_override: Option<&ProxySnapshot>,
quota_kind: &str,
) -> 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, quota_kind).await
}
async fn execute_windsurf_user_status_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
api_key: &str,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let spec = build_windsurf_pool_quota_request_with_base_url(
&transport.key.id,
&transport.endpoint.base_url,
api_key,
);
execute_windsurf_probe_plan(
state,
transport,
spec,
proxy_override,
"windsurf:user_status",
)
.await
}
async fn execute_windsurf_model_configs_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
api_key: &str,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let spec = build_windsurf_pool_model_configs_request_with_base_url(
&transport.key.id,
&transport.endpoint.base_url,
api_key,
);
execute_windsurf_probe_plan(
state,
transport,
spec,
proxy_override,
"windsurf:model_configs",
)
.await
}
async fn execute_windsurf_rate_limit_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
api_key: &str,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let spec = build_windsurf_pool_rate_limit_request_with_base_url(
&transport.key.id,
&transport.endpoint.base_url,
api_key,
);
execute_windsurf_probe_plan(
state,
transport,
spec,
proxy_override,
"windsurf:rate_limit",
)
.await
}
fn merge_windsurf_probe_metadata(
mut user_status_metadata: serde_json::Value,
model_configs_metadata: Option<serde_json::Value>,
rate_limit_metadata: Option<serde_json::Value>,
) -> serde_json::Value {
let Some(target) = user_status_metadata.as_object_mut() else {
return user_status_metadata;
};
for metadata in [model_configs_metadata, rate_limit_metadata]
.into_iter()
.flatten()
{
if let Some(source) = metadata.as_object() {
for (key, value) in source {
target.insert(key.clone(), value.clone());
}
}
}
user_status_metadata
}
fn append_windsurf_probe_warning(metadata: &mut serde_json::Value, probe: &str, message: String) {
let Some(target) = metadata.as_object_mut() else {
return;
};
let warnings = target
.entry("probe_warnings".to_string())
.or_insert_with(|| serde_json::Value::Array(Vec::new()));
if let Some(items) = warnings.as_array_mut() {
items.push(json!({
"probe": probe,
"message": message,
}));
}
}
fn build_windsurf_metadata_update(
current_upstream_metadata: Option<&serde_json::Value>,
patch: serde_json::Value,
) -> serde_json::Value {
let Some(patch_object) = patch.as_object() else {
return json!({ "windsurf": patch });
};
let mut merged_bucket = current_upstream_metadata
.and_then(|value| value.get("windsurf"))
.and_then(serde_json::Value::as_object)
.cloned()
.unwrap_or_default();
for (key, value) in patch_object {
merged_bucket.insert(key.clone(), value.clone());
}
json!({ "windsurf": merged_bucket })
}
fn sanitize_windsurf_probe_detail(detail: impl AsRef<str>) -> String {
let detail = detail.as_ref().trim();
if detail.is_empty() {
return "-".to_string();
}
if let Ok(mut value) = serde_json::from_str::<serde_json::Value>(detail) {
redact_windsurf_sensitive_json(&mut value);
return value.to_string().chars().take(500).collect();
}
if contains_windsurf_sensitive_marker(detail) {
"[REDACTED upstream error body]".to_string()
} else {
detail.chars().take(500).collect()
}
}
fn redact_windsurf_sensitive_json(value: &mut serde_json::Value) {
match value {
serde_json::Value::Object(object) => {
for (key, value) in object {
if is_windsurf_sensitive_key(key) {
*value = json!("[REDACTED]");
} else {
redact_windsurf_sensitive_json(value);
}
}
}
serde_json::Value::Array(items) => {
for item in items {
redact_windsurf_sensitive_json(item);
}
}
serde_json::Value::String(text) if looks_like_windsurf_secret(text) => {
*text = "[REDACTED]".to_string();
}
_ => {}
}
}
fn is_windsurf_sensitive_key(key: &str) -> bool {
let normalized = key
.chars()
.filter(|ch| ch.is_ascii_alphanumeric())
.collect::<String>()
.to_ascii_lowercase();
normalized.contains("token")
|| normalized.contains("apikey")
|| normalized.contains("password")
|| normalized.contains("authorization")
|| normalized.contains("secret")
}
fn looks_like_windsurf_secret(value: &str) -> bool {
let value = value.trim();
value.starts_with("devin-session-token$")
|| value.starts_with("sk-")
|| (value.len() > 80 && value.split('.').count() == 3)
}
fn contains_windsurf_sensitive_marker(value: &str) -> bool {
let lowered = value.to_ascii_lowercase();
[
"token",
"api_key",
"apikey",
"sessiontoken",
"firebase_id_token",
"idtoken",
"authorization",
"password",
"secret",
"devin-session-token$",
]
.iter()
.any(|marker| lowered.contains(marker))
|| value.contains("sk-")
}
pub(crate) async fn refresh_windsurf_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;
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;
}
};
let api_key = transport.key.decrypted_api_key.trim();
if api_key.is_empty() {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "缺少 Windsurf apiKey/sessionToken",
}));
continue;
}
let result = match execute_windsurf_user_status_plan(
state,
&transport,
api_key,
proxy_override.as_ref(),
)
.await?
{
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
failed_count += 1;
let detail = sanitize_windsurf_probe_detail(detail);
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("GetUserStatus 请求执行失败: {detail}"),
"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())
{
let mut windsurf_metadata =
aether_admin::provider::quota::parse_windsurf_user_status_response(
body_json,
now_unix_secs,
);
if let Some(mut metadata) = windsurf_metadata.take() {
let model_metadata = match execute_windsurf_model_configs_plan(
state,
&transport,
api_key,
proxy_override.as_ref(),
)
.await?
{
ProviderQuotaExecutionOutcome::Response(model_result)
if model_result.status_code == 200 =>
{
model_result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
.and_then(|body_json| {
aether_admin::provider::quota::parse_windsurf_model_configs_response(
body_json,
now_unix_secs,
)
})
}
ProviderQuotaExecutionOutcome::Response(model_result) => {
let detail = extract_execution_error_message(&model_result)
.unwrap_or_else(|| format!("HTTP {}", model_result.status_code));
let detail = sanitize_windsurf_probe_detail(detail);
append_windsurf_probe_warning(
&mut metadata,
"model_configs",
format!("GetCascadeModelConfigs 返回: {detail}"),
);
None
}
ProviderQuotaExecutionOutcome::Failure(detail) => {
let detail = sanitize_windsurf_probe_detail(detail);
append_windsurf_probe_warning(
&mut metadata,
"model_configs",
format!("GetCascadeModelConfigs 执行失败: {detail}"),
);
None
}
};
let rate_limit_metadata = match execute_windsurf_rate_limit_plan(
state,
&transport,
api_key,
proxy_override.as_ref(),
)
.await?
{
ProviderQuotaExecutionOutcome::Response(rate_limit_result)
if rate_limit_result.status_code == 200 =>
{
rate_limit_result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
.and_then(|body_json| {
aether_admin::provider::quota::parse_windsurf_rate_limit_response(
body_json,
now_unix_secs,
)
})
}
ProviderQuotaExecutionOutcome::Response(rate_limit_result) => {
let detail = extract_execution_error_message(&rate_limit_result)
.unwrap_or_else(|| format!("HTTP {}", rate_limit_result.status_code));
let detail = sanitize_windsurf_probe_detail(detail);
append_windsurf_probe_warning(
&mut metadata,
"rate_limit",
format!("CheckUserMessageRateLimit 返回: {detail}"),
);
None
}
ProviderQuotaExecutionOutcome::Failure(detail) => {
let detail = sanitize_windsurf_probe_detail(detail);
append_windsurf_probe_warning(
&mut metadata,
"rate_limit",
format!("CheckUserMessageRateLimit 执行失败: {detail}"),
);
None
}
};
metadata = merge_windsurf_probe_metadata(
metadata,
model_metadata,
rate_limit_metadata,
);
metadata_update = Some(build_windsurf_metadata_update(
key.upstream_metadata.as_ref(),
metadata,
));
status = "success".to_string();
} else {
status = "no_metadata".to_string();
message = Some("响应中未包含 Windsurf 限额信息".to_string());
}
} else {
status = "no_metadata".to_string();
message = Some("无法解析 GetUserStatus 响应".to_string());
}
} else {
let err_msg =
extract_execution_error_message(&result).map(sanitize_windsurf_probe_detail);
message = Some(match err_msg.as_deref() {
Some(detail) if !detail.is_empty() => {
format!(
"GetUserStatus 返回状态码 {}: {}",
result.status_code, detail
)
}
_ => format!("GetUserStatus 返回状态码 {}", result.status_code),
});
let detail = err_msg
.clone()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| format!("HTTP {}", result.status_code));
let mut metadata = serde_json::Map::new();
metadata.insert("updated_at".to_string(), json!(now_unix_secs));
metadata.insert("last_error".to_string(), json!(detail));
match result.status_code {
401 | 403 => {
oauth_invalid_at_unix_secs = Some(now_unix_secs);
oauth_invalid_reason =
Some(format!("Windsurf token 无效或已被拒绝: {}", detail));
metadata.insert("banned".to_string(), json!(result.status_code == 403));
status = if result.status_code == 401 {
"auth_invalid".to_string()
} else {
"forbidden".to_string()
};
}
429 => {
metadata.insert(
"rate_limit".to_string(),
json!({
"limited": true,
"message": metadata
.get("last_error")
.cloned()
.unwrap_or_else(|| json!("rate limited")),
}),
);
status = "rate_limited".to_string();
}
_ => {}
}
metadata_update = Some(build_windsurf_metadata_update(
key.upstream_metadata.as_ref(),
serde_json::Value::Object(metadata),
));
}
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 result.status_code != 200 {
payload.insert("status_code".to_string(), json!(result.status_code));
}
if let Some(metadata) = metadata_update
.as_ref()
.and_then(|value| value.get("windsurf"))
.cloned()
{
payload.insert("metadata".to_string(), metadata);
}
if let Some(quota_snapshot) = build_quota_snapshot_payload(
"windsurf",
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": success_count + failed_count,
"results": results,
"message": format!("已处理 {} 个 Key", success_count + failed_count),
"auto_removed": 0,
})))
}
#[cfg(test)]
mod tests {
use serde_json::json;
#[test]
fn windsurf_probe_metadata_merges_user_status_models_and_rate_limit() {
let metadata = super::merge_windsurf_probe_metadata(
json!({
"plan_name": "Pro",
"daily_remaining_percent": 42.0,
"updated_at": 1_770_000_000u64,
}),
Some(json!({
"allowed_models_count": 2u64,
"models": [
{"model_uid": "claude-sonnet-4-5"},
{"model_uid": "gpt-5-mini"}
],
"updated_at": 1_770_000_010u64,
})),
Some(json!({
"rate_limit": {
"limited": true,
"messages_remaining": 0.0,
"retry_after_ms": 60_000u64
},
"updated_at": 1_770_000_020u64,
})),
);
assert_eq!(metadata["plan_name"], json!("Pro"));
assert_eq!(metadata["daily_remaining_percent"], json!(42.0));
assert_eq!(metadata["allowed_models_count"], json!(2u64));
assert_eq!(metadata["rate_limit"]["limited"], json!(true));
assert_eq!(metadata["updated_at"], json!(1_770_000_020u64));
}
#[test]
fn windsurf_probe_detail_redacts_sensitive_values() {
let detail = super::sanitize_windsurf_probe_detail(
r#"{"error":{"message":"bad"},"apiKey":"sk-secret","sessionToken":"devin-session-token$secret"}"#,
);
assert!(detail.contains("[REDACTED]"));
assert!(!detail.contains("sk-secret"));
assert!(!detail.contains("devin-session-token$secret"));
}
#[test]
fn windsurf_metadata_update_preserves_existing_bucket_fields() {
let update = super::build_windsurf_metadata_update(
Some(&json!({
"windsurf": {
"daily_remaining_percent": 0.0,
"allowed_models_count": 3,
"updated_at": 1u64
}
})),
json!({
"last_error": "HTTP 429",
"rate_limit": {"limited": true},
"updated_at": 2u64
}),
);
assert_eq!(
update.pointer("/windsurf/daily_remaining_percent"),
Some(&json!(0.0))
);
assert_eq!(
update.pointer("/windsurf/allowed_models_count"),
Some(&json!(3))
);
assert_eq!(update.pointer("/windsurf/updated_at"), Some(&json!(2u64)));
assert_eq!(
update.pointer("/windsurf/rate_limit/limited"),
Some(&json!(true))
);
}
}
@@ -70,7 +70,7 @@ pub(super) async fn read_admin_provider_ops_balance_cache(
}
}
pub(super) async fn store_admin_provider_ops_balance_cache(
pub(crate) async fn store_admin_provider_ops_balance_cache(
state: &AdminAppState<'_>,
provider_id: &str,
payload: &Value,
@@ -1,4 +1,7 @@
use super::support::{AdminProviderOpsSaveConfigRequest, ADMIN_PROVIDER_OPS_SENSITIVE_FIELDS};
use super::support::{
AdminProviderOpsQuotaAlertConfigRequest, AdminProviderOpsSaveConfigRequest,
ADMIN_PROVIDER_OPS_SENSITIVE_FIELDS,
};
use crate::handlers::admin::request::AdminAppState;
use crate::GatewayError;
use aether_admin::provider::ops as admin_provider_ops_pure;
@@ -8,6 +11,10 @@ use aether_data_contracts::repository::provider_catalog::{
use serde_json::json;
use std::time::{SystemTime, UNIX_EPOCH};
const PROVIDER_OPS_QUOTA_ALERT_DEFAULT_FETCH_INTERVAL_SECS: u64 = 30;
const PROVIDER_OPS_QUOTA_ALERT_MIN_FETCH_INTERVAL_SECS: u64 = 30;
const PROVIDER_OPS_QUOTA_ALERT_MAX_FETCH_INTERVAL_SECS: u64 = 86_400;
pub(super) fn admin_provider_ops_config_object(
provider: &StoredProviderCatalogProvider,
) -> Option<&serde_json::Map<String, serde_json::Value>> {
@@ -276,6 +283,7 @@ pub(super) fn build_admin_provider_ops_saved_config_value(
)
})
.collect::<serde_json::Map<String, serde_json::Value>>();
let quota_alert = normalize_admin_provider_ops_quota_alert(payload.quota_alert)?;
Ok(json!({
"architecture_id": payload.architecture_id,
@@ -287,9 +295,48 @@ pub(super) fn build_admin_provider_ops_saved_config_value(
},
"actions": actions,
"schedule": payload.schedule,
"quota_alert": quota_alert,
}))
}
fn normalize_admin_provider_ops_quota_alert(
request: Option<AdminProviderOpsQuotaAlertConfigRequest>,
) -> Result<serde_json::Value, String> {
let Some(request) = request else {
return Ok(default_admin_provider_ops_quota_alert());
};
let threshold_amount = request.threshold_amount.unwrap_or(0.0);
if threshold_amount < 0.0 {
return Err("quota_alert.threshold_amount 必须大于等于 0".to_string());
}
let fetch_interval_seconds = request
.fetch_interval_seconds
.unwrap_or(PROVIDER_OPS_QUOTA_ALERT_DEFAULT_FETCH_INTERVAL_SECS);
if !(PROVIDER_OPS_QUOTA_ALERT_MIN_FETCH_INTERVAL_SECS
..=PROVIDER_OPS_QUOTA_ALERT_MAX_FETCH_INTERVAL_SECS)
.contains(&fetch_interval_seconds)
{
return Err(format!(
"quota_alert.fetch_interval_seconds 必须在 {} 到 {} 秒之间",
PROVIDER_OPS_QUOTA_ALERT_MIN_FETCH_INTERVAL_SECS,
PROVIDER_OPS_QUOTA_ALERT_MAX_FETCH_INTERVAL_SECS
));
}
Ok(json!({
"enabled": request.enabled,
"threshold_amount": threshold_amount,
"fetch_interval_seconds": fetch_interval_seconds,
}))
}
fn default_admin_provider_ops_quota_alert() -> serde_json::Value {
json!({
"enabled": false,
"threshold_amount": 0.0,
"fetch_interval_seconds": PROVIDER_OPS_QUOTA_ALERT_DEFAULT_FETCH_INTERVAL_SECS,
})
}
pub(super) fn resolve_admin_provider_ops_base_url(
provider: &StoredProviderCatalogProvider,
endpoints: &[StoredProviderCatalogEndpoint],
@@ -356,5 +403,10 @@ pub(super) fn build_admin_provider_ops_config_payload(
connector.and_then(|connector| connector.get("credentials")),
),
},
"quota_alert": provider_ops_config
.get("quota_alert")
.filter(|value| value.is_object())
.cloned()
.unwrap_or_else(default_admin_provider_ops_quota_alert),
})
}
@@ -4,4 +4,5 @@ mod config;
mod routes;
mod support;
mod verify;
pub(crate) use self::balance_cache::store_admin_provider_ops_balance_cache;
pub(super) use self::routes::maybe_build_local_admin_provider_ops_providers_response;
@@ -33,6 +33,8 @@ pub(super) struct AdminProviderOpsSaveConfigRequest {
pub(crate) actions: BTreeMap<String, AdminProviderOpsActionConfigRequest>,
#[serde(default)]
pub(crate) schedule: BTreeMap<String, String>,
#[serde(default)]
pub(crate) quota_alert: Option<AdminProviderOpsQuotaAlertConfigRequest>,
}
#[derive(Debug, Deserialize)]
@@ -52,6 +54,16 @@ pub(super) struct AdminProviderOpsActionConfigRequest {
pub(crate) config: serde_json::Map<String, serde_json::Value>,
}
#[derive(Debug, Deserialize)]
pub(super) struct AdminProviderOpsQuotaAlertConfigRequest {
#[serde(default)]
pub(crate) enabled: bool,
#[serde(default, deserialize_with = "deserialize_optional_f64_from_number")]
pub(crate) threshold_amount: Option<f64>,
#[serde(default)]
pub(crate) fetch_interval_seconds: Option<u64>,
}
#[derive(Debug, Deserialize)]
pub(super) struct AdminProviderOpsConnectRequest {
#[serde(default)]
@@ -71,3 +83,28 @@ fn default_admin_provider_ops_architecture_id() -> String {
fn default_admin_provider_ops_action_enabled() -> bool {
true
}
fn deserialize_optional_f64_from_number<'de, D>(deserializer: D) -> Result<Option<f64>, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = Option::<serde_json::Value>::deserialize(deserializer)?;
match value {
None | Some(serde_json::Value::Null) => Ok(None),
Some(serde_json::Value::Number(number)) => number
.as_f64()
.filter(|value| value.is_finite())
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected a finite number")),
Some(serde_json::Value::String(raw)) => raw
.trim()
.parse::<f64>()
.ok()
.filter(|value| value.is_finite())
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected a finite number or numeric string")),
Some(_) => Err(serde::de::Error::custom(
"expected a finite number or numeric string",
)),
}
}
@@ -14,7 +14,7 @@ use super::{provider_query_key_display_name, provider_query_provider_payload};
use crate::ai_serving::{
maybe_build_sync_finalize_outcome, GatewayControlDecision,
ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND,
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
};
use crate::clock::current_unix_ms;
use crate::execution_runtime;
@@ -1792,6 +1792,56 @@ async fn provider_query_execute_kiro_test_candidate(
})
}
async fn provider_query_finalize_windsurf_result(
route_path: &str,
trace_id: &str,
requested_model: &str,
mapped_model: &str,
original_request_body: &Value,
result: &aether_contracts::ExecutionResult,
) -> Result<Option<Value>, GatewayError> {
let decision = GatewayControlDecision::synthetic(
route_path,
Some("admin_proxy".to_string()),
Some("provider_query_manage".to_string()),
Some("test_model_failover".to_string()),
Some("openai:chat".to_string()),
);
let payload = GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind: OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND.to_string(),
report_context: Some(json!({
"client_api_format": "openai:chat",
"provider_api_format": "openai:chat",
"model": requested_model,
"mapped_model": mapped_model,
"needs_conversion": false,
"has_envelope": true,
"envelope_name": crate::provider_transport::windsurf::WINDSURF_ENVELOPE_NAME,
"original_request_body": original_request_body,
})),
status_code: result.status_code,
headers: result.headers.clone(),
body_json: result.body.as_ref().and_then(|body| body.json_body.clone()),
client_body_json: None,
body_base64: result
.body
.as_ref()
.and_then(|body| body.body_bytes_b64.clone()),
telemetry: result.telemetry.clone(),
};
let Some(outcome) = maybe_build_sync_finalize_outcome(trace_id, &decision, &payload)? else {
return Ok(None);
};
let bytes = to_bytes(outcome.response.into_body(), usize::MAX)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
serde_json::from_slice::<Value>(&bytes)
.map(Some)
.map_err(|err| GatewayError::Internal(err.to_string()))
}
fn provider_query_build_openai_image_test_request_body_for_route(
payload: &Value,
model: &str,
@@ -2665,6 +2715,22 @@ async fn provider_query_execute_standard_test_candidate(
route_path,
client_api_format,
);
if crate::provider_transport::is_windsurf_provider_transport(&transport)
&& provider_query_normalize_api_format_alias(candidate.endpoint.api_format.as_str())
== "openai:chat"
{
return provider_query_execute_windsurf_test_candidate(
state,
provider,
candidate,
payload,
route_path,
trace_id,
transport,
original_request_body,
)
.await;
}
if !provider_query_transport_supports_model_test_execution(
state,
&transport,
@@ -3133,6 +3199,182 @@ async fn provider_query_execute_standard_test_candidate(
})
}
#[allow(clippy::too_many_arguments)]
async fn provider_query_execute_windsurf_test_candidate(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
candidate: &ProviderQueryTestCandidate,
payload: &Value,
route_path: &str,
trace_id: &str,
transport: AdminGatewayProviderTransportSnapshot,
original_request_body: Value,
) -> Result<ProviderQueryExecutionOutcome, GatewayError> {
if let Some(_reason) =
crate::provider_transport::local_windsurf_request_transport_unsupported_reason_with_network(
&transport,
)
{
return Ok(provider_query_skipped_execution_outcome(
original_request_body,
provider_query_standard_test_unsupported_reason(
&transport,
candidate.endpoint.api_format.as_str(),
),
));
}
let incoming_request_headers = provider_query_extract_request_headers(payload);
let request_body = original_request_body.clone();
let request_model =
provider_query_request_body_model(&request_body, &candidate.effective_model);
let client_is_stream = request_body
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(false);
let hard_requires_streaming = crate::ai_serving::force_upstream_streaming_for_provider(
transport.provider.provider_type.as_str(),
candidate.endpoint.api_format.as_str(),
);
let upstream_is_stream = crate::ai_serving::resolve_upstream_is_stream_from_endpoint_config(
transport.endpoint.config.as_ref(),
client_is_stream,
hard_requires_streaming,
);
let Some((auth_header, auth_value)) =
crate::provider_transport::windsurf::resolve_windsurf_cascade_auth(&transport).or_else(
|| crate::provider_transport::auth::resolve_local_openai_bearer_auth(&transport),
)
else {
return Ok(provider_query_skipped_execution_outcome(
request_body,
"Provider auth is unavailable for windsurf".to_string(),
));
};
let mut synthetic_request = http::Request::builder()
.uri(route_path)
.body(())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
*synthetic_request.headers_mut() = incoming_request_headers;
let (parts, _) = synthetic_request.into_parts();
let Some(provider_request_body) =
crate::provider_transport::build_windsurf_cascade_request_body(
&request_body,
request_model,
&auth_value,
transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
upstream_is_stream,
)
else {
return Ok(provider_query_skipped_execution_outcome(
request_body,
"Provider request body could not be built for windsurf".to_string(),
));
};
let Some(request_url) = crate::provider_transport::build_windsurf_cascade_upstream_url(
transport.endpoint.base_url.as_str(),
parts.uri.query(),
) else {
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
"Provider request URL is unavailable for windsurf".to_string(),
));
};
let Some(request_headers) = crate::provider_transport::build_windsurf_cascade_headers(
&parts.headers,
&provider_request_body,
&request_body,
transport.endpoint.header_rules.as_ref(),
&auth_header,
&auth_value,
upstream_is_stream,
) else {
return Ok(ProviderQueryExecutionOutcome {
status: "failed",
skip_reason: None,
error_message: Some("provider request headers build failed".to_string()),
status_code: None,
latency_ms: None,
request_url,
request_headers: BTreeMap::new(),
request_body: provider_request_body,
response_headers: BTreeMap::new(),
response_body: None,
});
};
let plan = ExecutionPlan {
request_id: trace_id.to_string(),
candidate_id: Some(format!("provider-query-{}", candidate.key.id)),
provider_name: Some(provider.name.clone()),
provider_id: provider.id.clone(),
endpoint_id: candidate.endpoint.id.clone(),
key_id: candidate.key.id.clone(),
method: "POST".to_string(),
url: request_url.clone(),
headers: request_headers.clone(),
content_type: Some("application/connect+json".to_string()),
content_encoding: None,
body: RequestBody::from_json(provider_request_body.clone()),
stream: upstream_is_stream,
client_api_format: "openai:chat".to_string(),
provider_api_format: candidate.endpoint.api_format.clone(),
model_name: Some(request_model.to_string()),
proxy: state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
.await,
transport_profile: state.resolve_transport_profile(&transport),
timeouts: state.resolve_transport_execution_timeouts(&transport),
};
let result = state
.execute_execution_runtime_sync_plan(Some(trace_id), &plan)
.await?;
let response_body = if result.status_code < 400 {
provider_query_finalize_windsurf_result(
route_path,
trace_id,
request_model,
request_model,
&request_body,
&result,
)
.await?
} else {
result.body.as_ref().and_then(|body| body.json_body.clone())
};
let missing_success_body = result.status_code < 400 && response_body.is_none();
let did_fail = result.status_code >= 400 || missing_success_body;
let error_message = if did_fail {
provider_query_extract_error_message(&result).or_else(|| {
missing_success_body.then(|| {
format!(
"Provider returned HTTP {} without a model-test response body",
result.status_code
)
})
})
} else {
None
};
Ok(ProviderQueryExecutionOutcome {
status: if did_fail { "failed" } else { "success" },
skip_reason: None,
error_message,
status_code: Some(result.status_code),
latency_ms: result.telemetry.as_ref().and_then(|value| value.elapsed_ms),
request_url,
request_headers,
request_body: provider_request_body,
response_headers: result.headers,
response_body,
})
}
async fn build_admin_provider_query_kiro_failover_response(
state: &AdminAppState<'_>,
payload: &Value,
@@ -46,6 +46,22 @@ pub(super) fn provider_query_standard_test_unsupported_reason(
api_format: &str,
) -> String {
let normalized_api_format = crate::ai_serving::normalize_api_format_alias(api_format);
if crate::provider_transport::is_windsurf_provider_transport(transport)
&& normalized_api_format == "openai:chat"
{
let reason =
crate::provider_transport::local_windsurf_request_transport_unsupported_reason_with_network(
transport,
);
return match reason {
Some(reason) => format!(
"{} ({reason})",
provider_query_unsupported_test_api_format_message(api_format)
),
None => provider_query_unsupported_test_api_format_message(api_format),
};
}
let reason = match normalized_api_format.as_str() {
"openai:chat" => {
crate::provider_transport::policy::local_openai_chat_transport_unsupported_reason(
@@ -294,6 +310,15 @@ pub(super) fn provider_query_transport_supports_model_test_execution(
transport: &AdminGatewayProviderTransportSnapshot,
api_format: &str,
) -> bool {
if crate::provider_transport::is_windsurf_provider_transport(transport)
&& provider_query_normalize_api_format_alias(api_format) == "openai:chat"
{
return crate::provider_transport::local_windsurf_request_transport_unsupported_reason_with_network(
transport,
)
.is_none();
}
match provider_query_test_adapter_for_provider_api_format(
transport.provider.provider_type.as_str(),
api_format,
@@ -42,9 +42,9 @@ pub(super) fn provider_query_test_attempt_payload(
"status_code": execution.status_code,
"latency_ms": execution.latency_ms,
"request_url": execution.request_url,
"request_headers": provider_query_redact_diagnostic_headers(&execution.request_headers),
"request_body": execution.request_body,
"response_headers": provider_query_redact_diagnostic_headers(&execution.response_headers),
"request_headers": redacted_provider_query_headers(&execution.request_headers),
"request_body": redacted_provider_query_value(&execution.request_body),
"response_headers": redacted_provider_query_headers(&execution.response_headers),
"response_body": execution.response_body,
})
}
@@ -172,34 +172,84 @@ fn provider_query_endpoint_route_payload(
})
}
fn provider_query_redact_diagnostic_headers(
headers: &BTreeMap<String, String>,
) -> BTreeMap<String, String> {
fn redacted_provider_query_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
headers
.iter()
.map(|(name, value)| {
if provider_query_header_is_sensitive(name) {
(name.clone(), "<redacted>".to_string())
.map(|(key, value)| {
if provider_query_field_is_sensitive(key) {
(key.clone(), "[REDACTED]".to_string())
} else {
(name.clone(), value.clone())
(key.clone(), value.clone())
}
})
.collect()
}
fn provider_query_header_is_sensitive(name: &str) -> bool {
fn redacted_provider_query_value(value: &Value) -> Value {
match value {
Value::Object(object) => Value::Object(
object
.iter()
.map(|(key, value)| {
if provider_query_field_is_sensitive(key) {
(key.clone(), Value::String("[REDACTED]".to_string()))
} else {
(key.clone(), redacted_provider_query_value(value))
}
})
.collect(),
),
Value::Array(items) => Value::Array(
items
.iter()
.map(redacted_provider_query_value)
.collect::<Vec<_>>(),
),
other => other.clone(),
}
}
fn provider_query_field_is_sensitive(key: &str) -> bool {
let key = key.trim().to_ascii_lowercase();
let normalized = key
.chars()
.filter(|ch| ch.is_ascii_alphanumeric())
.collect::<String>();
if matches!(
normalized.as_str(),
"maxtokens"
| "maxoutputtokens"
| "inputtokens"
| "outputtokens"
| "prompttokens"
| "completiontokens"
| "totaltokens"
) {
return false;
}
matches!(
name.trim().to_ascii_lowercase().as_str(),
key.as_str(),
"authorization"
| "proxy-authorization"
| "cookie"
| "set-cookie"
| "x-api-key"
| "api_key"
| "apikey"
| "api-key"
| "x-api-key"
| "x-goog-api-key"
| "anthropic-api-key"
| "openai-api-key"
)
| "x-codeium-csrf-token"
| "access_token"
| "refresh_token"
| "id_token"
| "password"
| "secret"
) || normalized.ends_with("token")
|| normalized.contains("secret")
|| normalized.contains("apikey")
|| normalized.contains("authorization")
}
pub(super) fn provider_query_candidate_summary_payload(
@@ -309,34 +359,93 @@ pub(super) fn provider_query_candidate_summary_payload(
#[cfg(test)]
mod tests {
use super::*;
use super::{redacted_provider_query_headers, redacted_provider_query_value};
use serde_json::json;
use std::collections::BTreeMap;
#[test]
fn provider_query_diagnostic_headers_redact_credentials() {
fn redacts_sensitive_provider_query_headers() {
let headers = BTreeMap::from([
("cookie".to_string(), "sso=secret".to_string()),
("authorization".to_string(), "Bearer secret".to_string()),
(
"authorization".to_string(),
"Bearer secret-token".to_string(),
),
("x-goog-api-key".to_string(), "secret".to_string()),
("content-type".to_string(), "application/json".to_string()),
(
"x-codeium-csrf-token".to_string(),
"csrf-secret".to_string(),
),
]);
let redacted = provider_query_redact_diagnostic_headers(&headers);
let redacted = redacted_provider_query_headers(&headers);
assert_eq!(
redacted.get("cookie").map(String::as_str),
Some("<redacted>")
Some("[REDACTED]")
);
assert_eq!(
redacted.get("authorization").map(String::as_str),
Some("<redacted>")
Some("[REDACTED]")
);
assert_eq!(
redacted.get("x-goog-api-key").map(String::as_str),
Some("<redacted>")
Some("[REDACTED]")
);
assert_eq!(
redacted.get("x-codeium-csrf-token").map(String::as_str),
Some("[REDACTED]")
);
assert_eq!(
redacted.get("content-type").map(String::as_str),
Some("application/json")
);
}
#[test]
fn redacts_sensitive_provider_query_request_body_fields() {
let body = json!({
"metadata": {
"apiKey": "devin-session-token$secret",
"ideName": "windsurf"
},
"messages": [{"role": "user", "content": "hello"}],
"stream": true
});
let redacted = redacted_provider_query_value(&body);
assert_eq!(
redacted.pointer("/metadata/apiKey"),
Some(&json!("[REDACTED]"))
);
assert_eq!(
redacted.pointer("/metadata/ideName"),
Some(&json!("windsurf"))
);
assert_eq!(redacted.pointer("/stream"), Some(&json!(true)));
}
#[test]
fn keeps_non_secret_token_count_fields_visible() {
let body = json!({
"maxTokens": 64,
"usage": {
"inputTokens": 10,
"outputTokens": 2,
"accessToken": "secret"
}
});
let redacted = redacted_provider_query_value(&body);
assert_eq!(redacted.pointer("/maxTokens"), Some(&json!(64)));
assert_eq!(redacted.pointer("/usage/inputTokens"), Some(&json!(10)));
assert_eq!(redacted.pointer("/usage/outputTokens"), Some(&json!(2)));
assert_eq!(
redacted.pointer("/usage/accessToken"),
Some(&json!("[REDACTED]"))
);
}
}
@@ -138,6 +138,13 @@ pub(crate) fn build_admin_provider_summary_value(
.and_then(|cfg| cfg.get("simulated_cache_enabled"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let ops_quota_alert_enabled = provider_ops_config
.and_then(serde_json::Value::as_object)
.and_then(|cfg| cfg.get("quota_alert"))
.and_then(serde_json::Value::as_object)
.and_then(|cfg| cfg.get("enabled"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let billing_type = quota_snapshot
.map(|quota| quota.billing_type.clone())
.or_else(|| provider.billing_type.clone());
@@ -197,6 +204,7 @@ pub(crate) fn build_admin_provider_summary_value(
"ops_configured": ops_configured,
"ops_architecture_id": ops_architecture_id,
"kiro_simulated_cache_enabled": kiro_simulated_cache_enabled,
"ops_quota_alert_enabled": ops_quota_alert_enabled,
"created_at": endpoint_timestamp_or_now(provider.created_at_unix_ms, now_unix_secs),
"updated_at": endpoint_timestamp_or_now(provider.updated_at_unix_secs, now_unix_secs),
})
@@ -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" => Ok(normalized),
| "antigravity" | "vertex_ai" | "grok" | "windsurf" => Ok(normalized),
_ => Err(
"provider_type 仅支持 custom / claude_code / kiro / codex / chatgpt_web / gemini_cli / antigravity / vertex_ai / grok"
"provider_type 仅支持 custom / claude_code / kiro / codex / chatgpt_web / gemini_cli / antigravity / vertex_ai / grok / windsurf"
.to_string(),
),
}
@@ -49,7 +49,7 @@ use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;
const ADMIN_SYSTEM_IMPORT_MAX_SIZE_BYTES: usize = 10 * 1024 * 1024;
const ADMIN_SYSTEM_IMPORT_MAX_SIZE_BYTES: usize = 500 * 1024 * 1024;
fn invalid_request(detail: impl Into<String>) -> (http::StatusCode, Value) {
(
@@ -173,6 +173,9 @@ fn normalize_import_endpoint_format(value: &str) -> Result<String, String> {
let normalized = match value.trim().to_ascii_lowercase().as_str() {
"openai:cli" => "openai:responses",
"openai:compact" => "openai:responses:compact",
"openai_image" | "images" | "image" | "/v1/images/generations" | "/v1/images/edits" => {
"openai:image"
}
"claude:chat" | "claude:cli" => "claude:messages",
"gemini:chat" | "gemini:cli" => "gemini:generate_content",
_ => value.trim(),
@@ -956,7 +959,7 @@ impl<'a> AdminAppState<'a> {
}
if request_body.len() > ADMIN_SYSTEM_DATA_IMPORT_MAX_SIZE_BYTES {
return Ok(Err(invalid_request("请求体大小不能超过 20MB")));
return Ok(Err(invalid_request("请求体大小不能超过 500MB")));
}
let root = match serde_json::from_slice::<Value>(request_body) {
@@ -1053,7 +1056,7 @@ impl<'a> AdminAppState<'a> {
)));
}
if request_body.len() > ADMIN_SYSTEM_IMPORT_MAX_SIZE_BYTES {
return Ok(Err(invalid_request("请求体大小不能超过 10MB")));
return Ok(Err(invalid_request("请求体大小不能超过 500MB")));
}
let parsed = routed!(parse_admin_system_config_import_request(request_body));
@@ -1993,7 +1996,7 @@ impl<'a> AdminAppState<'a> {
)));
}
if request_body.len() > ADMIN_SYSTEM_IMPORT_MAX_SIZE_BYTES {
return Ok(Err(invalid_request("请求体大小不能超过 10MB")));
return Ok(Err(invalid_request("请求体大小不能超过 500MB")));
}
let root = match serde_json::from_slice::<Value>(request_body) {
@@ -3062,6 +3065,10 @@ mod tests {
for (raw, expected) in [
("openai:cli", "openai:responses"),
("openai:compact", "openai:responses:compact"),
("openai_image", "openai:image"),
("images", "openai:image"),
("/v1/images/generations", "openai:image"),
("/v1/images/edits", "openai:image"),
("claude:chat", "claude:messages"),
("claude:cli", "claude:messages"),
("gemini:chat", "gemini:generate_content"),
@@ -9,7 +9,7 @@ mod proxy_nodes;
mod templates;
const ADMIN_SYSTEM_DATA_EXPORT_VERSION: &str = "1.0";
const ADMIN_SYSTEM_DATA_IMPORT_MAX_SIZE_BYTES: usize = 20 * 1024 * 1024;
const ADMIN_SYSTEM_DATA_IMPORT_MAX_SIZE_BYTES: usize = 500 * 1024 * 1024;
impl<'a> AdminAppState<'a> {
pub(crate) async fn upsert_system_config_json_value(
@@ -17,6 +17,7 @@ use crate::handlers::admin::system::shared::settings::{
build_admin_system_stats_payload, current_aether_version, fetch_latest_admin_system_release,
};
use crate::handlers::admin::system::shared::smtp::build_admin_smtp_test_payload;
use crate::important_notification::build_important_notification_test_payload;
use crate::maintenance::{ManualUsageCleanupMode, ManualUsageCleanupOptions};
use crate::GatewayError;
use aether_data_contracts::repository::usage::UsageCleanupTargets;
@@ -241,6 +242,16 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
));
}
if decision.route_kind.as_deref() == Some("important_notification_test")
&& request_method == http::Method::POST
&& request_path == "/api/admin/system/important-notification/test"
{
return Ok(Some(
Json(build_important_notification_test_payload(state, request_body).await?)
.into_response(),
));
}
if decision.route_kind.as_deref() == Some("cleanup") && request_method == http::Method::POST {
return Ok(Some(attach_admin_audit_response(
Json(build_admin_system_cleanup_payload(state).await?).into_response(),
@@ -33,6 +33,21 @@ fn admin_system_config_default_value(key: &str) -> Option<serde_json::Value> {
admin_system_config_default_value_pure(key)
}
fn legacy_admin_system_config_fallback_key(normalized_key: &str) -> Option<&'static str> {
match normalized_key {
"module.server_chan_push.enabled" => {
Some("module.important_notification.server_chan_enabled")
}
"module.server_chan_push.send_key" => {
Some("module.important_notification.server_chan_send_key")
}
"module.server_chan_push.template" => {
Some("module.important_notification.server_chan_template")
}
_ => None,
}
}
pub(crate) fn build_admin_system_configs_payload(
entries: &[aether_data::repository::system::StoredSystemConfigEntry],
) -> serde_json::Value {
@@ -44,12 +59,14 @@ pub(crate) async fn build_admin_system_config_detail_payload(
requested_key: &str,
) -> Result<Result<serde_json::Value, (http::StatusCode, serde_json::Value)>, GatewayError> {
let requested_key = requested_key.trim();
let value = state
.read_system_config_json_value(&normalize_admin_system_config_key(requested_key))
.await?
.or_else(|| {
admin_system_config_default_value(&normalize_admin_system_config_key(requested_key))
});
let normalized_key = normalize_admin_system_config_key(requested_key);
let mut value = state.read_system_config_json_value(&normalized_key).await?;
if value.is_none() {
if let Some(legacy_key) = legacy_admin_system_config_fallback_key(&normalized_key) {
value = state.read_system_config_json_value(legacy_key).await?;
}
}
let value = value.or_else(|| admin_system_config_default_value(&normalized_key));
Ok(build_admin_system_config_detail_payload_pure(
requested_key,
value,
@@ -1,5 +1,11 @@
use crate::bark_push::bark_push_configured;
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::shared::{module_available_from_env, system_config_bool};
use crate::important_notification::{
important_notification_configured, IMPORTANT_NOTIFICATION_ENABLED_KEY,
LEGACY_NOTIFICATION_EMAIL_ENABLED_KEY,
};
use crate::server_chan_push::server_chan_push_configured;
use crate::system_features::ENABLE_MODEL_DIRECTIVES_CONFIG_KEY;
use crate::GatewayError;
use aether_admin::system as admin_system_kernel;
@@ -68,17 +74,41 @@ pub(crate) const ADMIN_MODULE_DEFINITIONS: &[AdminModuleDefinition] = &[
admin_menu_order: 59,
},
AdminModuleDefinition {
name: "notification_email",
display_name: "异常通知",
description: "为 5xx 异常发送邮件通知,可在模块管理中启用或禁",
name: "important_notification",
display_name: "通知服务",
description: "统一管理通知项、模板和推送服务选择,供后台任务和用户通知使",
category: "integration",
env_key: "NOTIFICATION_EMAIL_AVAILABLE",
env_key: "IMPORTANT_NOTIFICATION_AVAILABLE",
default_available: true,
admin_route: None,
admin_menu_icon: Some("Mail"),
admin_menu_group: Some("system"),
admin_route: Some("/admin/notification-service"),
admin_menu_icon: Some("BellRing"),
admin_menu_group: None,
admin_menu_order: 58,
},
AdminModuleDefinition {
name: "server_chan_push",
display_name: "Server 酱推送",
description: "第三方推送服务,配置 Server 酱 Turbo SendKey 并测试微信推送",
category: "integration",
env_key: "SERVER_CHAN_PUSH_AVAILABLE",
default_available: true,
admin_route: Some("/admin/modules/server-chan"),
admin_menu_icon: Some("Send"),
admin_menu_group: Some("system"),
admin_menu_order: 59,
},
AdminModuleDefinition {
name: "bark_push",
display_name: "Bark 推送",
description: "第三方推送服务,配置 Bark Device Key 并测试 iOS 推送",
category: "integration",
env_key: "BARK_PUSH_AVAILABLE",
default_available: true,
admin_route: Some("/admin/modules/bark"),
admin_menu_icon: Some("Send"),
admin_menu_group: Some("system"),
admin_menu_order: 59,
},
AdminModuleDefinition {
name: "model_directives",
display_name: "模型后缀参数",
@@ -150,10 +180,17 @@ pub(crate) struct AdminModuleRuntimeState {
oauth_providers: Vec<aether_data::repository::auth_modules::StoredOAuthProviderModuleConfig>,
ldap_config: Option<aether_data::repository::auth_modules::StoredLdapModuleConfig>,
gemini_files_has_capable_key: bool,
smtp_configured: bool,
important_notification_configured: bool,
server_chan_push_configured: bool,
bark_push_configured: bool,
}
pub(crate) fn admin_module_by_name(name: &str) -> Option<&'static AdminModuleDefinition> {
let name = if name == "notification_email" {
"important_notification"
} else {
name
};
ADMIN_MODULE_DEFINITIONS
.iter()
.find(|module| module.name == name)
@@ -170,11 +207,22 @@ pub(crate) fn admin_module_name_from_enabled_path(request_path: &str) -> Option<
pub(crate) fn admin_module_enabled_config_key(module: &AdminModuleDefinition) -> String {
if module.name == "model_directives" {
ENABLE_MODEL_DIRECTIVES_CONFIG_KEY.to_string()
} else if module.name == "important_notification" {
IMPORTANT_NOTIFICATION_ENABLED_KEY.to_string()
} else {
format!("module.{}.enabled", module.name)
}
}
fn admin_module_available(module: &AdminModuleDefinition) -> bool {
if module.name == "important_notification" {
let legacy_default =
module_available_from_env("NOTIFICATION_EMAIL_AVAILABLE", module.default_available);
return module_available_from_env(module.env_key, legacy_default);
}
module_available_from_env(module.env_key, module.default_available)
}
pub(crate) fn oauth_module_config_is_valid(
providers: &[aether_data::repository::auth_modules::StoredOAuthProviderModuleConfig],
) -> bool {
@@ -221,28 +269,17 @@ pub(crate) async fn build_admin_module_runtime_state(
})
};
let smtp_host = state.read_system_config_json_value("smtp_host").await?;
let smtp_from_email = state
.read_system_config_json_value("smtp_from_email")
.await?;
let smtp_configured = smtp_host
.as_ref()
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
&& smtp_from_email
.as_ref()
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some();
let notification_configured = important_notification_configured(state.app()).await?;
let server_chan_configured = server_chan_push_configured(state.app()).await?;
let bark_configured = bark_push_configured(state.app()).await?;
Ok(AdminModuleRuntimeState {
oauth_providers,
ldap_config,
gemini_files_has_capable_key,
smtp_configured,
important_notification_configured: notification_configured,
server_chan_push_configured: server_chan_configured,
bark_push_configured: bark_configured,
})
}
@@ -255,7 +292,9 @@ pub(crate) fn build_admin_module_validation_result(
&runtime.oauth_providers,
runtime.ldap_config.as_ref(),
runtime.gemini_files_has_capable_key,
runtime.smtp_configured,
runtime.important_notification_configured,
runtime.server_chan_push_configured,
runtime.bark_push_configured,
)
}
@@ -274,12 +313,19 @@ pub(crate) async fn build_admin_module_status_payload(
module: &AdminModuleDefinition,
runtime: &AdminModuleRuntimeState,
) -> Result<serde_json::Value, GatewayError> {
let available = module_available_from_env(module.env_key, module.default_available);
let available = admin_module_available(module);
let enabled = if available {
let enabled = state
let enabled_value = state
.read_system_config_json_value(&admin_module_enabled_config_key(module))
.await?;
system_config_bool(enabled.as_ref(), false)
let enabled_value = if module.name == "important_notification" && enabled_value.is_none() {
state
.read_system_config_json_value(LEGACY_NOTIFICATION_EMAIL_ENABLED_KEY)
.await?
} else {
enabled_value
};
system_config_bool(enabled_value.as_ref(), false)
} else {
false
};
@@ -1,14 +1,10 @@
use crate::email_delivery::{probe_smtp_connection, system_config_u16, SmtpDeliveryConfig};
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::shared::{system_config_bool, system_config_string};
use crate::GatewayError;
use axum::body::Bytes;
use base64::Engine;
use serde::Deserialize;
use serde_json::json;
use std::io::{BufRead, Write};
use std::time::Duration;
const SMTP_TIMEOUT_SECS: u64 = 30;
#[derive(Debug, Default, Deserialize)]
struct AdminSmtpTestRequest {
@@ -53,12 +49,12 @@ pub(crate) async fn build_admin_smtp_test_payload(
}));
}
let result = tokio::task::spawn_blocking(move || test_smtp_connection_blocking(config))
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let result = probe_smtp_connection(config.into_delivery_config()).await;
Ok(match result {
Ok(()) => json!({ "success": true, "message": "SMTP 连接测试成功" }),
Err(error) => json!({ "success": false, "message": translate_smtp_error(&error) }),
Err(error) => {
json!({ "success": false, "message": translate_smtp_error(&smtp_gateway_error_message(&error)) })
}
})
}
@@ -94,8 +90,8 @@ async fn resolve_admin_smtp_config(
port: request
.smtp_port
.as_ref()
.map(|value| system_config_u16(value, 587))
.unwrap_or_else(|| system_config_u16_opt(smtp_port.as_ref(), 587)),
.map(|value| system_config_u16(Some(value), 587))
.unwrap_or_else(|| system_config_u16(smtp_port.as_ref(), 587)),
user: request
.smtp_user
.as_ref()
@@ -130,6 +126,21 @@ async fn resolve_admin_smtp_config(
})
}
impl ResolvedSmtpConfig {
fn into_delivery_config(self) -> SmtpDeliveryConfig {
SmtpDeliveryConfig {
host: self.host.unwrap_or_default(),
port: self.port,
user: self.user,
password: self.password,
use_tls: self.use_tls,
use_ssl: self.use_ssl,
from_email: self.from_email.unwrap_or_default(),
from_name: self.from_name,
}
}
}
fn missing_smtp_fields(config: &ResolvedSmtpConfig) -> Vec<&'static str> {
let mut fields = Vec::new();
if config
@@ -171,178 +182,13 @@ fn missing_smtp_fields(config: &ResolvedSmtpConfig) -> Vec<&'static str> {
fields
}
fn system_config_u16_opt(value: Option<&serde_json::Value>, default: u16) -> u16 {
value
.map(|value| system_config_u16(value, default))
.unwrap_or(default)
}
fn system_config_u16(value: &serde_json::Value, default: u16) -> u16 {
match value {
serde_json::Value::Number(value) => value
.as_u64()
.and_then(|value| u16::try_from(value).ok())
.unwrap_or(default),
serde_json::Value::String(value) => value.trim().parse::<u16>().unwrap_or(default),
_ => default,
fn smtp_gateway_error_message(error: &GatewayError) -> String {
match error {
GatewayError::Internal(message) => message.clone(),
_ => format!("{error:?}"),
}
}
fn build_tls_config() -> std::sync::Arc<rustls::ClientConfig> {
let _ = rustls::crypto::ring::default_provider().install_default();
let root_store =
rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
std::sync::Arc::new(
rustls::ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth(),
)
}
fn resolve_server_name(host: &str) -> Result<rustls::pki_types::ServerName<'static>, String> {
let host = host.trim().trim_start_matches('[').trim_end_matches(']');
if let Ok(ip) = host.parse::<std::net::IpAddr>() {
return Ok(rustls::pki_types::ServerName::from(ip));
}
rustls::pki_types::ServerName::try_from(host.to_string()).map_err(|err| err.to_string())
}
fn connect_tcp_stream(config: &ResolvedSmtpConfig) -> Result<std::net::TcpStream, String> {
let host = config.host.as_deref().unwrap_or_default();
let stream =
std::net::TcpStream::connect((host, config.port)).map_err(|err| err.to_string())?;
stream
.set_read_timeout(Some(Duration::from_secs(SMTP_TIMEOUT_SECS)))
.map_err(|err| err.to_string())?;
stream
.set_write_timeout(Some(Duration::from_secs(SMTP_TIMEOUT_SECS)))
.map_err(|err| err.to_string())?;
Ok(stream)
}
fn wrap_tls_stream(
stream: std::net::TcpStream,
host: &str,
) -> Result<rustls::StreamOwned<rustls::ClientConnection, std::net::TcpStream>, String> {
let server_name = resolve_server_name(host)?;
let connection = rustls::ClientConnection::new(build_tls_config(), server_name)
.map_err(|err| err.to_string())?;
Ok(rustls::StreamOwned::new(connection, stream))
}
fn smtp_read_response<T: BufRead>(reader: &mut T) -> Result<(u16, String), String> {
let mut message = String::new();
let code = loop {
let mut line = String::new();
let bytes = reader.read_line(&mut line).map_err(|err| err.to_string())?;
if bytes == 0 {
return Err("smtp connection closed unexpectedly".to_string());
}
let trimmed = line.trim_end_matches(['\r', '\n']).to_string();
if trimmed.len() < 3 {
return Err("invalid smtp response".to_string());
}
let parsed_code = trimmed[..3].parse::<u16>().map_err(|err| err.to_string())?;
let continuation = trimmed.as_bytes().get(3).copied() == Some(b'-');
if !message.is_empty() {
message.push('\n');
}
message.push_str(&trimmed);
if !continuation {
break parsed_code;
}
};
Ok((code, message))
}
fn smtp_expect<T: BufRead>(reader: &mut T, allowed_codes: &[u16]) -> Result<String, String> {
let (code, message) = smtp_read_response(reader)?;
if allowed_codes.contains(&code) {
return Ok(message);
}
Err(format!("unexpected smtp response {code}: {message}"))
}
fn smtp_write_line<T: Write>(writer: &mut T, line: &str) -> Result<(), String> {
writer
.write_all(line.as_bytes())
.map_err(|err| err.to_string())?;
writer.write_all(b"\r\n").map_err(|err| err.to_string())?;
writer.flush().map_err(|err| err.to_string())
}
fn smtp_send_command<S: std::io::Read + Write>(
reader: &mut std::io::BufReader<S>,
command: &str,
allowed_codes: &[u16],
) -> Result<String, String> {
smtp_write_line(reader.get_mut(), command)?;
smtp_expect(reader, allowed_codes)
}
fn smtp_authenticate<S: std::io::Read + Write>(
reader: &mut std::io::BufReader<S>,
config: &ResolvedSmtpConfig,
) -> Result<(), String> {
let Some(username) = config
.user
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(());
};
let password = config.password.as_deref().unwrap_or_default();
smtp_send_command(reader, "AUTH LOGIN", &[334])?;
smtp_send_command(
reader,
&base64::engine::general_purpose::STANDARD.encode(username.as_bytes()),
&[334],
)?;
smtp_send_command(
reader,
&base64::engine::general_purpose::STANDARD.encode(password.as_bytes()),
&[235],
)?;
Ok(())
}
fn smtp_probe<S: std::io::Read + Write>(
reader: &mut std::io::BufReader<S>,
config: &ResolvedSmtpConfig,
) -> Result<(), String> {
smtp_send_command(reader, "EHLO aether.local", &[250])?;
smtp_authenticate(reader, config)?;
let _ = smtp_send_command(reader, "QUIT", &[221]);
Ok(())
}
fn test_smtp_connection_blocking(config: ResolvedSmtpConfig) -> Result<(), String> {
if config.use_ssl {
let stream = connect_tcp_stream(&config)?;
let tls_stream = wrap_tls_stream(stream, config.host.as_deref().unwrap_or_default())?;
let mut reader = std::io::BufReader::new(tls_stream);
smtp_expect(&mut reader, &[220])?;
return smtp_probe(&mut reader, &config);
}
let stream = connect_tcp_stream(&config)?;
let mut reader = std::io::BufReader::new(stream);
smtp_expect(&mut reader, &[220])?;
smtp_send_command(&mut reader, "EHLO aether.local", &[250])?;
if config.use_tls {
smtp_send_command(&mut reader, "STARTTLS", &[220])?;
let stream = reader.into_inner();
let tls_stream = wrap_tls_stream(stream, config.host.as_deref().unwrap_or_default())?;
let mut reader = std::io::BufReader::new(tls_stream);
return smtp_probe(&mut reader, &config);
}
smtp_authenticate(&mut reader, &config)?;
let _ = smtp_send_command(&mut reader, "QUIT", &[221]);
Ok(())
}
fn translate_smtp_error(error: &str) -> String {
let error_lower = error.to_ascii_lowercase();
@@ -1,11 +1,11 @@
use super::{
decrypt_catalog_secret_with_fallbacks, escape_admin_email_template_html, json,
read_admin_email_template_payload, render_admin_email_template_html, system_config_bool,
system_config_string, system_config_u16, AppState, GatewayError,
escape_admin_email_template_html, json, read_admin_email_template_payload,
render_admin_email_template_html, system_config_string, AppState, GatewayError,
AUTH_EMAIL_VERIFICATION_PREFIX, AUTH_EMAIL_VERIFIED_PREFIX, AUTH_EMAIL_VERIFIED_TTL_SECS,
AUTH_SMTP_TIMEOUT_SECS,
};
use base64::Engine;
use crate::email_delivery::{
read_smtp_delivery_config, send_smtp_email, ComposedEmail, SmtpDeliveryConfig,
};
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub(super) struct StoredAuthEmailVerificationCode {
@@ -13,25 +13,8 @@ pub(super) struct StoredAuthEmailVerificationCode {
pub(super) created_at: String,
}
#[derive(Debug, Clone)]
pub(super) struct AuthSmtpConfig {
pub(super) host: String,
pub(super) port: u16,
pub(super) user: Option<String>,
pub(super) password: Option<String>,
pub(super) use_tls: bool,
pub(super) use_ssl: bool,
pub(super) from_email: String,
pub(super) from_name: String,
}
#[derive(Debug, Clone)]
pub(super) struct AuthComposedEmail {
pub(super) to_email: String,
pub(super) subject: String,
pub(super) html_body: String,
pub(super) text_body: String,
}
pub(super) type AuthSmtpConfig = SmtpDeliveryConfig;
pub(super) type AuthComposedEmail = ComposedEmail;
pub(super) fn auth_email_verification_key(email: &str) -> String {
format!("{AUTH_EMAIL_VERIFICATION_PREFIX}{email}")
@@ -84,25 +67,6 @@ fn render_auth_template_string(
Ok(rendered)
}
fn auth_encode_mime_header(value: &str) -> String {
if value.is_ascii() {
return value.to_string();
}
format!(
"=?UTF-8?B?{}?=",
base64::engine::general_purpose::STANDARD.encode(value.as_bytes())
)
}
fn auth_wrap_base64(value: &str) -> String {
let mut wrapped = String::new();
for chunk in value.as_bytes().chunks(76) {
wrapped.push_str(std::str::from_utf8(chunk).unwrap_or_default());
wrapped.push_str("\r\n");
}
wrapped
}
fn auth_build_verification_text_body(
app_name: &str,
email: &str,
@@ -114,244 +78,6 @@ fn auth_build_verification_text_body(
)
}
fn auth_build_tls_config() -> std::sync::Arc<rustls::ClientConfig> {
let _ = rustls::crypto::ring::default_provider().install_default();
let root_store =
rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
let config = rustls::ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
std::sync::Arc::new(config)
}
fn auth_resolve_server_name(
host: &str,
) -> Result<rustls::pki_types::ServerName<'static>, GatewayError> {
let host = host.trim().trim_start_matches('[').trim_end_matches(']');
if let Ok(ip) = host.parse::<std::net::IpAddr>() {
return Ok(rustls::pki_types::ServerName::from(ip));
}
rustls::pki_types::ServerName::try_from(host.to_string())
.map_err(|err| GatewayError::Internal(err.to_string()))
}
fn auth_connect_tcp_stream(config: &AuthSmtpConfig) -> Result<std::net::TcpStream, GatewayError> {
let stream = std::net::TcpStream::connect((config.host.as_str(), config.port))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
stream
.set_read_timeout(Some(std::time::Duration::from_secs(AUTH_SMTP_TIMEOUT_SECS)))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
stream
.set_write_timeout(Some(std::time::Duration::from_secs(AUTH_SMTP_TIMEOUT_SECS)))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok(stream)
}
fn auth_wrap_tls_stream(
stream: std::net::TcpStream,
host: &str,
) -> Result<rustls::StreamOwned<rustls::ClientConnection, std::net::TcpStream>, GatewayError> {
let server_name = auth_resolve_server_name(host)?;
let connection = rustls::ClientConnection::new(auth_build_tls_config(), server_name)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok(rustls::StreamOwned::new(connection, stream))
}
fn auth_smtp_read_response<T: std::io::BufRead>(
reader: &mut T,
) -> Result<(u16, String), GatewayError> {
let mut message = String::new();
let code = loop {
let parsed_code;
let continuation;
let trimmed;
{
let mut line = String::new();
let bytes = reader
.read_line(&mut line)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if bytes == 0 {
return Err(GatewayError::Internal(
"smtp connection closed unexpectedly".to_string(),
));
}
trimmed = line.trim_end_matches(['\r', '\n']).to_string();
if trimmed.len() < 3 {
return Err(GatewayError::Internal("invalid smtp response".to_string()));
}
parsed_code = trimmed[..3]
.parse::<u16>()
.map_err(|err| GatewayError::Internal(err.to_string()))?;
continuation = trimmed.as_bytes().get(3).copied() == Some(b'-');
}
if !message.is_empty() {
message.push('\n');
}
message.push_str(&trimmed);
if !continuation {
break parsed_code;
}
};
Ok((code, message))
}
fn auth_smtp_expect<T: std::io::BufRead>(
reader: &mut T,
allowed_codes: &[u16],
) -> Result<String, GatewayError> {
let (code, message) = auth_smtp_read_response(reader)?;
if allowed_codes.contains(&code) {
return Ok(message);
}
Err(GatewayError::Internal(format!(
"unexpected smtp response {code}: {message}"
)))
}
fn auth_smtp_write_line<T: std::io::Write>(writer: &mut T, line: &str) -> Result<(), GatewayError> {
writer
.write_all(line.as_bytes())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
writer
.write_all(b"\r\n")
.map_err(|err| GatewayError::Internal(err.to_string()))?;
writer
.flush()
.map_err(|err| GatewayError::Internal(err.to_string()))
}
fn auth_smtp_send_command<S: std::io::Read + std::io::Write>(
reader: &mut std::io::BufReader<S>,
command: &str,
allowed_codes: &[u16],
) -> Result<String, GatewayError> {
auth_smtp_write_line(reader.get_mut(), command)?;
auth_smtp_expect(reader, allowed_codes)
}
fn auth_build_email_message(config: &AuthSmtpConfig, email: &AuthComposedEmail) -> String {
let boundary = format!("aether-{}", uuid::Uuid::new_v4().simple());
let text_body = auth_wrap_base64(
&base64::engine::general_purpose::STANDARD.encode(email.text_body.as_bytes()),
);
let html_body = auth_wrap_base64(
&base64::engine::general_purpose::STANDARD.encode(email.html_body.as_bytes()),
);
let from_header = if config.from_name.trim().is_empty() {
format!("<{}>", config.from_email)
} else {
format!(
"{} <{}>",
auth_encode_mime_header(config.from_name.trim()),
config.from_email
)
};
format!(
"From: {from_header}\r\nTo: <{to_email}>\r\nSubject: {subject}\r\nMIME-Version: 1.0\r\nContent-Type: multipart/alternative; boundary=\"{boundary}\"\r\n\r\n--{boundary}\r\nContent-Type: text/plain; charset=\"utf-8\"\r\nContent-Transfer-Encoding: base64\r\n\r\n{text_body}--{boundary}\r\nContent-Type: text/html; charset=\"utf-8\"\r\nContent-Transfer-Encoding: base64\r\n\r\n{html_body}--{boundary}--\r\n",
to_email = email.to_email,
subject = auth_encode_mime_header(&email.subject),
)
}
fn auth_smtp_authenticate<S: std::io::Read + std::io::Write>(
reader: &mut std::io::BufReader<S>,
config: &AuthSmtpConfig,
) -> Result<(), GatewayError> {
let Some(username) = config
.user
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(());
};
let password = config.password.as_deref().unwrap_or("");
auth_smtp_send_command(reader, "AUTH LOGIN", &[334])?;
auth_smtp_send_command(
reader,
&base64::engine::general_purpose::STANDARD.encode(username.as_bytes()),
&[334],
)?;
auth_smtp_send_command(
reader,
&base64::engine::general_purpose::STANDARD.encode(password.as_bytes()),
&[235],
)?;
Ok(())
}
fn auth_smtp_deliver_message<S: std::io::Read + std::io::Write>(
reader: &mut std::io::BufReader<S>,
config: &AuthSmtpConfig,
email: &AuthComposedEmail,
) -> Result<(), GatewayError> {
auth_smtp_send_command(
reader,
&format!("MAIL FROM:<{}>", config.from_email),
&[250],
)?;
auth_smtp_send_command(
reader,
&format!("RCPT TO:<{}>", email.to_email),
&[250, 251],
)?;
auth_smtp_send_command(reader, "DATA", &[354])?;
let message = auth_build_email_message(config, email);
reader
.get_mut()
.write_all(message.as_bytes())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
reader
.get_mut()
.write_all(b"\r\n.\r\n")
.map_err(|err| GatewayError::Internal(err.to_string()))?;
reader
.get_mut()
.flush()
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let _ = auth_smtp_expect(reader, &[250])?;
let _ = auth_smtp_send_command(reader, "QUIT", &[221]);
Ok(())
}
fn auth_smtp_send_message<S: std::io::Read + std::io::Write>(
reader: &mut std::io::BufReader<S>,
config: &AuthSmtpConfig,
email: &AuthComposedEmail,
) -> Result<(), GatewayError> {
auth_smtp_send_command(reader, "EHLO aether.local", &[250])?;
auth_smtp_authenticate(reader, config)?;
auth_smtp_deliver_message(reader, config, email)
}
fn send_auth_email_blocking(
config: AuthSmtpConfig,
email: AuthComposedEmail,
) -> Result<(), GatewayError> {
if config.use_ssl {
let stream = auth_connect_tcp_stream(&config)?;
let tls_stream = auth_wrap_tls_stream(stream, &config.host)?;
let mut reader = std::io::BufReader::new(tls_stream);
let _ = auth_smtp_expect(&mut reader, &[220])?;
return auth_smtp_send_message(&mut reader, &config, &email);
}
let stream = auth_connect_tcp_stream(&config)?;
let mut reader = std::io::BufReader::new(stream);
let _ = auth_smtp_expect(&mut reader, &[220])?;
let _ = auth_smtp_send_command(&mut reader, "EHLO aether.local", &[250])?;
if config.use_tls {
let _ = auth_smtp_send_command(&mut reader, "STARTTLS", &[220])?;
let stream = reader.into_inner();
let tls_stream = auth_wrap_tls_stream(stream, &config.host)?;
let mut reader = std::io::BufReader::new(tls_stream);
return auth_smtp_send_message(&mut reader, &config, &email);
}
auth_smtp_authenticate(&mut reader, &config)?;
auth_smtp_deliver_message(&mut reader, &config, &email)
}
pub(super) async fn read_auth_email_verification_code(
state: &AppState,
email: &str,
@@ -423,40 +149,7 @@ pub(super) async fn store_auth_email_verification_code(
pub(super) async fn read_auth_smtp_config(
state: &AppState,
) -> Result<Option<AuthSmtpConfig>, GatewayError> {
let smtp_host = state.read_system_config_json_value("smtp_host").await?;
let smtp_from_email = state
.read_system_config_json_value("smtp_from_email")
.await?;
let Some(host) = system_config_string(smtp_host.as_ref()) else {
return Ok(None);
};
let Some(from_email) = system_config_string(smtp_from_email.as_ref()) else {
return Ok(None);
};
let smtp_port = state.read_system_config_json_value("smtp_port").await?;
let smtp_user = state.read_system_config_json_value("smtp_user").await?;
let smtp_password = state.read_system_config_json_value("smtp_password").await?;
let smtp_use_tls = state.read_system_config_json_value("smtp_use_tls").await?;
let smtp_use_ssl = state.read_system_config_json_value("smtp_use_ssl").await?;
let smtp_from_name = state
.read_system_config_json_value("smtp_from_name")
.await?;
let password = system_config_string(smtp_password.as_ref()).map(|value| {
decrypt_catalog_secret_with_fallbacks(state.encryption_key(), &value).unwrap_or(value)
});
Ok(Some(AuthSmtpConfig {
host,
port: system_config_u16(smtp_port.as_ref(), 587),
user: system_config_string(smtp_user.as_ref()),
password,
use_tls: system_config_bool(smtp_use_tls.as_ref(), true),
use_ssl: system_config_bool(smtp_use_ssl.as_ref(), false),
from_email,
from_name: system_config_string(smtp_from_name.as_ref())
.unwrap_or_else(|| "Aether".to_string()),
}))
read_smtp_delivery_config(state).await
}
pub(super) async fn auth_email_app_name(state: &AppState) -> Result<String, GatewayError> {
@@ -516,27 +209,20 @@ pub(super) async fn send_auth_email(
if record_auth_email_delivery_for_tests(
state,
json!({
"to_email": email.to_email,
"subject": email.subject,
"html_body": email.html_body,
"text_body": email.text_body,
"to_email": email.to_email.clone(),
"subject": email.subject.clone(),
"html_body": email.html_body.clone(),
"text_body": email.text_body.clone(),
}),
) {
return Ok(());
}
tokio::task::spawn_blocking(move || send_auth_email_blocking(config, email))
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
send_smtp_email(config, email).await
}
pub(super) async fn auth_registration_email_configured(
state: &AppState,
) -> Result<bool, GatewayError> {
let smtp_host = state.read_system_config_json_value("smtp_host").await?;
let smtp_from_email = state
.read_system_config_json_value("smtp_from_email")
.await?;
Ok(system_config_string(smtp_host.as_ref()).is_some()
&& system_config_string(smtp_from_email.as_ref()).is_some())
Ok(read_smtp_delivery_config(state).await?.is_some())
}
@@ -121,7 +121,6 @@ pub(super) const AUTH_REFRESH_TOKEN_EXPIRATION_DAYS: i64 = 7;
pub(super) const AUTH_EMAIL_VERIFICATION_PREFIX: &str = "email:verification:";
pub(super) const AUTH_EMAIL_VERIFIED_PREFIX: &str = "email:verified:";
pub(super) const AUTH_EMAIL_VERIFIED_TTL_SECS: u64 = 3600;
pub(super) const AUTH_SMTP_TIMEOUT_SECS: u64 = 30;
pub(crate) fn build_auth_json_response(
status: http::StatusCode,
@@ -7,7 +7,9 @@ use axum::{
use serde::Deserialize;
use serde_json::json;
use crate::handlers::shared::{deserialize_optional_json_patch, normalize_feature_settings};
use crate::handlers::shared::{
deserialize_optional_json_patch, normalize_user_self_feature_settings_update,
};
use super::{
auth_password_policy_level, build_auth_error_response, resolve_authenticated_local_user,
@@ -65,12 +67,24 @@ pub(super) async fn handle_users_me_detail_put(
let email = normalize_users_me_optional_non_empty_string(payload.email);
let username = normalize_users_me_optional_non_empty_string(payload.username);
let feature_settings = match payload.feature_settings {
Some(value) => match normalize_feature_settings(value) {
Ok(value) => Some(value),
Err(detail) => {
return build_auth_error_response(http::StatusCode::BAD_REQUEST, detail, false);
Some(value) => {
let current = match state.read_user_feature_settings(&auth.user.id).await {
Ok(value) => value,
Err(err) => {
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user feature settings lookup failed: {err:?}"),
false,
)
}
};
match normalize_user_self_feature_settings_update(value, current) {
Ok(value) => Some(value),
Err(detail) => {
return build_auth_error_response(http::StatusCode::BAD_REQUEST, detail, false);
}
}
},
}
None => None,
};
@@ -1,5 +1,6 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_ai_serving::UPSTREAM_IS_STREAM_KEY;
use aether_billing::{
normalize_input_tokens_for_billing, normalize_total_input_context_for_cache_hit_rate,
};
@@ -314,7 +315,7 @@ fn users_me_usage_upstream_is_stream(item: &StoredRequestUsageAudit) -> bool {
item.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.or_else(|| users_me_usage_headers_stream_flag(item.response_headers.as_ref()))
.or_else(|| users_me_usage_infer_upstream_stream_from_captured_bodies(item))
@@ -1132,6 +1132,278 @@ fn build_chatgpt_web_quota_status_snapshot(
}))
}
fn windsurf_percent_quota_window_snapshot(
metadata: &Map<String, Value>,
code: &str,
label: &str,
remaining_percent_key: &str,
reset_at_key: &str,
observed_at_unix_secs: Option<u64>,
) -> Option<Value> {
let remaining_percent = metadata
.get(remaining_percent_key)
.and_then(admin_provider_quota_pure::coerce_json_f64);
let reset_at = provider_quota_timestamp_unix_secs(metadata.get(reset_at_key));
if remaining_percent.is_none() && reset_at.is_none() {
return None;
}
let remaining_ratio = remaining_percent.map(|value| (value / 100.0).clamp(0.0, 1.0));
let used_ratio = remaining_ratio.map(|value| (1.0 - value).clamp(0.0, 1.0));
let reset_seconds = quota_window_reset_seconds(observed_at_unix_secs, reset_at);
Some(json!({
"code": code,
"label": label,
"scope": "account",
"unit": "percent",
"used_ratio": used_ratio,
"remaining_ratio": remaining_ratio,
"reset_at": reset_at,
"reset_seconds": reset_seconds,
"is_exhausted": remaining_ratio.map(|value| value <= 1e-6),
}))
}
fn windsurf_count_quota_window_snapshot(
metadata: &Map<String, Value>,
code: &str,
label: &str,
used_key: &str,
limit_key: &str,
remaining_key: &str,
) -> Option<Value> {
let used = metadata
.get(used_key)
.and_then(admin_provider_quota_pure::coerce_json_f64);
let limit = metadata
.get(limit_key)
.and_then(admin_provider_quota_pure::coerce_json_f64);
let remaining = metadata
.get(remaining_key)
.and_then(admin_provider_quota_pure::coerce_json_f64)
.or_else(|| limit.zip(used).map(|(limit, used)| (limit - used).max(0.0)));
if used.is_none() && limit.is_none() && remaining.is_none() {
return None;
}
let used_ratio = used
.zip(limit)
.and_then(|(used, limit)| (limit > 0.0).then_some((used / limit).clamp(0.0, 1.0)));
let remaining_ratio = remaining.zip(limit).and_then(|(remaining, limit)| {
(limit > 0.0).then_some((remaining / limit).clamp(0.0, 1.0))
});
Some(json!({
"code": code,
"label": label,
"scope": "account",
"unit": "count",
"used_ratio": used_ratio,
"remaining_ratio": remaining_ratio,
"used_value": used,
"remaining_value": remaining,
"limit_value": limit,
"is_exhausted": remaining.is_some_and(|value| value <= 0.0),
}))
}
fn build_windsurf_quota_status_snapshot(
upstream_metadata: Option<&Value>,
source: &str,
) -> Option<Value> {
let metadata = provider_quota_metadata_bucket(upstream_metadata, "windsurf")?;
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
let plan_type = metadata
.get("plan_name")
.or_else(|| metadata.get("plan_type"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let rate_limit = metadata
.get("rate_limit")
.cloned()
.filter(|value| !value.is_null());
let last_error = metadata
.get("last_error")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let banned = metadata
.get("banned")
.or_else(|| metadata.get("is_banned"))
.and_then(admin_provider_quota_pure::coerce_json_bool)
== Some(true);
let quarantined = metadata
.get("quarantined")
.or_else(|| metadata.get("is_quarantined"))
.and_then(admin_provider_quota_pure::coerce_json_bool)
== Some(true);
let mut windows = [
windsurf_percent_quota_window_snapshot(
metadata,
"daily",
"",
"daily_remaining_percent",
"daily_reset_at",
observed_at_unix_secs,
),
windsurf_percent_quota_window_snapshot(
metadata,
"weekly",
"",
"weekly_remaining_percent",
"weekly_reset_at",
observed_at_unix_secs,
),
windsurf_count_quota_window_snapshot(
metadata,
"prompt",
"Prompt",
"prompt_used",
"prompt_limit",
"prompt_remaining",
),
windsurf_count_quota_window_snapshot(
metadata,
"flex",
"Flex",
"flex_used",
"flex_limit",
"flex_remaining",
),
]
.into_iter()
.flatten()
.collect::<Vec<_>>();
let mut rate_limit_cooling = false;
let mut rate_limit_reset_seconds = None;
let mut rate_limit_reason = None::<String>;
if let Some(rate_limit) = rate_limit.as_ref() {
if let Some(rate_limit_object) = rate_limit.as_object() {
let retry_after_ms = rate_limit_object
.get("retry_after_ms")
.or_else(|| rate_limit_object.get("retryAfterMs"))
.and_then(admin_provider_quota_pure::coerce_json_u64)
.filter(|value| *value > 0);
if let Some(retry_after_ms) = retry_after_ms {
rate_limit_cooling = true;
rate_limit_reset_seconds = Some(retry_after_ms.saturating_add(999) / 1000);
rate_limit_reason = rate_limit_object
.get("message")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
windows.push(json!({
"code": "rate_limit",
"label": "速率",
"scope": "account",
"unit": "count",
"is_exhausted": false,
"reset_seconds": rate_limit_reset_seconds,
}));
}
}
}
let allowed_models_count = metadata
.get("allowed_models_count")
.or_else(|| metadata.get("models_count"))
.and_then(admin_provider_quota_pure::coerce_json_u64);
if windows.is_empty()
&& plan_type.is_none()
&& observed_at_unix_secs.is_none()
&& rate_limit.is_none()
&& allowed_models_count.is_none()
&& !banned
&& !quarantined
{
return None;
}
let usage_ratio = quota_windows_usage_ratio(&windows);
let reset_seconds = if rate_limit_cooling {
rate_limit_reset_seconds.or_else(|| quota_windows_min_reset_seconds(&windows))
} else {
quota_windows_min_reset_seconds(&windows)
};
let reset_at = if rate_limit_cooling {
None
} else {
quota_windows_min_reset_at(&windows)
};
let exhausted_by_window = windows.iter().filter_map(Value::as_object).any(|window| {
window
.get("code")
.and_then(Value::as_str)
.is_some_and(|code| {
code.eq_ignore_ascii_case("daily") || code.eq_ignore_ascii_case("weekly")
})
&& window
.get("is_exhausted")
.and_then(admin_provider_quota_pure::coerce_json_bool)
.unwrap_or(false)
});
let exhausted = banned || quarantined || exhausted_by_window;
let (code, label, reason) = if banned {
(
"banned",
Some("账号已封禁"),
last_error
.clone()
.or_else(|| Some("账号被 Windsurf 标记为不可用".to_string())),
)
} else if quarantined {
(
"quarantined",
Some("账号隔离中"),
last_error
.clone()
.or_else(|| Some("账号处于隔离状态".to_string())),
)
} else if rate_limit_cooling {
(
"cooldown",
Some("冷却中"),
last_error.clone().or(rate_limit_reason),
)
} else if exhausted {
(
"exhausted",
Some("额度耗尽"),
Some("额度窗口已耗尽".to_string()),
)
} else {
("ok", None, last_error)
};
Some(json!({
"version": 2,
"provider_type": "windsurf",
"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": reset_at,
"reset_seconds": reset_seconds,
"plan_type": plan_type,
"allowed_models_count": allowed_models_count,
"rate_limit": rate_limit.unwrap_or(Value::Null),
"windows": windows,
}))
}
fn build_antigravity_quota_status_snapshot(
upstream_metadata: Option<&Value>,
source: &str,
@@ -1516,6 +1788,7 @@ pub(crate) fn sync_provider_key_quota_status_snapshot(
"codex" => build_codex_quota_status_snapshot(upstream_metadata, source),
"kiro" => build_kiro_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),
"grok" => build_grok_quota_status_snapshot(upstream_metadata, source),
"gemini_cli" => build_gemini_cli_quota_status_snapshot(upstream_metadata, source),
@@ -1552,6 +1825,12 @@ fn quota_snapshot_has_materialized_data(
return false;
}
if normalized_provider_type == "windsurf"
&& windsurf_quota_snapshot_has_stale_cooldown(quota_snapshot)
{
return false;
}
if quota_snapshot
.get("windows")
.and_then(Value::as_array)
@@ -1577,6 +1856,61 @@ fn quota_snapshot_has_materialized_data(
})
}
fn windsurf_quota_snapshot_has_stale_cooldown(quota_snapshot: &Map<String, Value>) -> bool {
let code = quota_snapshot
.get("code")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if !code.eq_ignore_ascii_case("cooldown") {
return false;
}
let rate_limit = quota_snapshot.get("rate_limit").and_then(Value::as_object);
let retry_after_ms = rate_limit
.and_then(|rate_limit| {
rate_limit
.get("retry_after_ms")
.or_else(|| rate_limit.get("retryAfterMs"))
.and_then(admin_provider_quota_pure::coerce_json_u64)
})
.unwrap_or(0);
if retry_after_ms > 0 {
return false;
}
let has_positive_rate_limit_reset = quota_snapshot
.get("windows")
.and_then(Value::as_array)
.is_some_and(|windows| {
windows.iter().filter_map(Value::as_object).any(|window| {
window
.get("code")
.and_then(Value::as_str)
.is_some_and(|code| code.eq_ignore_ascii_case("rate_limit"))
&& window
.get("reset_seconds")
.or_else(|| window.get("reset_at"))
.and_then(admin_provider_quota_pure::coerce_json_u64)
.is_some_and(|value| value > 0)
})
});
if has_positive_rate_limit_reset {
return false;
}
let exhausted = quota_snapshot
.get("exhausted")
.and_then(admin_provider_quota_pure::coerce_json_bool)
.unwrap_or(false);
let has_capacity = rate_limit
.and_then(|rate_limit| rate_limit.get("has_capacity"))
.and_then(admin_provider_quota_pure::coerce_json_bool)
.unwrap_or(false);
has_capacity || !exhausted
}
pub(crate) fn provider_key_status_snapshot_payload(
key: &StoredProviderCatalogKey,
provider_type: &str,
@@ -2504,6 +2838,274 @@ mod tests {
assert_eq!(windows[0].get("remaining_ratio"), Some(&json!(0.75)));
}
#[test]
fn provider_key_status_snapshot_payload_backfills_windsurf_daily_and_weekly_quota() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"windsurf": {
"updated_at": 1_778_067_246u64,
"plan_name": "Pro",
"daily_remaining_percent": 40.0,
"weekly_remaining_percent": 65.0,
"daily_reset_at": 1_778_100_000u64,
"weekly_reset_at": 1_778_600_000u64,
"prompt_used": 12.0,
"prompt_limit": 100.0,
"prompt_remaining": 88.0,
"flex_used": 3.0,
"flex_limit": 10.0,
"flex_remaining": 7.0,
"allowed_models_count": 82
}
}));
let payload = provider_key_status_snapshot_payload(&key, "windsurf");
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("windsurf quota windows should exist");
let daily = windows
.iter()
.filter_map(Value::as_object)
.find(|window| window.get("code") == Some(&json!("daily")))
.expect("daily quota window should exist");
let weekly = windows
.iter()
.filter_map(Value::as_object)
.find(|window| window.get("code") == Some(&json!("weekly")))
.expect("weekly quota window should exist");
assert_eq!(quota.get("provider_type"), Some(&json!("windsurf")));
assert_eq!(quota.get("code"), Some(&json!("ok")));
assert_eq!(quota.get("plan_type"), Some(&json!("Pro")));
assert_eq!(quota.get("usage_ratio"), Some(&json!(0.6)));
assert_eq!(quota.get("reset_at"), Some(&json!(1_778_100_000u64)));
assert_eq!(daily.get("remaining_ratio"), Some(&json!(0.4)));
assert_eq!(daily.get("used_ratio"), Some(&json!(0.6)));
assert_eq!(daily.get("reset_seconds"), Some(&json!(32_754u64)));
assert_eq!(weekly.get("remaining_ratio"), Some(&json!(0.65)));
assert_eq!(weekly.get("used_ratio"), Some(&json!(0.35)));
assert_eq!(weekly.get("reset_seconds"), Some(&json!(532_754u64)));
assert_eq!(quota.get("allowed_models_count"), Some(&json!(82)));
}
#[test]
fn provider_key_status_snapshot_payload_treats_windsurf_rate_limit_as_cooldown() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"windsurf": {
"updated_at": 1_778_067_246u64,
"daily_remaining_percent": 80.0,
"rate_limit": {
"limited": true,
"retry_after_ms": 60_001u64,
"message": "slow down"
},
"last_error": "slow down"
}
}));
let payload = provider_key_status_snapshot_payload(&key, "windsurf");
let quota = payload
.get("quota")
.and_then(Value::as_object)
.expect("quota snapshot should be object");
let rate_window = quota
.get("windows")
.and_then(Value::as_array)
.and_then(|windows| {
windows
.iter()
.filter_map(Value::as_object)
.find(|window| window.get("code") == Some(&json!("rate_limit")))
})
.expect("rate limit window should exist");
assert_eq!(quota.get("code"), Some(&json!("cooldown")));
assert_eq!(quota.get("exhausted"), Some(&json!(false)));
assert_eq!(quota.get("reset_seconds"), Some(&json!(61u64)));
assert_eq!(rate_window.get("is_exhausted"), Some(&json!(false)));
assert_eq!(rate_window.get("reset_seconds"), Some(&json!(61u64)));
}
#[test]
fn provider_key_status_snapshot_payload_keeps_windsurf_capacity_probe_without_retry_after_ok() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"windsurf": {
"updated_at": 1_778_067_246u64,
"daily_remaining_percent": 100.0,
"weekly_remaining_percent": 100.0,
"rate_limit": {
"limited": true,
"has_capacity": false,
"messages_remaining": 0.0,
"max_messages": 100.0
}
}
}));
let payload = provider_key_status_snapshot_payload(&key, "windsurf");
let quota = payload
.get("quota")
.and_then(Value::as_object)
.expect("quota snapshot should be object");
let has_rate_limit_window =
quota
.get("windows")
.and_then(Value::as_array)
.is_some_and(|windows| {
windows
.iter()
.filter_map(Value::as_object)
.any(|window| window.get("code") == Some(&json!("rate_limit")))
});
assert_eq!(quota.get("code"), Some(&json!("ok")));
assert_eq!(quota.get("label"), Some(&Value::Null));
assert_eq!(quota.get("exhausted"), Some(&json!(false)));
assert_eq!(
payload.pointer("/quota/rate_limit/limited"),
Some(&json!(true))
);
assert!(!has_rate_limit_window);
}
#[test]
fn provider_key_status_snapshot_payload_refreshes_stale_windsurf_cooldown_when_probe_has_capacity(
) {
let mut key = sample_catalog_key();
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "windsurf",
"code": "cooldown",
"label": "冷却中",
"exhausted": false,
"windows": [
{
"code": "daily",
"unit": "percent",
"label": "",
"scope": "account",
"remaining_ratio": 0.99,
"is_exhausted": false
},
{
"code": "rate_limit",
"unit": "count",
"label": "速率",
"scope": "account",
"is_exhausted": false,
"reset_seconds": null
}
],
"rate_limit": {
"limited": true,
"has_capacity": true,
"messages_remaining": -1,
"max_messages": -1
}
}
}));
key.upstream_metadata = Some(json!({
"windsurf": {
"updated_at": 1_778_067_246u64,
"daily_remaining_percent": 99.0,
"weekly_remaining_percent": 100.0,
"allowed_models_count": 118,
"rate_limit": {
"limited": true,
"has_capacity": true,
"messages_remaining": -1,
"max_messages": -1
}
}
}));
let payload = provider_key_status_snapshot_payload(&key, "windsurf");
let quota = payload
.get("quota")
.and_then(Value::as_object)
.expect("quota snapshot should be object");
let has_rate_limit_window =
quota
.get("windows")
.and_then(Value::as_array)
.is_some_and(|windows| {
windows
.iter()
.filter_map(Value::as_object)
.any(|window| window.get("code") == Some(&json!("rate_limit")))
});
assert_eq!(quota.get("code"), Some(&json!("ok")));
assert_eq!(quota.get("label"), Some(&Value::Null));
assert_eq!(quota.get("allowed_models_count"), Some(&json!(118u64)));
assert!(!has_rate_limit_window);
}
#[test]
fn provider_key_status_snapshot_payload_marks_windsurf_banned_and_quarantined_blocking() {
let mut banned_key = sample_catalog_key();
banned_key.upstream_metadata = Some(json!({
"windsurf": {
"updated_at": 1_778_067_246u64,
"banned": true,
"reason": "forbidden"
}
}));
let banned_payload = provider_key_status_snapshot_payload(&banned_key, "windsurf");
assert_eq!(
banned_payload.pointer("/quota/code"),
Some(&json!("banned"))
);
assert_eq!(
banned_payload.pointer("/quota/exhausted"),
Some(&json!(true))
);
assert_eq!(
banned_payload.pointer("/account/code"),
Some(&json!("account_banned"))
);
assert_eq!(
banned_payload.pointer("/account/blocked"),
Some(&json!(true))
);
let mut quarantined_key = sample_catalog_key();
quarantined_key.upstream_metadata = Some(json!({
"windsurf": {
"updated_at": 1_778_067_246u64,
"quarantined": true
}
}));
let quarantined_payload =
provider_key_status_snapshot_payload(&quarantined_key, "windsurf");
assert_eq!(
quarantined_payload.pointer("/quota/code"),
Some(&json!("quarantined"))
);
assert_eq!(
quarantined_payload.pointer("/quota/exhausted"),
Some(&json!(true))
);
assert_eq!(
quarantined_payload.pointer("/account/code"),
Some(&json!("account_quarantined"))
);
assert_eq!(
quarantined_payload.pointer("/account/blocked"),
Some(&json!(true))
);
}
#[test]
fn provider_key_status_snapshot_payload_preserves_existing_materialized_quota_snapshot() {
let mut key = sample_catalog_key();
@@ -38,7 +38,8 @@ pub(crate) use self::external_models::OFFICIAL_EXTERNAL_MODEL_PROVIDERS;
pub(crate) use self::normalize::{
deserialize_optional_json_patch, deserialize_optional_string_list_patch, ip_rules_allow,
json_ip_rules_allow, normalize_feature_settings, normalize_ip_rules, normalize_json_array,
normalize_json_object, normalize_string_list, parse_json_ip_rules,
normalize_json_object, normalize_string_list, normalize_user_self_feature_settings_update,
parse_json_ip_rules,
};
pub(crate) use self::payloads::{
InternalGatewayAuthContextRequest, InternalGatewayExecuteRequest,
@@ -54,6 +54,7 @@ pub(crate) fn normalize_feature_settings(value: Option<Value>) -> Result<Option<
Value::Null => Ok(None),
Value::Object(ref mut settings) => {
normalize_chat_pii_redaction_feature_settings(settings)?;
normalize_notification_push_service_feature_settings(settings)?;
if settings.is_empty() {
Ok(None)
} else {
@@ -64,6 +65,41 @@ pub(crate) fn normalize_feature_settings(value: Option<Value>) -> Result<Option<
}
}
pub(crate) fn normalize_user_self_feature_settings_update(
value: Option<Value>,
current: Option<Value>,
) -> Result<Option<Value>, String> {
let mut normalized = normalize_feature_settings(value)?;
let current_notification_push_service = current
.and_then(|value| match value {
Value::Object(mut settings) => settings.remove("notification_push_service"),
_ => None,
})
.and_then(|value| {
let mut wrapper = Map::new();
wrapper.insert("notification_push_service".to_string(), value);
normalize_notification_push_service_feature_settings(&mut wrapper)
.ok()
.and_then(|_| wrapper.remove("notification_push_service"))
});
match (&mut normalized, current_notification_push_service) {
(Some(Value::Object(settings)), Some(value)) => {
settings.insert("notification_push_service".to_string(), value);
}
(Some(Value::Object(settings)), None) => {
settings.remove("notification_push_service");
}
(None, Some(value)) => {
let mut settings = Map::new();
settings.insert("notification_push_service".to_string(), value);
normalized = Some(Value::Object(settings));
}
_ => {}
}
Ok(normalized)
}
pub(crate) fn normalize_ip_rules(
values: Option<Vec<String>>,
) -> Result<Option<Vec<String>>, String> {
@@ -326,9 +362,47 @@ fn normalize_chat_pii_redaction_feature_object(
Ok(())
}
fn normalize_notification_push_service_feature_settings(
settings: &mut Map<String, Value>,
) -> Result<(), String> {
let Some(value) = settings.get_mut("notification_push_service") else {
return Ok(());
};
match value {
Value::Null => {
settings.remove("notification_push_service");
Ok(())
}
Value::Object(feature) => {
normalize_notification_push_service_feature_object(feature)?;
if feature.is_empty() {
settings.remove("notification_push_service");
}
Ok(())
}
_ => Err("notification_push_service 必须是对象".to_string()),
}
}
fn normalize_notification_push_service_feature_object(
feature: &mut Map<String, Value>,
) -> Result<(), String> {
for key in ["enabled"] {
if let Some(value) = feature.get(key) {
if !value.is_boolean() {
return Err(format!("notification_push_service.{key} 必须是布尔值"));
}
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{ip_rules_allow, json_ip_rules_allow, normalize_ip_rules, parse_json_ip_rules};
use super::{
ip_rules_allow, json_ip_rules_allow, normalize_feature_settings, normalize_ip_rules,
normalize_user_self_feature_settings_update, parse_json_ip_rules,
};
use serde_json::json;
use std::net::{IpAddr, Ipv4Addr};
@@ -358,6 +432,41 @@ mod tests {
);
}
#[test]
fn normalize_feature_settings_accepts_notification_push_service_permission() {
let normalized = normalize_feature_settings(Some(json!({
"notification_push_service": {"enabled": true}
})))
.expect("feature settings should normalize")
.expect("feature settings should remain set");
assert_eq!(
normalized["notification_push_service"]["enabled"],
json!(true)
);
}
#[test]
fn user_self_feature_update_preserves_notification_push_permission() {
let normalized = normalize_user_self_feature_settings_update(
Some(json!({
"chat_pii_redaction": {"enabled": true, "inject_model_instruction": false},
"notification_push_service": {"enabled": false}
})),
Some(json!({
"notification_push_service": {"enabled": true}
})),
)
.expect("feature settings should normalize")
.expect("feature settings should remain set");
assert_eq!(
normalized["notification_push_service"]["enabled"],
json!(true)
);
assert_eq!(normalized["chat_pii_redaction"]["enabled"], json!(true));
}
#[test]
fn ip_rules_allow_applies_allow_rules_and_deny_overrides() {
let rules = vec![
@@ -254,6 +254,11 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
| (Some("system_manage"), http::Method::PUT, Some("config_set"))
| (Some("system_manage"), http::Method::PUT, Some("email_template_set"))
| (Some("system_manage"), http::Method::POST, Some("email_template_preview"))
| (
Some("system_manage"),
http::Method::POST,
Some("important_notification_test"),
)
| (
Some("provider_models_manage"),
http::Method::POST,
@@ -0,0 +1,890 @@
use crate::admin_api::AdminAppState;
use crate::bark_push::{read_bark_push_config, send_bark_push, BarkPushConfig};
use crate::email_delivery::{
read_smtp_delivery_config, send_smtp_email, ComposedEmail, SmtpDeliveryConfig,
};
use crate::handlers::shared::{system_config_bool, system_config_string};
use crate::server_chan_push::{
read_server_chan_push_config, send_server_chan_push, ServerChanPushConfig,
};
use crate::{AppState, GatewayError};
use axum::body::Bytes;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tracing::warn;
pub(crate) const IMPORTANT_NOTIFICATION_ENABLED_KEY: &str = "module.important_notification.enabled";
pub(crate) const LEGACY_NOTIFICATION_EMAIL_ENABLED_KEY: &str = "module.notification_email.enabled";
pub(crate) const IMPORTANT_NOTIFICATION_EMAIL_ENABLED_KEY: &str =
"module.important_notification.email_enabled";
pub(crate) const IMPORTANT_NOTIFICATION_EMAIL_RECIPIENTS_KEY: &str =
"module.important_notification.email_recipients";
pub(crate) const IMPORTANT_NOTIFICATION_DEFAULT_CHANNEL_KEY: &str =
"module.important_notification.default_channel";
pub(crate) const IMPORTANT_NOTIFICATION_ITEMS_KEY: &str = "module.important_notification.items";
pub(crate) const PROVIDER_QUOTA_ALERT_ITEM_KEY: &str = "provider_quota_alert";
#[derive(Debug, Clone)]
pub(crate) struct ImportantNotification {
pub(crate) title: String,
pub(crate) markdown_body: String,
pub(crate) text_body: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ImportantNotificationChannelFilter {
All,
Email,
ServerChan,
Bark,
}
#[derive(Debug, Clone)]
struct ImportantNotificationConfig {
module_enabled: bool,
email_enabled: bool,
email_recipients: Vec<String>,
default_channel: ImportantNotificationChannelFilter,
items: Vec<ImportantNotificationItemConfig>,
server_chan: ServerChanPushConfig,
bark: BarkPushConfig,
}
#[derive(Debug, Clone)]
struct ImportantNotificationItemConfig {
key: String,
name: String,
enabled: bool,
channel: Option<ImportantNotificationChannelFilter>,
title_template: Option<String>,
markdown_template: Option<String>,
text_template: Option<String>,
user_email_enabled: bool,
}
#[derive(Debug, Clone, Copy)]
struct NotificationChannelReadiness {
email: bool,
server_chan: bool,
bark: bool,
}
#[derive(Debug, Clone, Serialize)]
pub(crate) struct ImportantNotificationChannelReport {
pub(crate) channel: &'static str,
pub(crate) success: bool,
pub(crate) message: String,
}
#[derive(Debug, Clone, Serialize)]
pub(crate) struct ImportantNotificationDeliveryReport {
pub(crate) success: bool,
pub(crate) channels: Vec<ImportantNotificationChannelReport>,
}
#[derive(Debug, Deserialize)]
struct ImportantNotificationTestRequest {
#[serde(default)]
channel: Option<String>,
#[serde(default)]
item_key: Option<String>,
}
pub(crate) async fn important_notification_module_enabled(
state: &AppState,
) -> Result<bool, GatewayError> {
let canonical = state
.read_system_config_json_value(IMPORTANT_NOTIFICATION_ENABLED_KEY)
.await?;
if canonical.is_some() {
return Ok(system_config_bool(canonical.as_ref(), false));
}
let legacy = state
.read_system_config_json_value(LEGACY_NOTIFICATION_EMAIL_ENABLED_KEY)
.await?;
Ok(system_config_bool(legacy.as_ref(), false))
}
pub(crate) async fn important_notification_configured(
state: &AppState,
) -> Result<bool, GatewayError> {
let config = read_important_notification_config(state).await?;
important_notification_has_configured_channel(state, &config).await
}
pub(crate) async fn important_notification_dispatch_ready_for_item(
state: &AppState,
item_key: &str,
) -> Result<bool, GatewayError> {
let config = read_important_notification_config(state).await?;
if !config.module_enabled {
return Ok(false);
}
if let Some(item) = find_notification_item(&config, item_key) {
if !item.enabled {
return Ok(false);
}
}
let readiness = read_notification_channel_readiness(state, &config).await?;
Ok(channel_filter_has_ready_channel(
notification_item_channel_filter(&config, item_key),
readiness,
))
}
pub(crate) async fn send_important_notification(
state: &AppState,
notification: ImportantNotification,
) -> Result<ImportantNotificationDeliveryReport, GatewayError> {
send_important_notification_with_filter(
state,
notification,
ImportantNotificationChannelFilter::All,
)
.await
}
pub(crate) async fn send_important_notification_for_item(
state: &AppState,
item_key: &str,
notification: ImportantNotification,
variables: &[(&str, String)],
) -> Result<ImportantNotificationDeliveryReport, GatewayError> {
dispatch_important_notification(state, Some(item_key), notification, variables, None, false)
.await
}
pub(crate) async fn send_important_notification_with_filter(
state: &AppState,
notification: ImportantNotification,
channel_filter: ImportantNotificationChannelFilter,
) -> Result<ImportantNotificationDeliveryReport, GatewayError> {
dispatch_important_notification(state, None, notification, &[], Some(channel_filter), false)
.await
}
pub(crate) async fn send_user_important_notification_email(
state: &AppState,
item_key: &str,
user_email: &str,
notification: ImportantNotification,
variables: &[(&str, String)],
) -> Result<ImportantNotificationDeliveryReport, GatewayError> {
let config = read_important_notification_config(state).await?;
if !config.module_enabled {
return Ok(single_report("module", false, "通知服务未启用"));
}
let Some(item) = find_notification_item(&config, item_key) else {
return Ok(single_report("item", false, "通知项未定义"));
};
if !item.enabled {
return Ok(single_report("item", false, "通知项未启用"));
}
if !item.user_email_enabled {
return Ok(single_report("user_email", false, "通知项未启用用户邮件"));
}
let notification = apply_notification_item_template(Some(item), notification, variables);
let smtp_config = match read_smtp_delivery_config(state).await? {
Some(config) => config,
None => return Ok(single_report("user_email", false, "SMTP 配置不完整")),
};
let user_email = user_email.trim();
if user_email.is_empty() {
return Ok(single_report("user_email", false, "用户邮箱为空"));
}
match send_single_email_notification(smtp_config, user_email, &notification).await {
Ok(()) => Ok(single_report("user_email", true, "用户邮件通知已发送")),
Err(err) => {
warn!(error = ?err, user_email = %user_email, "failed to send user notification email");
Ok(single_report(
"user_email",
false,
format!("用户邮件通知发送失败: {err:?}"),
))
}
}
}
async fn important_notification_has_configured_channel(
state: &AppState,
config: &ImportantNotificationConfig,
) -> Result<bool, GatewayError> {
let readiness = read_notification_channel_readiness(state, config).await?;
if channel_filter_has_ready_channel(config.default_channel, readiness) {
return Ok(true);
}
Ok(config.items.iter().any(|item| {
item.enabled
&& channel_filter_has_ready_channel(
item.channel.unwrap_or(config.default_channel),
readiness,
)
}))
}
async fn read_notification_channel_readiness(
state: &AppState,
config: &ImportantNotificationConfig,
) -> Result<NotificationChannelReadiness, GatewayError> {
let smtp_config = read_smtp_delivery_config(state).await?;
Ok(NotificationChannelReadiness {
email: config.email_enabled && !config.email_recipients.is_empty() && smtp_config.is_some(),
server_chan: config.server_chan.enabled && config.server_chan.send_key.is_some(),
bark: config.bark.enabled && config.bark.device_key.is_some(),
})
}
fn channel_filter_has_ready_channel(
filter: ImportantNotificationChannelFilter,
readiness: NotificationChannelReadiness,
) -> bool {
match filter {
ImportantNotificationChannelFilter::All => {
readiness.email || readiness.server_chan || readiness.bark
}
ImportantNotificationChannelFilter::Email => readiness.email,
ImportantNotificationChannelFilter::ServerChan => readiness.server_chan,
ImportantNotificationChannelFilter::Bark => readiness.bark,
}
}
async fn dispatch_important_notification(
state: &AppState,
item_key: Option<&str>,
notification: ImportantNotification,
variables: &[(&str, String)],
channel_override: Option<ImportantNotificationChannelFilter>,
bypass_enable_checks: bool,
) -> Result<ImportantNotificationDeliveryReport, GatewayError> {
let config = read_important_notification_config(state).await?;
if !bypass_enable_checks && !config.module_enabled {
return Ok(single_report("module", false, "通知服务未启用"));
}
let item = item_key.and_then(|key| find_notification_item(&config, key));
if !bypass_enable_checks && item.is_some_and(|item| !item.enabled) {
return Ok(single_report("item", false, "通知项未启用"));
}
let notification = apply_notification_item_template(item, notification, variables);
let channel_filter = channel_override.unwrap_or_else(|| {
item.and_then(|item| item.channel)
.unwrap_or(config.default_channel)
});
let mut reports = Vec::new();
if matches!(
channel_filter,
ImportantNotificationChannelFilter::All | ImportantNotificationChannelFilter::Email
) {
maybe_send_email_notification(
state,
&config,
&notification,
bypass_enable_checks,
&mut reports,
)
.await;
}
if matches!(
channel_filter,
ImportantNotificationChannelFilter::All | ImportantNotificationChannelFilter::ServerChan
) {
maybe_send_server_chan_notification(
state,
&config,
&notification,
bypass_enable_checks,
&mut reports,
)
.await;
}
if matches!(
channel_filter,
ImportantNotificationChannelFilter::All | ImportantNotificationChannelFilter::Bark
) {
maybe_send_bark_notification(
state,
&config,
&notification,
bypass_enable_checks,
&mut reports,
)
.await;
}
if reports.is_empty() {
reports.push(ImportantNotificationChannelReport {
channel: "none",
success: false,
message: "未启用可用的通知通道".to_string(),
});
}
let success = reports.iter().any(|report| report.success);
Ok(ImportantNotificationDeliveryReport {
success,
channels: reports,
})
}
pub(crate) async fn build_important_notification_test_payload(
state: &AdminAppState<'_>,
request_body: Option<&Bytes>,
) -> Result<Value, GatewayError> {
let request = match request_body.filter(|body| !body.is_empty()) {
Some(body) => serde_json::from_slice::<ImportantNotificationTestRequest>(body).unwrap_or(
ImportantNotificationTestRequest {
channel: None,
item_key: None,
},
),
None => ImportantNotificationTestRequest {
channel: None,
item_key: None,
},
};
let filter = request.channel.as_deref().and_then(parse_channel_filter);
let item_key = request
.item_key
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
let variables = vec![
("provider_name", "示例 Provider".to_string()),
("provider_id", "provider-demo".to_string()),
("total_available", "8.0000".to_string()),
("threshold_amount", "10.0000".to_string()),
("user_email", "user@example.com".to_string()),
("balance", "1.0000".to_string()),
];
let report = dispatch_important_notification(
state.app(),
item_key,
ImportantNotification {
title: "Aether 通知服务测试".to_string(),
markdown_body: "这是一条来自 Aether 的通知服务测试。".to_string(),
text_body: "这是一条来自 Aether 的通知服务测试。".to_string(),
},
&variables,
filter,
true,
)
.await?;
Ok(json!({
"success": report.success,
"message": if report.success { "测试通知已发送" } else { "测试通知发送失败" },
"channels": report.channels,
}))
}
async fn read_important_notification_config(
state: &AppState,
) -> Result<ImportantNotificationConfig, GatewayError> {
let module_enabled = important_notification_module_enabled(state).await?;
let email_enabled = state
.read_system_config_json_value(IMPORTANT_NOTIFICATION_EMAIL_ENABLED_KEY)
.await?;
let email_recipients = state
.read_system_config_json_value(IMPORTANT_NOTIFICATION_EMAIL_RECIPIENTS_KEY)
.await?;
let default_channel = state
.read_system_config_json_value(IMPORTANT_NOTIFICATION_DEFAULT_CHANNEL_KEY)
.await?;
let items = state
.read_system_config_json_value(IMPORTANT_NOTIFICATION_ITEMS_KEY)
.await?;
Ok(ImportantNotificationConfig {
module_enabled,
email_enabled: system_config_bool(email_enabled.as_ref(), false),
email_recipients: parse_recipient_list(email_recipients.as_ref()),
default_channel: default_channel
.as_ref()
.and_then(|value| value.as_str())
.and_then(parse_channel_filter)
.unwrap_or(ImportantNotificationChannelFilter::All),
items: parse_notification_items(items.as_ref()),
server_chan: read_server_chan_push_config(state).await?,
bark: read_bark_push_config(state).await?,
})
}
fn parse_channel_filter(raw: &str) -> Option<ImportantNotificationChannelFilter> {
match raw.trim().to_ascii_lowercase().as_str() {
"all" => Some(ImportantNotificationChannelFilter::All),
"email" => Some(ImportantNotificationChannelFilter::Email),
"server_chan" | "serverchan" | "serve_chan" => {
Some(ImportantNotificationChannelFilter::ServerChan)
}
"bark" => Some(ImportantNotificationChannelFilter::Bark),
"global" | "" => None,
_ => None,
}
}
fn parse_notification_items(value: Option<&Value>) -> Vec<ImportantNotificationItemConfig> {
let Some(Value::Array(items)) = value else {
return default_notification_items();
};
items
.iter()
.filter_map(parse_notification_item)
.collect::<Vec<_>>()
}
fn parse_notification_item(value: &Value) -> Option<ImportantNotificationItemConfig> {
let item = value.as_object()?;
let key = item.get("key")?.as_str()?.trim();
if key.is_empty() {
return None;
}
let name = item
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(key);
Some(ImportantNotificationItemConfig {
key: key.to_string(),
name: name.to_string(),
enabled: item.get("enabled").and_then(Value::as_bool).unwrap_or(true),
channel: item
.get("channel")
.and_then(Value::as_str)
.and_then(parse_channel_filter),
title_template: optional_non_empty_string(item.get("title_template")),
markdown_template: optional_non_empty_string(item.get("markdown_template")),
text_template: optional_non_empty_string(item.get("text_template")),
user_email_enabled: item
.get("user_email_enabled")
.and_then(Value::as_bool)
.unwrap_or(false),
})
}
fn default_notification_items() -> Vec<ImportantNotificationItemConfig> {
vec![
ImportantNotificationItemConfig {
key: PROVIDER_QUOTA_ALERT_ITEM_KEY.to_string(),
name: "号池额度不足".to_string(),
enabled: true,
channel: None,
title_template: None,
markdown_template: None,
text_template: None,
user_email_enabled: false,
},
ImportantNotificationItemConfig {
key: "provider_pool_abnormal".to_string(),
name: "号池异常".to_string(),
enabled: true,
channel: None,
title_template: Some("号池异常:{provider_name}".to_string()),
markdown_template: Some(
"号池 `{provider_name}` 出现异常,请检查服务状态。".to_string(),
),
text_template: Some("号池 {provider_name} 出现异常,请检查服务状态。".to_string()),
user_email_enabled: false,
},
ImportantNotificationItemConfig {
key: "user_balance_low".to_string(),
name: "用户余额不足".to_string(),
enabled: true,
channel: Some(ImportantNotificationChannelFilter::Email),
title_template: Some("余额不足提醒".to_string()),
markdown_template: Some("你的账户余额已低于提醒阈值,请及时处理。".to_string()),
text_template: Some("你的账户余额已低于提醒阈值,请及时处理。".to_string()),
user_email_enabled: true,
},
]
}
fn optional_non_empty_string(value: Option<&Value>) -> Option<String> {
value
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn find_notification_item<'a>(
config: &'a ImportantNotificationConfig,
item_key: &str,
) -> Option<&'a ImportantNotificationItemConfig> {
let item_key = item_key.trim();
config.items.iter().find(|item| item.key == item_key)
}
fn notification_item_channel_filter(
config: &ImportantNotificationConfig,
item_key: &str,
) -> ImportantNotificationChannelFilter {
find_notification_item(config, item_key)
.and_then(|item| item.channel)
.unwrap_or(config.default_channel)
}
fn apply_notification_item_template(
item: Option<&ImportantNotificationItemConfig>,
notification: ImportantNotification,
variables: &[(&str, String)],
) -> ImportantNotification {
let Some(item) = item else {
return notification;
};
let title = render_template(
item.title_template.as_deref(),
&notification.title,
&notification,
variables,
);
let markdown_body = render_template(
item.markdown_template.as_deref(),
&notification.markdown_body,
&notification,
variables,
);
let text_body = render_template(
item.text_template.as_deref(),
&notification.text_body,
&notification,
variables,
);
ImportantNotification {
title,
markdown_body,
text_body,
}
}
fn render_template(
template: Option<&str>,
fallback: &str,
notification: &ImportantNotification,
variables: &[(&str, String)],
) -> String {
let mut rendered = template
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(fallback)
.to_string();
rendered = rendered
.replace("{title}", &notification.title)
.replace("{body}", &notification.markdown_body)
.replace("{text_body}", &notification.text_body);
for (key, value) in variables {
rendered = rendered.replace(&format!("{{{}}}", key.trim()), value);
}
rendered
}
async fn maybe_send_email_notification(
state: &AppState,
config: &ImportantNotificationConfig,
notification: &ImportantNotification,
bypass_channel_toggle: bool,
reports: &mut Vec<ImportantNotificationChannelReport>,
) {
if !bypass_channel_toggle && !config.email_enabled {
return;
}
if config.email_recipients.is_empty() {
reports.push(ImportantNotificationChannelReport {
channel: "email",
success: false,
message: "未配置邮件收件人".to_string(),
});
return;
}
let smtp_config = match read_smtp_delivery_config(state).await {
Ok(Some(config)) => config,
Ok(None) => {
reports.push(ImportantNotificationChannelReport {
channel: "email",
success: false,
message: "SMTP 配置不完整".to_string(),
});
return;
}
Err(err) => {
reports.push(ImportantNotificationChannelReport {
channel: "email",
success: false,
message: format!("读取 SMTP 配置失败: {err:?}"),
});
return;
}
};
let mut sent = 0usize;
let mut failed = 0usize;
for recipient in &config.email_recipients {
match send_single_email_notification(smtp_config.clone(), recipient, notification).await {
Ok(()) => sent += 1,
Err(err) => {
failed += 1;
warn!(
error = ?err,
recipient = %recipient,
"failed to send important notification email"
);
}
}
}
reports.push(ImportantNotificationChannelReport {
channel: "email",
success: sent > 0,
message: if failed == 0 {
format!("邮件通知已发送给 {sent} 个收件人")
} else {
format!("邮件通知成功 {sent} 个,失败 {failed}")
},
});
}
async fn send_single_email_notification(
smtp_config: SmtpDeliveryConfig,
recipient: &str,
notification: &ImportantNotification,
) -> Result<(), GatewayError> {
send_smtp_email(
smtp_config,
ComposedEmail {
to_email: recipient.to_string(),
subject: notification.title.clone(),
html_body: build_notification_html(notification),
text_body: notification.text_body.clone(),
},
)
.await
}
async fn maybe_send_server_chan_notification(
state: &AppState,
config: &ImportantNotificationConfig,
notification: &ImportantNotification,
bypass_channel_toggle: bool,
reports: &mut Vec<ImportantNotificationChannelReport>,
) {
if !bypass_channel_toggle && !config.server_chan.enabled {
return;
}
if config.server_chan.send_key.is_none() {
reports.push(ImportantNotificationChannelReport {
channel: "server_chan",
success: false,
message: "未配置 Server 酱 SendKey".to_string(),
});
return;
};
match send_server_chan_push(
state,
&config.server_chan,
&notification.title,
&notification.markdown_body,
)
.await
{
Ok(()) => reports.push(ImportantNotificationChannelReport {
channel: "server_chan",
success: true,
message: "Server 酱通知已发送".to_string(),
}),
Err(err) => {
warn!(error = ?err, "failed to send server chan important notification");
reports.push(ImportantNotificationChannelReport {
channel: "server_chan",
success: false,
message: format!("Server 酱通知发送失败: {err:?}"),
});
}
}
}
async fn maybe_send_bark_notification(
state: &AppState,
config: &ImportantNotificationConfig,
notification: &ImportantNotification,
bypass_channel_toggle: bool,
reports: &mut Vec<ImportantNotificationChannelReport>,
) {
if !bypass_channel_toggle && !config.bark.enabled {
return;
}
if config.bark.device_key.is_none() {
reports.push(ImportantNotificationChannelReport {
channel: "bark",
success: false,
message: "未配置 Bark Device Key".to_string(),
});
return;
};
match send_bark_push(
state,
&config.bark,
&notification.title,
&notification.markdown_body,
)
.await
{
Ok(()) => reports.push(ImportantNotificationChannelReport {
channel: "bark",
success: true,
message: "Bark 通知已发送".to_string(),
}),
Err(err) => {
warn!(error = ?err, "failed to send bark important notification");
reports.push(ImportantNotificationChannelReport {
channel: "bark",
success: false,
message: format!("Bark 通知发送失败: {err:?}"),
});
}
}
}
fn single_report(
channel: &'static str,
success: bool,
message: impl Into<String>,
) -> ImportantNotificationDeliveryReport {
ImportantNotificationDeliveryReport {
success,
channels: vec![ImportantNotificationChannelReport {
channel,
success,
message: message.into(),
}],
}
}
fn parse_recipient_list(value: Option<&Value>) -> Vec<String> {
let mut recipients = Vec::new();
match value {
Some(Value::Array(items)) => {
for item in items {
if let Some(raw) = item.as_str() {
push_recipient_parts(&mut recipients, raw);
}
}
}
Some(Value::String(raw)) => push_recipient_parts(&mut recipients, raw),
_ => {}
}
recipients.sort();
recipients.dedup();
recipients
}
fn push_recipient_parts(recipients: &mut Vec<String>, raw: &str) {
for item in raw
.split([',', ';', '\n', '\r'])
.map(str::trim)
.filter(|value| !value.is_empty())
{
recipients.push(item.to_string());
}
}
fn build_notification_html(notification: &ImportantNotification) -> String {
format!(
"<!doctype html><html><body><h2>{}</h2><pre style=\"font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;white-space:pre-wrap;line-height:1.6\">{}</pre></body></html>",
escape_html(&notification.title),
escape_html(&notification.text_body),
)
}
fn escape_html(value: &str) -> String {
value
.replace('&', "&amp;")
.replace('<', "&lt;")
.replace('>', "&gt;")
.replace('"', "&quot;")
.replace('\'', "&#39;")
}
#[cfg(test)]
mod tests {
use super::{
apply_notification_item_template, parse_channel_filter, parse_notification_items,
parse_recipient_list, ImportantNotification, ImportantNotificationChannelFilter,
};
use serde_json::json;
#[test]
fn parse_recipient_list_accepts_arrays_and_delimiters() {
assert_eq!(
parse_recipient_list(Some(&json!([
"ops@example.com, admin@example.com",
"ops@example.com"
]))),
vec![
"admin@example.com".to_string(),
"ops@example.com".to_string()
]
);
}
#[test]
fn parse_notification_items_reads_channel_and_user_email_flag() {
let items = parse_notification_items(Some(&json!([
{
"key": "user_balance_low",
"name": "用户余额不足",
"enabled": true,
"channel": "email",
"title_template": "余额提醒",
"markdown_template": "{user_email}: {balance}",
"user_email_enabled": true
}
])));
assert_eq!(items.len(), 1);
assert_eq!(items[0].key, "user_balance_low");
assert_eq!(
items[0].channel,
Some(ImportantNotificationChannelFilter::Email)
);
assert!(items[0].user_email_enabled);
}
#[test]
fn parse_channel_filter_accepts_bark() {
assert_eq!(
parse_channel_filter("bark"),
Some(ImportantNotificationChannelFilter::Bark)
);
}
#[test]
fn item_template_renders_fallback_and_variables() {
let items = parse_notification_items(Some(&json!([
{
"key": "provider_quota_alert",
"name": "号池额度不足",
"title_template": "额度提醒:{provider_name}",
"markdown_template": "{body}\n剩余:{total_available}",
"text_template": "{text_body}\n剩余:{total_available}"
}
])));
let rendered = apply_notification_item_template(
Some(&items[0]),
ImportantNotification {
title: "默认标题".to_string(),
markdown_body: "默认正文".to_string(),
text_body: "默认文本".to_string(),
},
&[
("provider_name", "示例 Provider".to_string()),
("total_available", "8.0000".to_string()),
],
);
assert_eq!(rendered.title, "额度提醒:示例 Provider");
assert_eq!(rendered.markdown_body, "默认正文\n剩余:8.0000");
assert_eq!(rendered.text_body, "默认文本\n剩余:8.0000");
}
}
+4
View File
@@ -30,6 +30,7 @@ mod api;
mod async_task;
mod audit;
mod auth;
mod bark_push;
mod cache;
mod client_session_affinity;
mod clock;
@@ -37,6 +38,7 @@ mod constants;
mod control;
mod data;
mod dispatch;
mod email_delivery;
mod error;
mod execution_runtime;
mod executor;
@@ -46,6 +48,7 @@ mod handlers;
mod headers;
mod hooks;
mod image_capabilities;
mod important_notification;
mod log_ids;
mod maintenance;
pub(crate) mod middleware;
@@ -62,6 +65,7 @@ mod roles;
mod router;
mod routing;
mod scheduler;
mod server_chan_push;
mod state;
mod system_features;
mod task_runtime;
+1
View File
@@ -497,6 +497,7 @@ impl GatewayUsageArgs {
fn to_config(&self) -> UsageRuntimeConfig {
UsageRuntimeConfig {
enabled: true,
queue_terminal_events: true,
stream_key: self.queue_stream_key.trim().to_string(),
consumer_group: self.queue_group.trim().to_string(),
dlq_stream_key: self.queue_dlq_stream_key.trim().to_string(),
+13 -12
View File
@@ -7,8 +7,8 @@ pub(crate) use runtime::{
ensure_provider_key_pool_scores_for_keys, inspect_proxy_upgrade_rollout,
list_admin_cleanup_run_records, perform_account_self_check_once,
perform_oauth_token_refresh_once, perform_pool_quota_probe_once, perform_provider_checkin_once,
pool_quota_probe_target_count, preview_manual_usage_cleanup, rebuild_admin_stats_once,
record_completed_cleanup_run, record_proxy_upgrade_traffic_success,
perform_provider_quota_alert_once, pool_quota_probe_target_count, preview_manual_usage_cleanup,
rebuild_admin_stats_once, record_completed_cleanup_run, record_proxy_upgrade_traffic_success,
restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node,
run_admin_system_cleanup_once, run_manual_usage_cleanup_once, skip_proxy_upgrade_rollout_node,
spawn_account_self_check_worker, spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
@@ -16,16 +16,17 @@ pub(crate) use runtime::{
spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_replenish_for_request, spawn_pool_quota_probe_worker,
spawn_pool_score_rebuild_worker, spawn_provider_checkin_worker,
spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_usage_counter_flush_worker,
spawn_wallet_daily_usage_aggregation_worker, start_admin_request_body_cleanup_task,
start_admin_system_purge_task, start_manual_usage_cleanup_task, start_proxy_upgrade_rollout,
AccountSelfCheckRunSummary, AdminCleanupRunRecord, AdminCleanupTaskKind,
AdminStatsRebuildSummary, AdminSystemCleanupSummary, ManualUsageCleanupError,
ManualUsageCleanupMode, ManualUsageCleanupOptions, OAuthTokenRefreshRunSummary,
PoolQuotaProbeRunSummary, PoolQuotaProbeWorkerConfig, ProviderCheckinRunSummary,
spawn_provider_quota_alert_worker, spawn_proxy_node_metrics_cleanup_worker,
spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker,
spawn_request_candidate_cleanup_worker, spawn_stats_aggregation_worker,
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
spawn_usage_counter_flush_worker, spawn_wallet_daily_usage_aggregation_worker,
start_admin_request_body_cleanup_task, start_admin_system_purge_task,
start_manual_usage_cleanup_task, start_proxy_upgrade_rollout, AccountSelfCheckRunSummary,
AdminCleanupRunRecord, AdminCleanupTaskKind, AdminStatsRebuildSummary,
AdminSystemCleanupSummary, ManualUsageCleanupError, ManualUsageCleanupMode,
ManualUsageCleanupOptions, OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary,
PoolQuotaProbeWorkerConfig, ProviderCheckinRunSummary, ProviderQuotaAlertRunSummary,
ProxyUpgradeRolloutCancelSummary, ProxyUpgradeRolloutConflictClearSummary,
ProxyUpgradeRolloutNodeActionSummary, ProxyUpgradeRolloutProbeConfig,
ProxyUpgradeRolloutSkippedRestoreSummary, ProxyUpgradeRolloutStatus,
@@ -26,6 +26,8 @@ mod pool_quota_probe;
mod pool_score_rebuild;
#[path = "runtime/provider_checkin.rs"]
mod provider_checkin;
#[path = "runtime/provider_quota_alert.rs"]
mod provider_quota_alert;
#[path = "runtime/proxy_node_metrics_cleanup.rs"]
mod proxy_node_metrics_cleanup;
#[path = "runtime/proxy_node_staleness.rs"]
@@ -85,6 +87,9 @@ pub(crate) use pool_score_rebuild::{
PoolScoreRebuildRunSummary, PoolScoreRebuildWorkerConfig,
};
pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary};
pub(crate) use provider_quota_alert::{
perform_provider_quota_alert_once, ProviderQuotaAlertRunSummary,
};
use proxy_node_metrics_cleanup::*;
use proxy_node_staleness::*;
use proxy_upgrade_rollout::*;
@@ -138,6 +143,8 @@ const USAGE_COUNTER_DELTA_CLEANUP_INTERVAL: Duration = Duration::from_secs(60);
const USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE: usize = 5_000;
const USAGE_COUNTER_DELTA_RETENTION_SECS: u64 = 7 * 24 * 60 * 60;
const PROVIDER_CHECKIN_CONCURRENCY: usize = 3;
const PROVIDER_QUOTA_ALERT_CONCURRENCY: usize = 3;
const PROVIDER_QUOTA_ALERT_INTERVAL: Duration = Duration::from_secs(5);
const PROVIDER_CHECKIN_DEFAULT_TIME: &str = "01:05";
const REQUEST_CANDIDATE_CLEANUP_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
const STATS_DAILY_AGGREGATION_HOUR: u32 = 0;
@@ -738,8 +738,21 @@ pub(crate) fn spawn_account_self_check_worker(
let mut interval = tokio::time::interval(config.scan_interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut deferred_since = None;
loop {
interval.tick().await;
if state
.data
.should_defer_maintenance_for_database_pool_pressure(&mut deferred_since)
{
debug!(
event_name = "maintenance_worker_deferred",
log_type = "ops",
worker = "account_self_check",
"gateway account self-check deferred because database pool has no idle reserve"
);
continue;
}
if let Err(err) = perform_account_self_check_once_with_config(&state, config).await {
warn!(
error = ?err,
@@ -393,8 +393,21 @@ pub(crate) fn spawn_pool_score_rebuild_worker(
}
let mut interval = tokio::time::interval(config.interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let mut deferred_since = None;
loop {
interval.tick().await;
if state
.data
.should_defer_maintenance_for_database_pool_pressure(&mut deferred_since)
{
debug!(
event_name = "maintenance_worker_deferred",
log_type = "ops",
worker = "pool_score_rebuild",
"gateway pool score rebuild deferred because database pool has no idle reserve"
);
continue;
}
match perform_pool_score_rebuild_once_with_config(&state, config).await {
Ok(summary) if summary.scores_upserted > 0 => {
info!(
@@ -0,0 +1,486 @@
use std::collections::HashMap;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
use futures_util::stream::{self, StreamExt};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tracing::warn;
use crate::admin_api::{
admin_provider_ops_local_action_response, store_admin_provider_ops_balance_cache, AdminAppState,
};
use crate::important_notification::{
important_notification_dispatch_ready_for_item, send_important_notification_for_item,
ImportantNotification, PROVIDER_QUOTA_ALERT_ITEM_KEY,
};
use crate::{AppState, GatewayError};
use super::PROVIDER_QUOTA_ALERT_CONCURRENCY;
const PROVIDER_QUOTA_ALERT_STATE_PREFIX: &str = "provider_ops:quota_alert:";
const PROVIDER_QUOTA_ALERT_DEFAULT_FETCH_INTERVAL_SECS: u64 = 30;
const PROVIDER_QUOTA_ALERT_MIN_FETCH_INTERVAL_SECS: u64 = 30;
const PROVIDER_QUOTA_ALERT_REPEAT_COOLDOWN_SECS: u64 = 24 * 60 * 60;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ProviderQuotaAlertRunSummary {
pub(crate) checked: usize,
pub(crate) alerted: usize,
pub(crate) skipped: usize,
pub(crate) failed: usize,
}
#[derive(Debug, Clone)]
struct ProviderQuotaAlertTarget {
provider: StoredProviderCatalogProvider,
config: ProviderQuotaAlertConfig,
}
#[derive(Debug, Clone, Copy)]
struct ProviderQuotaAlertConfig {
enabled: bool,
threshold_amount: f64,
fetch_interval_seconds: u64,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
struct ProviderQuotaAlertRuntimeState {
#[serde(default)]
last_checked_at: u64,
#[serde(default)]
last_available: Option<f64>,
#[serde(default)]
below_threshold: bool,
#[serde(default)]
last_notified_at: Option<u64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ProviderQuotaAlertStatus {
Checked,
Alerted,
Skipped,
Failed,
}
pub(crate) async fn perform_provider_quota_alert_once(
state: &AppState,
) -> Result<ProviderQuotaAlertRunSummary, GatewayError> {
if !state.has_provider_catalog_data_reader() {
return Ok(ProviderQuotaAlertRunSummary {
checked: 0,
alerted: 0,
skipped: 0,
failed: 0,
});
}
if !important_notification_dispatch_ready_for_item(state, PROVIDER_QUOTA_ALERT_ITEM_KEY).await?
{
return Ok(ProviderQuotaAlertRunSummary {
checked: 0,
alerted: 0,
skipped: 0,
failed: 0,
});
}
let now_unix_secs = now_unix_secs();
let targets = select_provider_quota_alert_targets(state, now_unix_secs).await?;
if targets.is_empty() {
return Ok(ProviderQuotaAlertRunSummary {
checked: 0,
alerted: 0,
skipped: 0,
failed: 0,
});
}
let provider_ids = targets
.iter()
.map(|target| target.provider.id.clone())
.collect::<Vec<_>>();
let mut endpoints_by_provider = HashMap::<String, Vec<StoredProviderCatalogEndpoint>>::new();
for endpoint in state
.list_provider_catalog_endpoints_by_provider_ids(&provider_ids)
.await?
{
endpoints_by_provider
.entry(endpoint.provider_id.clone())
.or_default()
.push(endpoint);
}
let mut results = stream::iter(targets.into_iter().map(|target| {
let state = state.clone();
let provider_id = target.provider.id.clone();
let endpoints = endpoints_by_provider
.get(&provider_id)
.cloned()
.unwrap_or_default();
async move { run_provider_quota_alert_for_provider(&state, target, endpoints).await }
}))
.buffer_unordered(PROVIDER_QUOTA_ALERT_CONCURRENCY);
let mut summary = ProviderQuotaAlertRunSummary {
checked: 0,
alerted: 0,
skipped: 0,
failed: 0,
};
while let Some(status) = results.next().await {
match status {
ProviderQuotaAlertStatus::Checked => summary.checked += 1,
ProviderQuotaAlertStatus::Alerted => {
summary.checked += 1;
summary.alerted += 1;
}
ProviderQuotaAlertStatus::Skipped => summary.skipped += 1,
ProviderQuotaAlertStatus::Failed => summary.failed += 1,
}
}
Ok(summary)
}
async fn select_provider_quota_alert_targets(
state: &AppState,
now_unix_secs: u64,
) -> Result<Vec<ProviderQuotaAlertTarget>, GatewayError> {
let providers = state
.list_provider_catalog_providers(true)
.await?
.into_iter()
.filter_map(|provider| {
let config = provider_quota_alert_config(&provider)?;
(config.enabled).then_some(ProviderQuotaAlertTarget { provider, config })
})
.collect::<Vec<_>>();
let mut due = Vec::new();
for target in providers {
let runtime = read_quota_alert_runtime_state(state, &target.provider.id).await;
let last_checked_at = runtime
.as_ref()
.map(|state| state.last_checked_at)
.unwrap_or(0);
if now_unix_secs.saturating_sub(last_checked_at)
>= target
.config
.fetch_interval_seconds
.max(PROVIDER_QUOTA_ALERT_MIN_FETCH_INTERVAL_SECS)
{
due.push(target);
}
}
Ok(due)
}
async fn run_provider_quota_alert_for_provider(
state: &AppState,
target: ProviderQuotaAlertTarget,
endpoints: Vec<StoredProviderCatalogEndpoint>,
) -> ProviderQuotaAlertStatus {
let provider_id = target.provider.id.clone();
let admin_state = AdminAppState::new(state);
let payload = admin_provider_ops_local_action_response(
&admin_state,
&provider_id,
Some(&target.provider),
&endpoints,
"query_balance",
None,
)
.await;
store_admin_provider_ops_balance_cache(&admin_state, &provider_id, &payload).await;
let now_unix_secs = now_unix_secs();
let payload_status_success = payload.get("status").and_then(Value::as_str) == Some("success");
let Some(total_available) = extract_total_available(&payload) else {
warn!(
provider_id = %provider_id,
payload = %payload,
"provider quota alert skipped because balance payload has no total_available"
);
write_checked_runtime_state_without_balance(state, &provider_id, now_unix_secs).await;
return if payload_status_success {
ProviderQuotaAlertStatus::Skipped
} else {
ProviderQuotaAlertStatus::Failed
};
};
let previous = read_quota_alert_runtime_state(state, &provider_id).await;
let should_notify = provider_quota_alert_should_notify(
now_unix_secs,
total_available,
target.config.threshold_amount,
previous.as_ref(),
);
let mut next = ProviderQuotaAlertRuntimeState {
last_checked_at: now_unix_secs,
last_available: Some(total_available),
below_threshold: total_available <= target.config.threshold_amount,
last_notified_at: previous.and_then(|state| state.last_notified_at),
};
if should_notify {
let notification = build_provider_quota_alert_notification(
&target.provider,
total_available,
target.config.threshold_amount,
);
let variables = provider_quota_alert_notification_variables(
&target.provider,
total_available,
target.config.threshold_amount,
);
let report = send_important_notification_for_item(
state,
PROVIDER_QUOTA_ALERT_ITEM_KEY,
notification,
&variables,
)
.await;
let delivered = match &report {
Ok(report) if report.success => true,
Ok(report) => {
warn!(
provider_id = %provider_id,
report = ?report,
"provider quota alert notification did not reach any channel"
);
false
}
Err(err) => {
warn!(
provider_id = %provider_id,
error = ?err,
"provider quota alert notification failed"
);
false
}
};
if delivered {
next.last_notified_at = Some(now_unix_secs);
write_quota_alert_runtime_state(state, &provider_id, &next).await;
return ProviderQuotaAlertStatus::Alerted;
}
write_quota_alert_runtime_state(state, &provider_id, &next).await;
return ProviderQuotaAlertStatus::Failed;
}
write_quota_alert_runtime_state(state, &provider_id, &next).await;
ProviderQuotaAlertStatus::Checked
}
fn provider_quota_alert_config(
provider: &StoredProviderCatalogProvider,
) -> Option<ProviderQuotaAlertConfig> {
let quota_alert = provider
.config
.as_ref()
.and_then(Value::as_object)
.and_then(|config| config.get("provider_ops"))
.and_then(Value::as_object)
.and_then(|provider_ops| provider_ops.get("quota_alert"))
.and_then(Value::as_object)?;
let enabled = quota_alert
.get("enabled")
.and_then(Value::as_bool)
.unwrap_or(false);
let threshold_amount = quota_alert
.get("threshold_amount")
.and_then(value_as_f64)
.filter(|value| value.is_finite() && *value >= 0.0)
.unwrap_or(0.0);
let fetch_interval_seconds = quota_alert
.get("fetch_interval_seconds")
.and_then(Value::as_u64)
.unwrap_or(PROVIDER_QUOTA_ALERT_DEFAULT_FETCH_INTERVAL_SECS)
.max(PROVIDER_QUOTA_ALERT_MIN_FETCH_INTERVAL_SECS);
Some(ProviderQuotaAlertConfig {
enabled,
threshold_amount,
fetch_interval_seconds,
})
}
fn provider_quota_alert_should_notify(
now_unix_secs: u64,
total_available: f64,
threshold_amount: f64,
previous: Option<&ProviderQuotaAlertRuntimeState>,
) -> bool {
if total_available > threshold_amount {
return false;
}
let Some(previous) = previous else {
return true;
};
if !previous.below_threshold {
return true;
}
previous
.last_notified_at
.map(|last| now_unix_secs.saturating_sub(last) >= PROVIDER_QUOTA_ALERT_REPEAT_COOLDOWN_SECS)
.unwrap_or(true)
}
fn extract_total_available(payload: &Value) -> Option<f64> {
if payload.get("status").and_then(Value::as_str) != Some("success") {
return None;
}
payload
.get("data")
.and_then(|data| data.get("total_available"))
.and_then(value_as_f64)
.filter(|value| value.is_finite())
}
fn value_as_f64(value: &Value) -> Option<f64> {
value.as_f64().or_else(|| {
value
.as_str()
.and_then(|raw| raw.trim().parse::<f64>().ok())
})
}
fn build_provider_quota_alert_notification(
provider: &StoredProviderCatalogProvider,
total_available: f64,
threshold_amount: f64,
) -> ImportantNotification {
let title = format!("提供商额度提醒:{}", provider.name);
let body = format!(
"提供商 `{}` 当前剩余额度为 `{:.4}`,已低于或等于提醒阈值 `{:.4}`。\n\nProvider ID: `{}`",
provider.name, total_available, threshold_amount, provider.id
);
let text_body = format!(
"提供商 {} 当前剩余额度为 {:.4},已低于或等于提醒阈值 {:.4}。\n\nProvider ID: {}",
provider.name, total_available, threshold_amount, provider.id
);
ImportantNotification {
title,
markdown_body: body,
text_body,
}
}
fn provider_quota_alert_notification_variables(
provider: &StoredProviderCatalogProvider,
total_available: f64,
threshold_amount: f64,
) -> Vec<(&'static str, String)> {
vec![
("provider_name", provider.name.clone()),
("provider_id", provider.id.clone()),
("total_available", format!("{total_available:.4}")),
("threshold_amount", format!("{threshold_amount:.4}")),
]
}
async fn read_quota_alert_runtime_state(
state: &AppState,
provider_id: &str,
) -> Option<ProviderQuotaAlertRuntimeState> {
let key = provider_quota_alert_state_key(provider_id);
state
.runtime_kv_get(&key)
.await
.ok()
.flatten()
.and_then(|raw| serde_json::from_str::<ProviderQuotaAlertRuntimeState>(&raw).ok())
}
async fn write_checked_runtime_state_without_balance(
state: &AppState,
provider_id: &str,
now_unix_secs: u64,
) {
let mut next = read_quota_alert_runtime_state(state, provider_id)
.await
.unwrap_or_default();
next.last_checked_at = now_unix_secs;
write_quota_alert_runtime_state(state, provider_id, &next).await;
}
async fn write_quota_alert_runtime_state(
state: &AppState,
provider_id: &str,
runtime_state: &ProviderQuotaAlertRuntimeState,
) {
let Ok(serialized) = serde_json::to_string(runtime_state) else {
return;
};
if let Err(err) = state
.runtime_state()
.kv_set(
&provider_quota_alert_state_key(provider_id),
serialized,
None,
)
.await
{
warn!(
error = %err,
provider_id,
"failed to write provider quota alert runtime state"
);
}
}
fn provider_quota_alert_state_key(provider_id: &str) -> String {
format!("{PROVIDER_QUOTA_ALERT_STATE_PREFIX}{provider_id}")
}
fn now_unix_secs() -> u64 {
chrono::Utc::now().timestamp().max(0) as u64
}
#[cfg(test)]
mod tests {
use super::{
provider_quota_alert_should_notify, ProviderQuotaAlertRuntimeState,
PROVIDER_QUOTA_ALERT_REPEAT_COOLDOWN_SECS,
};
#[test]
fn quota_alert_notifies_on_first_drop_and_after_cooldown() {
assert!(provider_quota_alert_should_notify(100, 3.0, 5.0, None));
assert!(provider_quota_alert_should_notify(
100,
3.0,
5.0,
Some(&ProviderQuotaAlertRuntimeState {
below_threshold: false,
..ProviderQuotaAlertRuntimeState::default()
})
));
assert!(!provider_quota_alert_should_notify(
100,
3.0,
5.0,
Some(&ProviderQuotaAlertRuntimeState {
below_threshold: true,
last_notified_at: Some(90),
..ProviderQuotaAlertRuntimeState::default()
})
));
assert!(provider_quota_alert_should_notify(
100 + PROVIDER_QUOTA_ALERT_REPEAT_COOLDOWN_SECS,
3.0,
5.0,
Some(&ProviderQuotaAlertRuntimeState {
below_threshold: true,
last_notified_at: Some(100),
..ProviderQuotaAlertRuntimeState::default()
})
));
}
#[test]
fn quota_alert_does_not_notify_above_threshold() {
assert!(!provider_quota_alert_should_notify(100, 6.0, 5.0, None));
}
}
@@ -1,7 +1,8 @@
use std::sync::Arc;
use std::time::Instant;
use chrono::Utc;
use tracing::warn;
use tracing::{debug, warn};
use crate::data::GatewayDataState;
use crate::AppState;
@@ -10,22 +11,22 @@ use super::{
cleanup_processed_usage_counter_deltas_once, duration_until_next_daily_run,
duration_until_next_db_maintenance_run, duration_until_next_stats_aggregation_run,
duration_until_next_stats_hourly_aggregation_run, maintenance_timezone, parse_hhmm_time,
perform_oauth_token_refresh_once, provider_checkin_schedule, run_audit_cleanup_once,
run_db_maintenance_once, run_gemini_file_mapping_cleanup_once, run_pending_cleanup_once,
run_pool_monitor_once, run_provider_checkin_once, run_proxy_node_metrics_cleanup_once,
run_proxy_node_stale_cleanup_once, run_proxy_upgrade_rollout_once,
run_request_candidate_cleanup_once, run_stats_aggregation_once,
perform_oauth_token_refresh_once, perform_provider_quota_alert_once, provider_checkin_schedule,
run_audit_cleanup_once, run_db_maintenance_once, run_gemini_file_mapping_cleanup_once,
run_pending_cleanup_once, run_pool_monitor_once, run_provider_checkin_once,
run_proxy_node_metrics_cleanup_once, run_proxy_node_stale_cleanup_once,
run_proxy_upgrade_rollout_once, run_request_candidate_cleanup_once, run_stats_aggregation_once,
run_stats_hourly_aggregation_once, run_usage_cleanup_once, run_usage_counter_flush_once,
run_wallet_daily_usage_aggregation_once, AUDIT_LOG_CLEANUP_INTERVAL,
GEMINI_FILE_MAPPING_CLEANUP_INTERVAL, OAUTH_TOKEN_REFRESH_INTERVAL, PENDING_CLEANUP_INTERVAL,
POOL_MONITOR_INTERVAL, PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_METRICS_CLEANUP_HOUR,
PROXY_NODE_METRICS_CLEANUP_MINUTE, PROXY_NODE_STALE_SWEEP_INTERVAL,
PROXY_UPGRADE_ROLLOUT_INTERVAL, REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR,
USAGE_CLEANUP_MINUTE, USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE,
USAGE_COUNTER_DELTA_CLEANUP_INTERVAL, USAGE_COUNTER_DELTA_RETENTION_SECS,
USAGE_COUNTER_FLUSH_BATCH_SIZE, USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT,
USAGE_COUNTER_FLUSH_INTERVAL, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
POOL_MONITOR_INTERVAL, PROVIDER_CHECKIN_DEFAULT_TIME, PROVIDER_QUOTA_ALERT_INTERVAL,
PROXY_NODE_METRICS_CLEANUP_HOUR, PROXY_NODE_METRICS_CLEANUP_MINUTE,
PROXY_NODE_STALE_SWEEP_INTERVAL, PROXY_UPGRADE_ROLLOUT_INTERVAL,
REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE, USAGE_COUNTER_DELTA_CLEANUP_INTERVAL,
USAGE_COUNTER_DELTA_RETENTION_SECS, USAGE_COUNTER_FLUSH_BATCH_SIZE,
USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT, USAGE_COUNTER_FLUSH_INTERVAL,
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
const STATS_DAILY_CATCH_UP_BURST_LIMIT: usize = 14;
@@ -46,6 +47,37 @@ fn log_maintenance_worker_failure(
);
}
fn should_defer_for_database_pressure(
data: &GatewayDataState,
worker: &'static str,
deferred_since: &mut Option<Instant>,
) -> bool {
let Some(summary) = data.database_pool_summary() else {
*deferred_since = None;
return false;
};
if !GatewayDataState::should_defer_maintenance_for_pool_pressure_state(
GatewayDataState::database_pool_summary_under_maintenance_pressure(&summary),
deferred_since,
) {
return false;
}
debug!(
event_name = "maintenance_worker_deferred",
log_type = "ops",
worker,
driver = %summary.driver,
checked_out = summary.checked_out,
pool_size = summary.pool_size,
idle = summary.idle,
max_connections = summary.max_connections,
usage_rate = summary.usage_rate,
"gateway maintenance worker deferred because database pool has no idle reserve"
);
true
}
pub(crate) fn spawn_audit_cleanup_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
@@ -177,8 +209,19 @@ pub(crate) fn spawn_usage_counter_flush_worker(
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut last_delta_cleanup = tokio::time::Instant::now();
let mut usage_counter_flush_deferred_since = None;
let mut usage_counter_delta_cleanup_deferred_since = None;
loop {
if should_defer_for_database_pressure(
&data,
"usage_counter_flush",
&mut usage_counter_flush_deferred_since,
) {
interval.tick().await;
continue;
}
let mut batches = 0_usize;
while batches < USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT {
match run_usage_counter_flush_once(&data, USAGE_COUNTER_FLUSH_BATCH_SIZE).await {
@@ -197,7 +240,18 @@ pub(crate) fn spawn_usage_counter_flush_worker(
}
if last_delta_cleanup.elapsed() >= USAGE_COUNTER_DELTA_CLEANUP_INTERVAL {
if let Err(err) = cleanup_processed_usage_counter_deltas_once(
if should_defer_for_database_pressure(
&data,
"usage_counter_delta_cleanup",
&mut usage_counter_delta_cleanup_deferred_since,
) {
debug!(
event_name = "maintenance_worker_deferred",
log_type = "ops",
worker = "usage_counter_delta_cleanup",
"gateway maintenance worker deferred cleanup under database pressure"
);
} else if let Err(err) = cleanup_processed_usage_counter_deltas_once(
&data,
USAGE_COUNTER_DELTA_RETENTION_SECS,
USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE,
@@ -254,6 +308,34 @@ pub(crate) fn spawn_provider_checkin_worker(
}))
}
pub(crate) fn spawn_provider_quota_alert_worker(
state: AppState,
) -> Option<tokio::task::JoinHandle<()>> {
if !state.has_provider_catalog_data_reader() {
return None;
}
Some(tokio::spawn(async move {
let mut interval = tokio::time::interval(PROVIDER_QUOTA_ALERT_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut deferred_since = None;
loop {
interval.tick().await;
if should_defer_for_database_pressure(
&state.data,
"provider_quota_alert",
&mut deferred_since,
) {
continue;
}
if let Err(err) = perform_provider_quota_alert_once(&state).await {
log_maintenance_worker_failure("provider_quota_alert", "tick", &err);
}
}
}))
}
pub(crate) fn spawn_oauth_token_refresh_worker(
state: AppState,
) -> Option<tokio::task::JoinHandle<()>> {
@@ -268,8 +350,16 @@ pub(crate) fn spawn_oauth_token_refresh_worker(
let mut interval = tokio::time::interval(OAUTH_TOKEN_REFRESH_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut deferred_since = None;
loop {
interval.tick().await;
if should_defer_for_database_pressure(
&state.data,
"oauth_token_refresh",
&mut deferred_since,
) {
continue;
}
if let Err(err) = perform_oauth_token_refresh_once(&state).await {
log_maintenance_worker_failure("oauth_token_refresh", "tick", &err);
}
@@ -314,8 +404,12 @@ pub(crate) fn spawn_pending_cleanup_worker(
let mut interval = tokio::time::interval(PENDING_CLEANUP_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut deferred_since = None;
loop {
interval.tick().await;
if should_defer_for_database_pressure(&data, "pending_cleanup", &mut deferred_since) {
continue;
}
if let Err(err) = run_pending_cleanup_once(&data).await {
log_maintenance_worker_failure("pending_cleanup", "tick", &err);
}
@@ -337,8 +431,16 @@ pub(crate) fn spawn_proxy_node_stale_cleanup_worker(
let mut interval = tokio::time::interval(PROXY_NODE_STALE_SWEEP_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut deferred_since = None;
loop {
interval.tick().await;
if should_defer_for_database_pressure(
&data,
"proxy_node_stale_cleanup",
&mut deferred_since,
) {
continue;
}
if let Err(err) = run_proxy_node_stale_cleanup_once(&data).await {
log_maintenance_worker_failure("proxy_node_stale_cleanup", "tick", &err);
}
@@ -387,8 +489,16 @@ pub(crate) fn spawn_proxy_upgrade_rollout_worker(
let mut interval = tokio::time::interval(PROXY_UPGRADE_ROLLOUT_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut deferred_since = None;
loop {
interval.tick().await;
if should_defer_for_database_pressure(
&state.data,
"proxy_upgrade_rollout",
&mut deferred_since,
) {
continue;
}
if let Err(err) = run_proxy_upgrade_rollout_once(&state).await {
log_maintenance_worker_failure("proxy_upgrade_rollout", "tick", &err);
}
+135 -1
View File
@@ -426,6 +426,7 @@ mod tests {
keys: Arc<Mutex<Vec<StoredProviderCatalogKey>>>,
transports: Arc<HashMap<(String, String, String), GatewayProviderTransportSnapshot>>,
execution_results: Arc<Mutex<VecDeque<ExecutionResult>>>,
executed_plans: Arc<Mutex<Vec<ExecutionPlan>>>,
cached_models: Arc<Mutex<HashMap<(String, String), Vec<Value>>>>,
}
@@ -443,6 +444,7 @@ mod tests {
keys: Arc::new(Mutex::new(keys)),
transports: Arc::new(transports),
execution_results: Arc::new(Mutex::new(VecDeque::from(execution_results))),
executed_plans: Arc::new(Mutex::new(Vec::new())),
cached_models: Arc::new(Mutex::new(HashMap::new())),
}
}
@@ -482,8 +484,12 @@ mod tests {
async fn execute_model_fetch_execution_plan(
&self,
_plan: &ExecutionPlan,
plan: &ExecutionPlan,
) -> Result<ExecutionResult, String> {
self.executed_plans
.lock()
.expect("executed plans mutex")
.push(plan.clone());
self.execution_results
.lock()
.expect("execution result mutex")
@@ -911,6 +917,134 @@ mod tests {
);
}
#[tokio::test]
async fn model_fetch_fetches_windsurf_model_configs_and_persists_allowed_models() {
let provider = sample_provider("provider-windsurf", "windsurf");
let endpoint = StoredProviderCatalogEndpoint::new(
"endpoint-windsurf-chat".to_string(),
"provider-windsurf".to_string(),
"openai:chat".to_string(),
None,
None,
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://server.codeium.com".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build");
let key = sample_key(
"key-windsurf",
"provider-windsurf",
"api_key",
&["openai:chat"],
);
let mut transport = sample_transport(
"windsurf",
"provider-windsurf",
"endpoint-windsurf-chat",
"key-windsurf",
"openai:chat",
"api_key",
Some(r#"{"provider_type":"windsurf"}"#),
);
transport.endpoint.base_url = "https://server.codeium.com".to_string();
transport.key.decrypted_api_key = "devin-session-token$abc".to_string();
let state = TestState::new(
vec![provider],
vec![endpoint],
vec![key],
HashMap::from([(
(
"provider-windsurf".to_string(),
"endpoint-windsurf-chat".to_string(),
"key-windsurf".to_string(),
),
transport,
)]),
vec![execution_result(json!({
"clientModelConfigs": [
{
"modelUid": "claude-sonnet-4-6",
"label": "Claude Sonnet 4.6",
"provider": "anthropic",
"supportsImages": true,
"creditMultiplier": 4
},
{
"modelUid": "gpt-5.4",
"label": "GPT-5.4",
"provider": "openai"
}
],
"defaultOverrideModelConfig": {
"modelUid": "claude-sonnet-4-6"
}
}))],
);
let summary = perform_model_fetch_once_with_state(&state)
.await
.expect("fetch should succeed");
assert_eq!(summary.succeeded, 1);
let plans = state.executed_plans.lock().expect("executed plans mutex");
assert_eq!(plans.len(), 1);
assert_eq!(
plans[0].url,
"https://server.codeium.com/exa.api_server_pb.ApiServerService/GetCascadeModelConfigs"
);
assert_eq!(plans[0].method, "POST");
assert_eq!(plans[0].provider_api_format, "windsurf:model_configs");
assert_eq!(
plans[0]
.body
.json_body
.as_ref()
.and_then(|body| body.get("metadata"))
.and_then(|metadata| metadata.get("apiKey")),
Some(&json!("devin-session-token$abc"))
);
drop(plans);
let updated = state.key("key-windsurf");
assert_eq!(
updated.allowed_models,
Some(json!(["claude-sonnet-4-6", "gpt-5.4"]))
);
assert_eq!(
updated
.upstream_metadata
.as_ref()
.and_then(|value| value.get("windsurf"))
.and_then(|value| value.get("allowed_models_count")),
Some(&json!(2))
);
assert_eq!(
updated
.upstream_metadata
.as_ref()
.and_then(|value| value.get("windsurf"))
.and_then(|value| value.get("default_model_uid")),
Some(&json!("claude-sonnet-4-6"))
);
let cached = state.cached_models.lock().expect("cache mutex");
let cached_models = cached
.get(&("provider-windsurf".to_string(), "key-windsurf".to_string()))
.expect("cached models should be written");
assert_eq!(
cached_models[0]["api_formats"],
json!(["openai:chat", "openai:responses", "claude:messages"])
);
}
#[tokio::test]
async fn model_fetch_failure_keeps_existing_allowed_models() {
let provider = sample_provider("provider-openai", "openai");
@@ -243,6 +243,35 @@ fn local_scheduler_affinity_target(plan: &ExecutionPlan) -> Option<SchedulerAffi
})
}
async fn local_execution_plan_uses_pool(state: &AppState, plan: &ExecutionPlan) -> bool {
let Ok(Some(transport)) = state
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
.await
else {
return false;
};
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref()).is_some()
}
async fn local_scheduler_affinity_matches_failed_target(
state: &AppState,
plan: &ExecutionPlan,
cached_target: &SchedulerAffinityTarget,
failed_target: &SchedulerAffinityTarget,
) -> bool {
if cached_target == failed_target {
return true;
}
if cached_target.provider_id != failed_target.provider_id
|| cached_target.endpoint_id != failed_target.endpoint_id
{
return false;
}
local_execution_plan_uses_pool(state, plan).await
}
async fn scheduler_cache_affinity_enabled(state: &AppState) -> bool {
match read_scheduler_ordering_config(state).await {
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
@@ -360,7 +389,24 @@ async fn record_attempt_failure_effect(
}
if let Some(cache_key) = local_scheduler_affinity_cache_key(context.report_context) {
let _ = state.remove_scheduler_affinity_cache_entry(&cache_key);
let Some(failed_target) = local_scheduler_affinity_target(context.plan) else {
return;
};
let Some(cached_target) =
state.read_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL)
else {
return;
};
if local_scheduler_affinity_matches_failed_target(
state,
context.plan,
&cached_target,
&failed_target,
)
.await
{
let _ = state.remove_scheduler_affinity_cache_entry(&cache_key);
}
}
}
@@ -448,7 +494,10 @@ async fn record_adaptive_rate_limit_effect(
updated_key.status_snapshot = Some(projection.status_snapshot);
updated_key.updated_at_unix_secs = Some(observed_at_unix_secs);
if let Err(err) = state.update_provider_catalog_key(&updated_key).await {
if let Err(err) = state
.update_provider_catalog_key_runtime_state(&updated_key)
.await
{
warn!(
"gateway orchestration effects: failed to persist adaptive rate-limit projection for provider {} endpoint {} key {}: {:?}",
context.plan.provider_id, context.plan.endpoint_id, context.plan.key_id, err
@@ -496,7 +545,10 @@ async fn record_adaptive_success_effect(
updated_key.status_snapshot = Some(projection.status_snapshot);
updated_key.updated_at_unix_secs = Some(observed_at_unix_secs);
if let Err(err) = state.update_provider_catalog_key(&updated_key).await {
if let Err(err) = state
.update_provider_catalog_key_runtime_state(&updated_key)
.await
{
warn!(
"gateway orchestration effects: failed to persist adaptive success projection for provider {} endpoint {} key {}: {:?}",
context.plan.provider_id, context.plan.endpoint_id, context.plan.key_id, err
@@ -1204,6 +1256,20 @@ mod tests {
.expect("provider should build")
}
fn sample_pool_health_provider() -> StoredProviderCatalogProvider {
sample_health_provider().with_transport_fields(
true,
false,
false,
None,
None,
None,
None,
None,
Some(json!({"pool_advanced": {}})),
)
}
fn sample_health_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"ep-1".to_string(),
@@ -1266,6 +1332,20 @@ mod tests {
)
}
fn pool_health_state() -> AppState {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_pool_health_provider()],
vec![sample_health_endpoint()],
vec![sample_health_key()],
));
AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
}
fn health_state_with_key(key: StoredProviderCatalogKey) -> AppState {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_health_provider()],
@@ -1429,6 +1509,137 @@ mod tests {
.is_some());
}
#[tokio::test]
async fn attempt_failure_keeps_scheduler_affinity_for_non_affinity_candidate() {
let state = AppState::new().expect("gateway state should build");
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
let affinity_target = SchedulerAffinityTarget {
provider_id: "prov-2".to_string(),
endpoint_id: "ep-2".to_string(),
key_id: "key-2".to_string(),
};
state.remember_scheduler_affinity_target(
&cache_key,
affinity_target.clone(),
SCHEDULER_AFFINITY_TTL,
16,
);
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: 524,
classification: LocalFailoverClassification::RetryUpstreamFailure,
}),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(affinity_target)
);
}
#[tokio::test]
async fn attempt_failure_keeps_scheduler_affinity_for_non_pool_sibling_key() {
let state = health_state();
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
let affinity_target = SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-2".to_string(),
};
state.remember_scheduler_affinity_target(
&cache_key,
affinity_target.clone(),
SCHEDULER_AFFINITY_TTL,
16,
);
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: 524,
classification: LocalFailoverClassification::RetryUpstreamFailure,
}),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(affinity_target)
);
}
#[tokio::test]
async fn attempt_failure_invalidates_scheduler_affinity_for_same_pool_candidate() {
let state = pool_health_state();
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
state.remember_scheduler_affinity_target(
&cache_key,
SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-2".to_string(),
},
SCHEDULER_AFFINITY_TTL,
16,
);
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: 524,
classification: LocalFailoverClassification::RetryUpstreamFailure,
}),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
None
);
}
#[tokio::test]
async fn attempt_failure_keeps_scheduler_affinity_for_non_failure_status() {
let state = AppState::new().expect("gateway state should build");
@@ -1546,6 +1757,39 @@ mod tests {
);
}
#[tokio::test]
async fn health_success_keeps_scheduler_affinity_after_health_state_update() {
let state = health_state();
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-1".to_string(),
})
);
}
#[tokio::test]
async fn load_balance_success_does_not_remember_scheduler_affinity_cache() {
let state = AppState::new()
@@ -1985,6 +2229,21 @@ mod tests {
async fn adaptive_rate_limit_effect_updates_adaptive_key_observation() {
let state = adaptive_state();
let plan = sample_plan();
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
let target = SchedulerAffinityTarget {
provider_id: plan.provider_id.clone(),
endpoint_id: plan.endpoint_id.clone(),
key_id: plan.key_id.clone(),
};
state.remember_scheduler_affinity_target(
&cache_key,
target.clone(),
SCHEDULER_AFFINITY_TTL,
16,
);
let initial_epoch = state.scheduler_affinity_epoch();
apply_local_execution_effect(
&state,
@@ -2048,6 +2307,11 @@ mod tests {
.and_then(|value| value.get("enforcement_active")),
Some(&json!(false))
);
assert_eq!(state.scheduler_affinity_epoch(), initial_epoch);
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(target)
);
}
#[tokio::test]
@@ -2178,6 +2442,21 @@ mod tests {
.expect("request candidate should build")],
);
let plan = sample_plan();
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
let target = SchedulerAffinityTarget {
provider_id: plan.provider_id.clone(),
endpoint_id: plan.endpoint_id.clone(),
key_id: plan.key_id.clone(),
};
state.remember_scheduler_affinity_target(
&cache_key,
target.clone(),
SCHEDULER_AFFINITY_TTL,
16,
);
let initial_epoch = state.scheduler_affinity_epoch();
apply_local_execution_effect(
&state,
@@ -2209,5 +2488,10 @@ mod tests {
.and_then(Value::as_str),
Some("high_utilization")
);
assert_eq!(state.scheduler_affinity_epoch(), initial_epoch);
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(target)
);
}
}
@@ -540,7 +540,7 @@ async fn sync_grok_quota_from_report_context(
updated_key.updated_at_unix_secs = Some(now_unix_secs);
Ok(state
.update_provider_catalog_key(&updated_key)
.update_provider_catalog_key_runtime_state(&updated_key)
.await?
.is_some())
}
@@ -870,7 +870,7 @@ async fn sync_codex_quota_from_response_headers(
updated_key.updated_at_unix_secs = Some(now_unix_secs);
let updated = state
.update_provider_catalog_key(&updated_key)
.update_provider_catalog_key_runtime_state(&updated_key)
.await?
.is_some();
if updated {
+27 -3
View File
@@ -50,6 +50,7 @@ pub(crate) struct ProviderKeyAuthSemantics {
credential_kind: ProviderKeyCredentialKind,
runtime_auth_kind: ProviderKeyRuntimeAuthKind,
oauth_managed: bool,
can_refresh_oauth: bool,
}
impl ProviderKeyAuthSemantics {
@@ -66,7 +67,7 @@ impl ProviderKeyAuthSemantics {
}
pub(crate) const fn can_refresh_oauth(self) -> bool {
self.oauth_managed
self.can_refresh_oauth
}
pub(crate) const fn can_export_oauth(self) -> bool {
@@ -115,7 +116,13 @@ fn key_has_auth_type_overrides(key: &StoredProviderCatalogKey) -> bool {
fn provider_uses_bearer_oauth_runtime(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"claude_code" | "codex" | "chatgpt_web" | "gemini_cli" | "antigravity" | "kiro"
"claude_code"
| "codex"
| "chatgpt_web"
| "gemini_cli"
| "antigravity"
| "kiro"
| "windsurf"
)
}
@@ -173,10 +180,12 @@ pub(crate) fn provider_key_auth_semantics(
}
};
let provider_type_normalized = provider_type.trim().to_ascii_lowercase();
ProviderKeyAuthSemantics {
credential_kind,
runtime_auth_kind,
oauth_managed,
can_refresh_oauth: oauth_managed && provider_type_normalized != "windsurf",
}
}
@@ -319,7 +328,6 @@ mod tests {
key.encrypted_auth_config = Some(r#"{"sso_token":"abc"}"#.to_string());
let semantics = provider_key_auth_semantics(&key, "grok");
assert!(semantics.oauth_managed());
assert_eq!(
semantics.credential_kind(),
@@ -331,6 +339,22 @@ mod tests {
);
}
#[test]
fn recognizes_windsurf_oauth_as_bearer_runtime() {
let semantics = provider_key_auth_semantics(&sample_key("oauth"), "windsurf");
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");
+44 -30
View File
@@ -150,11 +150,20 @@ impl FrontdoorUserRpmLimiter {
scope_key: &str,
bucket: u64,
) -> Result<u32, GatewayError> {
if !state.runtime_state.is_memory() {
let raw = state.runtime_state.kv_get(scope_key).await.map_err(|err| {
GatewayError::Internal(format!("frontdoor user rpm runtime read failed: {err}"))
})?;
return Ok(raw.and_then(|value| value.parse::<u32>().ok()).unwrap_or(0));
match state.runtime_state.kv_get(scope_key).await {
Ok(raw) => return Ok(raw.and_then(|value| value.parse::<u32>().ok()).unwrap_or(0)),
Err(err) if !self.config.allow_local_fallback() => {
return Err(GatewayError::Internal(format!(
"frontdoor user rpm runtime read failed: {err}"
)));
}
Err(err) => {
warn!(
error = ?err,
scope_key = %scope_key,
"frontdoor user rpm runtime count read failed; using local fallback"
);
}
}
let counts = self.memory_counts.lock().await;
@@ -187,24 +196,22 @@ impl FrontdoorUserRpmLimiter {
return Ok(FrontdoorUserRpmOutcome::Allowed);
}
if !state.runtime_state.is_memory() {
match self.check_and_consume_runtime(state, &plan).await {
Ok(outcome) => return Ok(outcome),
Err(err) => {
warn!(
error = ?err,
user_rpm_key = %plan.user_rpm_key,
key_rpm_key = %plan.key_rpm_key,
"frontdoor user rpm redis check failed"
);
if self.config.fail_open() {
return Ok(FrontdoorUserRpmOutcome::NotApplicable);
}
if !self.config.allow_local_fallback() {
return Err(GatewayError::Internal(
"frontdoor user rpm runtime backend is unavailable and local fallback is disabled for the current deployment mode".to_string(),
));
}
match self.check_and_consume_runtime(state, &plan).await {
Ok(outcome) => return Ok(outcome),
Err(err) => {
warn!(
error = ?err,
user_rpm_key = %plan.user_rpm_key,
key_rpm_key = %plan.key_rpm_key,
"frontdoor user rpm runtime check failed"
);
if self.config.fail_open() {
return Ok(FrontdoorUserRpmOutcome::NotApplicable);
}
if !self.config.allow_local_fallback() {
return Err(GatewayError::Internal(
"frontdoor user rpm runtime backend is unavailable and local fallback is disabled for the current deployment mode".to_string(),
));
}
}
}
@@ -605,7 +612,7 @@ mod tests {
}
#[tokio::test]
async fn limiter_rejects_missing_shared_runtime_when_local_fallback_disabled() {
async fn limiter_uses_runtime_state_when_local_fallback_disabled() {
let limiter = FrontdoorUserRpmLimiter::new(
FrontdoorUserRpmConfig::new(60, 120, false).with_local_fallback(false),
);
@@ -626,15 +633,22 @@ mod tests {
});
let state = AppState::new().expect("state should build for tests");
let err = limiter
let first = limiter
.check_and_consume(&state, Some(&decision))
.await
.expect_err("missing shared runtime should fail in strict mode");
match err {
crate::GatewayError::Internal(message) => {
assert!(message.contains("requires shared runtime state"));
.expect("runtime check should succeed");
assert_eq!(first, FrontdoorUserRpmOutcome::Allowed);
let second = limiter
.check_and_consume(&state, Some(&decision))
.await
.expect("runtime check should succeed");
match second {
FrontdoorUserRpmOutcome::Rejected(rejection) => {
assert_eq!(rejection.scope, "user");
assert_eq!(rejection.limit, 1);
}
other => panic!("expected internal error, got {other:?}"),
other => panic!("expected rejection, got {other:?}"),
}
}
}
@@ -191,7 +191,7 @@ pub(crate) fn snapshot_local_request_candidate_status(
.map(str::trim)
.filter(|value| !value.is_empty())?;
let metadata = parse_request_candidate_report_context(report_context)?;
let candidate_index = metadata.candidate_index?;
let candidate_index = metadata.candidate_index.unwrap_or(0);
Some(LocalRequestCandidateStatusSnapshot {
candidate_id: candidate_id.to_string(),
+175
View File
@@ -0,0 +1,175 @@
use crate::handlers::shared::{
decrypt_catalog_secret_with_fallbacks, system_config_bool, system_config_string,
};
use crate::{AppState, GatewayError};
use serde_json::Value;
pub(crate) const SERVER_CHAN_PUSH_ENABLED_KEY: &str = "module.server_chan_push.enabled";
pub(crate) const SERVER_CHAN_PUSH_SEND_KEY_KEY: &str = "module.server_chan_push.send_key";
pub(crate) const SERVER_CHAN_PUSH_TEMPLATE_KEY: &str = "module.server_chan_push.template";
pub(crate) const LEGACY_SERVER_CHAN_ENABLED_KEY: &str =
"module.important_notification.server_chan_enabled";
pub(crate) const LEGACY_SERVER_CHAN_SEND_KEY_KEY: &str =
"module.important_notification.server_chan_send_key";
pub(crate) const LEGACY_SERVER_CHAN_TEMPLATE_KEY: &str =
"module.important_notification.server_chan_template";
const SERVER_CHAN_API_BASE: &str = "https://sctapi.ftqq.com";
#[derive(Debug, Clone)]
pub(crate) struct ServerChanPushConfig {
pub(crate) enabled: bool,
pub(crate) send_key: Option<String>,
pub(crate) template: Option<String>,
}
pub(crate) async fn server_chan_push_module_enabled(
state: &AppState,
) -> Result<bool, GatewayError> {
let canonical = state
.read_system_config_json_value(SERVER_CHAN_PUSH_ENABLED_KEY)
.await?;
if canonical.is_some() {
return Ok(system_config_bool(canonical.as_ref(), false));
}
let legacy = state
.read_system_config_json_value(LEGACY_SERVER_CHAN_ENABLED_KEY)
.await?;
Ok(system_config_bool(legacy.as_ref(), false))
}
pub(crate) async fn server_chan_push_configured(state: &AppState) -> Result<bool, GatewayError> {
Ok(read_server_chan_push_config(state)
.await?
.send_key
.is_some())
}
pub(crate) async fn read_server_chan_push_config(
state: &AppState,
) -> Result<ServerChanPushConfig, GatewayError> {
let enabled = server_chan_push_module_enabled(state).await?;
let send_key = read_server_chan_value(
state,
SERVER_CHAN_PUSH_SEND_KEY_KEY,
LEGACY_SERVER_CHAN_SEND_KEY_KEY,
)
.await?
.and_then(|value| system_config_string(Some(&value)))
.map(|value| {
decrypt_catalog_secret_with_fallbacks(state.encryption_key(), &value).unwrap_or(value)
});
let template = read_server_chan_value(
state,
SERVER_CHAN_PUSH_TEMPLATE_KEY,
LEGACY_SERVER_CHAN_TEMPLATE_KEY,
)
.await?
.and_then(|value| system_config_string(Some(&value)));
Ok(ServerChanPushConfig {
enabled,
send_key,
template,
})
}
async fn read_server_chan_value(
state: &AppState,
canonical_key: &str,
legacy_key: &str,
) -> Result<Option<Value>, GatewayError> {
let canonical = state.read_system_config_json_value(canonical_key).await?;
if canonical.is_some() {
return Ok(canonical);
}
state.read_system_config_json_value(legacy_key).await
}
pub(crate) async fn send_server_chan_push(
state: &AppState,
config: &ServerChanPushConfig,
title: &str,
markdown_body: &str,
) -> Result<(), GatewayError> {
let Some(send_key) = config.send_key.as_deref() else {
return Err(GatewayError::Internal(
"未配置 Server 酱 SendKey".to_string(),
));
};
let send_key = send_key.trim();
if send_key.is_empty() {
return Err(GatewayError::Internal(
"Server 酱 SendKey 不能为空".to_string(),
));
}
let desp = render_server_chan_desp(config.template.as_deref(), title, markdown_body);
let url = format!("{SERVER_CHAN_API_BASE}/{send_key}.send");
let response = state
.client
.post(url)
.form(&[("title", title), ("desp", desp.as_str())])
.send()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if !status.is_success() {
return Err(GatewayError::Internal(format!(
"Server 酱返回 HTTP {status}: {text}"
)));
}
if let Ok(payload) = serde_json::from_str::<Value>(&text) {
let code_is_ok = payload
.get("code")
.and_then(|value| {
value
.as_i64()
.map(|code| code == 0)
.or_else(|| value.as_str().map(|code| code.trim() == "0"))
})
.unwrap_or(true);
if !code_is_ok {
return Err(GatewayError::Internal(format!(
"Server 酱返回失败: {payload}"
)));
}
}
Ok(())
}
fn render_server_chan_desp(template: Option<&str>, title: &str, markdown_body: &str) -> String {
match template {
Some(template) if !template.trim().is_empty() => template
.replace("{title}", title)
.replace("{body}", markdown_body),
_ => markdown_body.to_string(),
}
}
#[cfg(test)]
mod tests {
use super::render_server_chan_desp;
#[test]
fn server_chan_desp_uses_template_when_provided() {
let rendered =
render_server_chan_desp(Some("**{title}**\n\n{body}\n\n--end--"), "告警", "原始正文");
assert_eq!(rendered, "**告警**\n\n原始正文\n\n--end--");
}
#[test]
fn server_chan_desp_falls_back_to_markdown_body_for_empty_template() {
assert_eq!(
render_server_chan_desp(None, "告警", "原始正文"),
"原始正文"
);
assert_eq!(
render_server_chan_desp(Some(" "), "告警", "原始正文"),
"原始正文"
);
}
}
+59 -1
View File
@@ -615,6 +615,21 @@ impl AppState {
Ok(updated)
}
pub(crate) async fn update_provider_catalog_key_runtime_state(
&self,
key: &provider_catalog::StoredProviderCatalogKey,
) -> Result<Option<provider_catalog::StoredProviderCatalogKey>, GatewayError> {
let updated = self
.data
.update_provider_catalog_key(key)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated.is_some() {
self.invalidate_provider_health_routing_caches();
}
Ok(updated)
}
pub(crate) async fn update_provider_catalog_key_upstream_metadata(
&self,
key_id: &str,
@@ -806,7 +821,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated {
self.invalidate_provider_routing_caches();
self.invalidate_provider_health_routing_caches();
}
Ok(updated)
}
@@ -930,4 +945,47 @@ mod tests {
.expect("provider transport should exist after update");
assert!(snapshot.provider.keep_priority_on_conversion);
}
#[tokio::test]
async fn provider_catalog_health_update_keeps_scheduler_affinity_cache() {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository)
.with_encryption_key_for_tests("test-encryption-key"),
);
let cache_key = "scheduler_affinity:api-key-1:openai:chat:gpt-5";
let ttl = Duration::from_secs(300);
let target = SchedulerAffinityTarget {
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
};
state.remember_scheduler_affinity_target(cache_key, target.clone(), ttl, 128);
let initial_epoch = state.scheduler_affinity_epoch();
let health_by_format = serde_json::json!({
"openai:chat": {
"last_success_at_unix_secs": 1,
"consecutive_failures": 0
}
});
let updated = state
.update_provider_catalog_key_health_state("key-1", true, Some(&health_by_format), None)
.await
.expect("key health update should succeed");
assert!(updated);
assert_eq!(state.scheduler_affinity_epoch(), initial_epoch);
assert_eq!(
state.read_scheduler_affinity_target(cache_key, ttl),
Some(target)
);
}
}
+14 -11
View File
@@ -49,6 +49,7 @@ use crate::maintenance::spawn_pending_cleanup_worker;
use crate::maintenance::spawn_pool_monitor_worker;
use crate::maintenance::spawn_pool_score_rebuild_worker;
use crate::maintenance::spawn_provider_checkin_worker;
use crate::maintenance::spawn_provider_quota_alert_worker;
use crate::maintenance::spawn_proxy_node_metrics_cleanup_worker;
use crate::maintenance::spawn_proxy_node_stale_cleanup_worker;
use crate::maintenance::spawn_proxy_upgrade_rollout_worker;
@@ -89,24 +90,17 @@ impl AppState {
fn usage_worker_queue_for(
runtime_state: &Arc<RuntimeState>,
) -> Option<Arc<dyn RuntimeQueueStore>> {
if runtime_state.is_redis() {
let queue: Arc<dyn RuntimeQueueStore> = runtime_state.clone();
Some(queue)
} else {
None
}
let queue: Arc<dyn RuntimeQueueStore> = runtime_state.clone();
Some(queue)
}
fn spawn_scheduler_affinity_redis_write(
fn spawn_scheduler_affinity_runtime_write(
&self,
cache_key: &str,
target: &SchedulerAffinityTarget,
ttl: Duration,
epoch: u64,
) {
if self.runtime_state.is_memory() {
return;
}
let Ok(handle) = tokio::runtime::Handle::try_current() else {
return;
};
@@ -578,6 +572,11 @@ impl AppState {
self.invalidate_scheduler_affinity_cache();
}
pub(crate) fn invalidate_provider_health_routing_caches(&self) {
self.data.clear_minimal_candidate_selection_cache();
self.clear_provider_transport_snapshot_cache();
}
pub(crate) fn invalidate_auth_context_cache(&self) {
self.auth_context_cache.clear();
}
@@ -1088,7 +1087,7 @@ impl AppState {
if self.scheduler_affinity_epoch() != epoch {
return false;
}
self.spawn_scheduler_affinity_redis_write(cache_key, &target, ttl, epoch);
self.spawn_scheduler_affinity_runtime_write(cache_key, &target, ttl, epoch);
self.scheduler_affinity_cache.insert_for_epoch(
cache_key.to_string(),
target,
@@ -1210,6 +1209,10 @@ impl AppState {
crate::task_runtime::TASK_KEY_PROVIDER_CHECKIN,
spawn_provider_checkin_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROVIDER_QUOTA_ALERT,
spawn_provider_quota_alert_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_OAUTH_TOKEN_REFRESH,
spawn_oauth_token_refresh_worker(self.clone()),
@@ -36,6 +36,7 @@ pub(crate) const TASK_KEY_PROXY_NODE_METRICS_CLEANUP: &str =
"maintenance.proxy.node.metrics.cleanup";
pub(crate) const TASK_KEY_PROXY_UPGRADE_ROLLOUT: &str = "maintenance.proxy.upgrade.rollout";
pub(crate) const TASK_KEY_PROVIDER_CHECKIN: &str = "maintenance.provider.checkin";
pub(crate) const TASK_KEY_PROVIDER_QUOTA_ALERT: &str = "maintenance.provider.quota_alert";
pub(crate) const TASK_KEY_USAGE_CLEANUP: &str = "maintenance.usage.cleanup";
pub(crate) const TASK_KEY_WALLET_DAILY_USAGE_AGG: &str = "maintenance.wallet.daily.usage.agg";
pub(crate) const TASK_KEY_STATS_DAILY_AGG: &str = "maintenance.stats.daily.agg";
@@ -207,6 +208,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_PROVIDER_QUOTA_ALERT,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_USAGE_CLEANUP,
TaskKind::Scheduled,
@@ -884,15 +884,15 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo
);
assert_eq!(
seen_execution_runtime_request.prompt_cache_key,
"b6741389-8b9e-5c00-bef6-fbce92aee45a"
"bc749eb7-a9e2-5793-8d14-abd659c700b0"
);
assert_eq!(
seen_execution_runtime_request.session_id,
"9fa08f4f14ccba13"
"d1e9b802644e1f52"
);
assert_eq!(
seen_execution_runtime_request.conversation_id,
"9fa08f4f14ccba13"
"d1e9b802644e1f52"
);
assert_eq!(
seen_execution_runtime_request.instructions,
@@ -104,9 +104,16 @@ fn runtime_state_owns_redis_runtime_boundaries() {
let mut violations = Vec::new();
for root in [
"apps/aether-gateway/src",
"apps/aether-tunnel/src",
"crates/aether-admin/src",
"crates/aether-billing/src",
"crates/aether-model-fetch/src",
"crates/aether-provider-pool/src",
"crates/aether-runtime/src",
"crates/aether-task-runtime/src",
"crates/aether-usage-runtime/src",
"crates/aether-provider-transport/src",
"crates/aether-wallet/src",
] {
for path in collect_workspace_rust_files(root) {
if path
@@ -132,6 +139,33 @@ fn runtime_state_owns_redis_runtime_boundaries() {
violations.join("\n")
);
let mut dependency_violations = Vec::new();
for manifest in [
"apps/aether-gateway/Cargo.toml",
"apps/aether-tunnel/Cargo.toml",
"crates/aether-admin/Cargo.toml",
"crates/aether-billing/Cargo.toml",
"crates/aether-model-fetch/Cargo.toml",
"crates/aether-provider-pool/Cargo.toml",
"crates/aether-provider-transport/Cargo.toml",
"crates/aether-runtime/Cargo.toml",
"crates/aether-task-runtime/Cargo.toml",
"crates/aether-usage-runtime/Cargo.toml",
"crates/aether-wallet/Cargo.toml",
] {
let cargo = read_workspace_file(manifest);
for forbidden in ["redis.workspace", "redis ="] {
if cargo.contains(forbidden) {
dependency_violations.push(format!("{manifest} -> {forbidden}"));
}
}
}
assert!(
dependency_violations.is_empty(),
"business/runtime crates must not depend on redis directly:\n{}",
dependency_violations.join("\n")
);
let mut runtime_state_violations = Vec::new();
for path in collect_workspace_rust_files("crates/aether-runtime-state/src") {
if path
@@ -160,6 +194,22 @@ fn runtime_state_owns_redis_runtime_boundaries() {
"only crates/aether-runtime-state/src/redis may depend on the redis crate directly:\n{}",
runtime_state_violations.join("\n")
);
let mut runtime_connection_violations = Vec::new();
for path in collect_workspace_rust_files("crates/aether-runtime-state/src") {
if path.ends_with("crates/aether-runtime-state/src/redis/client.rs") {
continue;
}
let source = production_workspace_source(&path);
if source.contains("get_multiplexed_async_connection") {
runtime_connection_violations.push(path.display().to_string());
}
}
assert!(
runtime_connection_violations.is_empty(),
"runtime Redis connections must be initialized only by redis/client.rs:\n{}",
runtime_connection_violations.join("\n")
);
}
#[test]
@@ -347,6 +347,78 @@ async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpo
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_counts_inherited_windsurf_key_formats_for_admin_provider_endpoints() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-windsurf/endpoints",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut key = sample_key(
"key-windsurf-a",
"provider-windsurf",
"openai:chat",
"oauth-secret",
);
key.auth_type = "oauth".to_string();
key.api_formats = None;
let mut provider = sample_provider("provider-windsurf", "windsurf", 10);
provider.provider_type = "windsurf".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-windsurf-chat",
"provider-windsurf",
"openai:chat",
"https://server.codeium.com",
)],
vec![key],
));
let (_, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-windsurf/endpoints?skip=0&limit=50"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let items = payload.as_array().expect("payload should be an array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["id"], "endpoint-windsurf-chat");
assert_eq!(items[0]["total_keys"], 1);
assert_eq!(items[0]["active_keys"], 1);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_service_unavailable_for_admin_provider_endpoint_create_when_catalog_writer_unavailable(
) {
@@ -754,6 +754,14 @@ async fn gateway_handles_admin_modules_status_locally_with_trusted_admin_princip
.with_system_config_values_for_tests(vec![
("module.oauth.enabled".to_string(), json!(true)),
("module.management_tokens.enabled".to_string(), json!(true)),
(
"module.important_notification.email_enabled".to_string(),
json!(true),
),
(
"module.important_notification.email_recipients".to_string(),
json!("ops@example.com"),
),
("smtp_host".to_string(), json!("smtp.example.com")),
("smtp_from_email".to_string(), json!("ops@example.com")),
]);
@@ -796,9 +804,20 @@ async fn gateway_handles_admin_modules_status_locally_with_trusted_admin_princip
"/admin/modules/chat-pii-redaction"
);
assert_eq!(
payload["notification_email"]["config_validated"],
payload["important_notification"]["config_validated"],
json!(true)
);
assert_eq!(
payload["important_notification"]["admin_route"],
"/admin/notification-service"
);
assert_eq!(payload["server_chan_push"]["display_name"], "Server 酱推送");
assert_eq!(
payload["server_chan_push"]["admin_route"],
"/admin/modules/server-chan"
);
assert_eq!(payload["bark_push"]["display_name"], "Bark 推送");
assert_eq!(payload["bark_push"]["admin_route"], "/admin/modules/bark");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -828,6 +847,14 @@ async fn gateway_handles_admin_modules_status_locally_with_bearer_admin_session(
.with_system_config_values_for_tests(vec![
("module.oauth.enabled".to_string(), json!(true)),
("module.management_tokens.enabled".to_string(), json!(true)),
(
"module.important_notification.email_enabled".to_string(),
json!(true),
),
(
"module.important_notification.email_recipients".to_string(),
json!("ops@example.com"),
),
("smtp_host".to_string(), json!("smtp.example.com")),
("smtp_from_email".to_string(), json!("ops@example.com")),
]);
@@ -855,7 +882,7 @@ async fn gateway_handles_admin_modules_status_locally_with_bearer_admin_session(
assert_eq!(payload["oauth"]["config_validated"], json!(true));
assert_eq!(payload["management_tokens"]["active"], json!(true));
assert_eq!(
payload["notification_email"]["config_validated"],
payload["important_notification"]["config_validated"],
json!(true)
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -216,6 +216,23 @@ fn codex_quota_execution_result(request_id: &str) -> serde_json::Value {
})
}
fn windsurf_register_user_execution_result(request_id: &str) -> serde_json::Value {
json!({
"request_id": request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"apiKey": "devin-session-token$registered",
"name": "Windsurf User",
"apiServerUrl": "https://server.codeium.com"
}
}
})
}
fn assert_single_provider_oauth_refresh_token_plan<'a>(
plans: &'a [ExecutionPlan],
) -> &'a ExecutionPlan {
@@ -281,18 +298,404 @@ async fn gateway_handles_admin_provider_oauth_supported_types_locally_with_trust
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let items = payload.as_array().expect("items should be array");
assert_eq!(items.len(), 5);
assert_eq!(items.len(), 6);
assert_eq!(items[0]["provider_type"], "claude_code");
assert_eq!(items[1]["provider_type"], "codex");
assert_eq!(items[2]["provider_type"], "chatgpt_web");
assert_eq!(items[3]["provider_type"], "gemini_cli");
assert_eq!(items[4]["provider_type"], "antigravity");
assert_eq!(items[5]["provider_type"], "windsurf");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_oauth_device_authorize_for_windsurf_browser() {
let mut provider = sample_provider("provider-windsurf", "windsurf", 10);
provider.provider_type = "windsurf".to_string();
let endpoint = sample_endpoint(
"endpoint-windsurf-chat",
"provider-windsurf",
"openai:chat",
"https://server.codeium.com",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
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-windsurf/device-authorize",
Some(json!({
"auth_type": "browser",
"login_option": "github",
"proxy_node_id": "proxy-node-windsurf"
})),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
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");
let session_id = payload["session_id"]
.as_str()
.expect("session_id should exist");
assert_eq!(payload["auth_type"], "browser");
assert_eq!(payload["login_option"], "github");
assert_eq!(payload["redirect_uri"], "show-auth-token");
assert_eq!(payload["callback_required"], true);
let authorization_url = payload["verification_uri_complete"]
.as_str()
.expect("authorization url should exist");
let parsed = url::Url::parse(authorization_url).expect("authorization url should parse");
let params = parsed
.query_pairs()
.into_owned()
.collect::<std::collections::BTreeMap<_, _>>();
assert_eq!(
parsed.as_str().split('?').next(),
Some("https://windsurf.com/windsurf/signin")
);
assert_eq!(
params.get("response_type").map(String::as_str),
Some("token")
);
assert_eq!(params.get("state").map(String::as_str), Some(session_id));
assert_eq!(
params.get("redirect_uri").map(String::as_str),
Some("show-auth-token")
);
assert_eq!(
params.get("redirect_parameters_type").map(String::as_str),
Some("query")
);
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-windsurf");
assert_eq!(stored["auth_type"], "browser");
assert_eq!(stored["social_provider"], "github");
assert_eq!(stored["redirect_uri"], "show-auth-token");
assert_eq!(stored["proxy_node_id"], "proxy-node-windsurf");
assert_eq!(stored["status"], "pending");
}
#[tokio::test]
async fn gateway_rejects_generic_oauth_start_for_windsurf_provider() {
let mut provider = sample_provider("provider-windsurf", "windsurf", 10);
provider.provider_type = "windsurf".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-windsurf/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("浏览器登录")),
"payload={payload}"
);
}
#[tokio::test]
async fn gateway_handles_admin_provider_oauth_device_poll_for_windsurf_callback_token() {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
if plan.request_id == "provider-oauth:windsurf-register:new" {
return Json(windsurf_register_user_execution_result(&plan.request_id));
}
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
}
}))
}
}),
);
let mut provider = sample_provider("provider-windsurf", "windsurf", 10);
provider.provider_type = "windsurf".to_string();
let endpoint = sample_endpoint(
"endpoint-windsurf-chat",
"provider-windsurf",
"openai:chat",
"https://server.codeium.com",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let mut proxy_node = sample_proxy_node("proxy-node-windsurf");
proxy_node.status = "online".to_string();
proxy_node.is_manual = true;
proxy_node.tunnel_mode = false;
proxy_node.tunnel_connected = false;
proxy_node.proxy_url = Some("http://proxy.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![proxy_node]));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY)
.attach_proxy_node_repository_for_tests(proxy_node_repository),
)
.with_provider_oauth_device_session_entry_for_tests(
"session-windsurf",
json!({
"provider_id": "provider-windsurf",
"region": "",
"client_id": "",
"client_secret": "",
"device_code": "",
"auth_type": "browser",
"social_provider": "google",
"code_verifier": null,
"redirect_uri": "show-auth-token",
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
"interval": 5,
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": "proxy-node-windsurf",
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
"error_msg": null,
}),
);
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-windsurf/device-poll",
Some(json!({
"session_id": "session-windsurf",
"callback_url": "https://windsurf.com/show-auth-token?token=firebase-id-token&state=session-windsurf&provider=google"
})),
)
.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}");
assert_eq!(payload["status"], "authorized");
assert_eq!(payload["replaced"], false);
let stored = state
.load_provider_oauth_device_session_for_tests("device_auth_session:session-windsurf")
.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");
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-windsurf", "enabled": true}))
);
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, "devin-session-token$registered");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["provider_type"], "windsurf");
assert_eq!(auth_config["auth_method"], "browser");
assert_eq!(auth_config["register_source"], "new");
assert_eq!(auth_config["social_provider"], "google");
{
let plans = execution_plans.lock().expect("mutex should lock");
let register_plan = plans
.iter()
.find(|plan| plan.request_id == "provider-oauth:windsurf-register:new")
.expect("register plan should execute");
assert_eq!(register_plan.method, "POST");
assert_eq!(
register_plan
.body
.json_body
.as_ref()
.and_then(|body| body.get("firebase_id_token"))
.and_then(serde_json::Value::as_str),
Some("firebase-id-token")
);
assert_eq!(
register_plan
.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-windsurf")
);
}
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_windsurf_callback_state_mismatch_and_missing_token() {
let mut provider = sample_provider("provider-windsurf", "windsurf", 10);
provider.provider_type = "windsurf".to_string();
let endpoint = sample_endpoint(
"endpoint-windsurf-chat",
"provider-windsurf",
"openai:chat",
"https://server.codeium.com",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
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_device_session_entry_for_tests(
"session-windsurf",
json!({
"provider_id": "provider-windsurf",
"region": "",
"client_id": "",
"client_secret": "",
"device_code": "",
"auth_type": "browser",
"social_provider": "google",
"code_verifier": null,
"redirect_uri": "show-auth-token",
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
"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,
}),
);
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-windsurf/device-poll",
Some(json!({
"session_id": "session-windsurf",
"callback_url": "https://windsurf.com/show-auth-token?token=firebase-id-token&state=wrong-state"
})),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
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!(payload["status"], "error");
assert!(payload["error"]
.as_str()
.is_some_and(|error| error.contains("state")));
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-windsurf/device-poll",
Some(json!({
"session_id": "session-windsurf",
"callback_url": "https://windsurf.com/show-auth-token?state=session-windsurf"
})),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
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!(payload["status"], "error");
assert!(payload["error"]
.as_str()
.is_some_and(|error| error.contains("token")));
}
#[tokio::test]
async fn gateway_handles_admin_provider_oauth_device_authorize_locally_with_trusted_admin_principal(
) {
@@ -2477,6 +2880,7 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p
.expect("account_state_recheck_error should be string when recheck is attempted");
assert!(
account_state_recheck_error == "wham/usage API 返回状态码 401"
|| account_state_recheck_error == "wham/usage API 返回状态码 403"
|| account_state_recheck_error.starts_with("wham/usage 请求执行失败:"),
"unexpected account_state_recheck_error: {account_state_recheck_error}"
);
@@ -4948,6 +5352,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
.expect("account_state_recheck_error should be string when attempted");
assert!(
account_state_recheck_error == "wham/usage API 返回状态码 401"
|| account_state_recheck_error == "wham/usage API 返回状态码 403"
|| account_state_recheck_error.starts_with("wham/usage 请求执行失败:"),
"unexpected account_state_recheck_error: {account_state_recheck_error}"
);
@@ -4997,6 +5402,14 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
stored_key.oauth_invalid_reason.as_deref(),
Some("[OAUTH_EXPIRED] Codex Token 无效或已过期 (401)")
);
} else if account_state_recheck_attempted
&& payload["account_state_recheck_error"] == "wham/usage API 返回状态码 403"
{
assert!(stored_key.oauth_invalid_at_unix_secs.is_some());
assert!(stored_key
.oauth_invalid_reason
.as_deref()
.is_some_and(|reason| reason.contains("(403)")));
} else {
assert_eq!(stored_key.oauth_invalid_at_unix_secs, None);
assert_eq!(stored_key.oauth_invalid_reason, None);
@@ -244,6 +244,144 @@ async fn gateway_handles_admin_provider_query_models_fetches_upstream_for_select
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_models_fetches_windsurf_model_configs() {
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_runtime_hits_inner = Arc::clone(&execution_runtime_hits_clone);
async move {
*execution_runtime_hits_inner
.lock()
.expect("mutex should lock") += 1;
assert_eq!(plan.method, "POST");
assert_eq!(
plan.url,
"https://server.codeium.com/exa.api_server_pb.ApiServerService/GetCascadeModelConfigs"
);
assert_eq!(plan.client_api_format, "openai:chat");
assert_eq!(plan.provider_api_format, "windsurf:model_configs");
assert_eq!(plan.model_name.as_deref(), Some("GetCascadeModelConfigs"));
assert_eq!(
plan.headers.get("connect-protocol-version").map(String::as_str),
Some("1")
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("metadata"))
.and_then(|metadata| metadata.get("apiKey")),
Some(&json!("devin-session-token$abc"))
);
Json(json!({
"request_id": "req-provider-query-windsurf",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"clientModelConfigs": [{
"modelUid": "claude-sonnet-4-6",
"label": "Claude Sonnet 4.6",
"provider": "anthropic",
"supportsImages": true,
"creditMultiplier": 4
}],
"defaultOverrideModelConfig": {
"modelUid": "claude-sonnet-4-6"
}
}
}
}))
}
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-windsurf", "Windsurf", 10);
provider.provider_type = "windsurf".to_string();
let mut windsurf_key = sample_key(
"key-windsurf-selected",
"provider-windsurf",
"openai:chat",
"devin-session-token$abc",
);
windsurf_key.auth_type = "oauth".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![StoredProviderCatalogEndpoint::new(
"endpoint-windsurf-chat".to_string(),
"provider-windsurf".to_string(),
"openai:chat".to_string(),
Some("chat".to_string()),
Some("primary".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://server.codeium.com".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")],
vec![windsurf_key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/models"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-windsurf",
"api_key_id": "key-windsurf-selected"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["data"]["error"], serde_json::Value::Null);
assert_eq!(payload["data"]["from_cache"], json!(false));
let models = payload["data"]["models"]
.as_array()
.expect("models should be an array");
assert_eq!(models.len(), 1);
assert_eq!(models[0]["id"], json!("claude-sonnet-4-6"));
assert_eq!(
models[0]["api_formats"],
json!(["openai:chat", "openai:responses", "claude:messages"])
);
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
1
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_models_with_openai_responses_endpoint() {
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
@@ -2247,6 +2385,122 @@ async fn gateway_routes_grok_responses_admin_pool_model_test_through_grok_runtim
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_streams_windsurf_connect_upstream_for_admin_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-windsurf");
assert_eq!(plan.endpoint_id, "endpoint-windsurf-chat");
assert_eq!(plan.key_id, "key-windsurf-primary");
assert_eq!(plan.provider_api_format, "openai:chat");
assert_eq!(plan.content_type.as_deref(), Some("application/connect+json"));
assert!(plan.stream, "Windsurf Connect model test must stream upstream");
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("stream")),
Some(&json!(true))
);
let windsurf_payload = serde_json::to_vec(&json!({
"chatMessage": {
"text": "ok"
}
}))
.expect("windsurf payload should encode");
let mut windsurf_frame = vec![0u8];
windsurf_frame.extend_from_slice(&(windsurf_payload.len() as u32).to_be_bytes());
windsurf_frame.extend_from_slice(&windsurf_payload);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/connect+json"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(windsurf_frame)
},
"telemetry": {
"elapsed_ms": 24
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-windsurf", "Windsurf", 10);
provider.provider_type = "windsurf".to_string();
let mut key = sample_key(
"key-windsurf-primary",
"provider-windsurf",
"openai:chat",
"devin-session-token$abc",
);
key.auth_type = "oauth".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-windsurf-chat",
"provider-windsurf",
"openai:chat",
"https://server.codeium.com",
)],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-windsurf",
"model": "claude-opus-4-7-medium",
"api_format": "openai:chat",
"endpoint_id": "endpoint-windsurf-chat",
"request_body": {
"model": "claude-opus-4-7-medium",
"messages": [{
"role": "user",
"content": "Hello! This is a test message."
}],
"max_tokens": 30,
"temperature": 0.7,
"stream": true
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["attempts"][0]["request_body"]["stream"],
json!(true)
);
assert_eq!(
payload["attempts"][0]["response_body"]["choices"][0]["message"]["content"],
json!("ok")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_uses_pool_scheduler_order_for_admin_pool_model_test() {
let execution_runtime = Router::new().route(
@@ -478,6 +478,73 @@ async fn gateway_imports_admin_system_config_locally_and_persists_data() {
let _ = upstream_url;
}
#[tokio::test]
async fn gateway_imports_admin_system_config_openai_image_aliases() {
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
Vec::new(),
Vec::new(),
Vec::new(),
));
let global_model_repository = Arc::new(InMemoryGlobalModelReadRepository::seed(Vec::<
StoredPublicGlobalModel,
>::new()));
let data_state = build_admin_system_data_state_with_repositories(
Arc::clone(&provider_catalog_repository),
Arc::clone(&global_model_repository),
);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let mut payload = sample_system_import_payload();
payload["providers"][0]["endpoints"][0]["api_format"] = json!("openai_image");
payload["providers"][0]["api_keys"][0]["api_formats"] = json!(["images"]);
payload["providers"][0]["api_keys"][0]["supported_endpoints"] = json!(["openai:image"]);
payload["providers"][0]["models"][0]["supports_image_generation"] = json!(true);
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/config/import"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&payload)
.send()
.await
.expect("request should succeed");
let status = response.status();
let body: Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={body}");
assert_eq!(body["stats"]["endpoints"]["created"], json!(1));
assert_eq!(body["stats"]["keys"]["created"], json!(1));
let providers = provider_catalog_repository
.list_providers(false)
.await
.expect("providers should load");
let provider_ids = providers
.iter()
.map(|provider| provider.id.clone())
.collect::<Vec<_>>();
let endpoints = provider_catalog_repository
.list_endpoints_by_provider_ids(&provider_ids)
.await
.expect("endpoints should load");
assert_eq!(endpoints[0].api_format, "openai:image");
let keys = provider_catalog_repository
.list_keys_by_provider_ids(&provider_ids)
.await
.expect("keys should load");
assert_eq!(keys[0].api_formats, Some(json!(["openai:image"])));
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_returns_503_for_admin_system_config_import_when_local_data_is_unavailable() {
let upstream_hits = Arc::new(Mutex::new(0usize));
+12 -4
View File
@@ -625,10 +625,18 @@ fn assert_usage_and_pricing(
expected.cache_read_tokens
);
assert_eq!(stored_usage.total_tokens, expected.total_tokens());
assert_eq!(
stored_usage.response_time_ms,
Some(expected_response_time_ms)
);
if expected_ttfb_ms.is_some() {
assert!(
stored_usage.response_time_ms >= Some(expected_response_time_ms),
"stream response_time_ms should be at least reported telemetry: expected >= {expected_response_time_ms:?}, got {:?}",
stored_usage.response_time_ms
);
} else {
assert_eq!(
stored_usage.response_time_ms,
Some(expected_response_time_ms)
);
}
assert_eq!(stored_usage.first_byte_time_ms, expected_ttfb_ms);
assert_eq!(
stored_usage.settlement_input_price_per_1m(),
+23 -4
View File
@@ -18,7 +18,7 @@ use context::{report_context_is_locally_actionable, resolve_locally_actionable_r
use aether_usage_runtime::{
is_local_ai_stream_report_kind, is_local_ai_sync_report_kind, report_request_id,
should_handle_local_stream_report, should_handle_local_sync_report,
sync_report_represents_failure,
stream_report_represents_failure, sync_report_represents_failure,
};
pub(crate) use aether_usage_runtime::{GatewayStreamReportRequest, GatewaySyncReportRequest};
@@ -256,14 +256,33 @@ async fn handle_local_stream_report(state: &AppState, payload: &GatewayStreamRep
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms);
let failed = stream_report_represents_failure(payload);
record_report_request_candidate_status(
state,
payload.report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
status: if failed {
RequestCandidateStatus::Failed
} else {
RequestCandidateStatus::Success
},
status_code: Some(payload.status_code),
error_type: None,
error_message: None,
error_type: failed.then(|| {
if payload.status_code >= 400 {
"stream_http_error".to_string()
} else {
"stream_terminal_error".to_string()
}
}),
error_message: failed.then(|| {
payload
.terminal_summary
.as_ref()
.and_then(|summary| summary.parser_error.clone())
.unwrap_or_else(|| {
"execution runtime stream ended with a terminal error".to_string()
})
}),
latency_ms,
started_at_unix_ms: None,
finished_at_unix_ms: Some(terminal_unix_ms),
@@ -9,6 +9,7 @@ use axum::{
response::{IntoResponse, Response},
Json,
};
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
use serde_json::{json, Value};
use std::collections::BTreeMap;
@@ -548,16 +549,69 @@ fn admin_monitoring_trace_response_data(
return None;
}
let body = admin_monitoring_trace_response_body(headers, body);
Some(json!({
"source": source,
"status_code": status_code,
"headers": headers.cloned().unwrap_or(Value::Null),
"body": body.cloned().unwrap_or(Value::Null),
"body": body.unwrap_or(Value::Null),
"body_ref": body_ref,
"body_state": body_state.map(|state| state.as_str()),
}))
}
fn admin_monitoring_trace_response_body(
headers: Option<&Value>,
body: Option<&Value>,
) -> Option<Value> {
let body = body?;
admin_monitoring_decode_connect_json_error_body(headers, body).or_else(|| Some(body.clone()))
}
fn admin_monitoring_decode_connect_json_error_body(
headers: Option<&Value>,
body: &Value,
) -> Option<Value> {
if !admin_monitoring_headers_indicate_connect_json(headers) {
return None;
}
let body_base64 = match body {
Value::String(value) => Some(value.as_str()),
Value::Object(object) => object
.get("encoding")
.and_then(Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("base64"))
.then(|| object.get("data").and_then(Value::as_str))
.flatten(),
_ => None,
}?
.trim();
if body_base64.is_empty() {
return None;
}
let body_bytes = BASE64_STANDARD.decode(body_base64).ok()?;
aether_ai_formats::api::extract_provider_private_stream_error_body(None, &body_bytes)
}
fn admin_monitoring_headers_indicate_connect_json(headers: Option<&Value>) -> bool {
headers
.and_then(Value::as_object)
.and_then(|object| {
object.iter().find_map(|(key, value)| {
key.eq_ignore_ascii_case("content-type")
.then(|| value.as_str())
.flatten()
})
})
.map(str::trim)
.is_some_and(|value| {
let value = value.to_ascii_lowercase();
value.contains("application/connect+json") || value.contains("+connect+json")
})
}
fn merge_admin_monitoring_trace_response(
extra_object: &mut serde_json::Map<String, Value>,
key: &str,
@@ -1,5 +1,6 @@
use crate::observability::stats::{aggregate_usage_stats, parse_bounded_u32, round_to};
use aether_ai_formats::api::request_path_implies_stream_request;
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_billing::{
normalize_input_tokens_for_billing, normalize_total_input_context_for_cache_hit_rate,
};
@@ -1052,7 +1053,7 @@ fn admin_usage_upstream_is_stream(item: &StoredRequestUsageAudit) -> bool {
item.request_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool)
.or_else(|| admin_usage_headers_stream_flag(item.response_headers.as_ref()))
.or_else(|| admin_usage_infer_upstream_stream_from_captured_bodies(item))
@@ -1256,7 +1257,10 @@ pub fn admin_usage_record_json(
.as_object_mut()
.expect("admin usage record payload should be an object");
object.insert("is_stream".to_string(), json!(item.is_stream));
object.insert("upstream_is_stream".to_string(), json!(upstream_is_stream));
object.insert(
UPSTREAM_IS_STREAM_KEY.to_string(),
json!(upstream_is_stream),
);
object.insert(
"client_requested_stream".to_string(),
json!(client_is_stream),
+38 -14
View File
@@ -797,22 +797,13 @@ pub fn admin_provider_ops_sub2api_verify_payload(
}
}
let username_or_email = admin_provider_ops_sub2api_non_empty_string(user_data, "username")
.or_else(|| admin_provider_ops_sub2api_non_empty_string(user_data, "email"));
admin_provider_ops_verify_success(
admin_provider_ops_verify_user_payload(
user_data
.get("username")
.or_else(|| user_data.get("email"))
.and_then(Value::as_str)
.map(ToOwned::to_owned),
user_data
.get("username")
.or_else(|| user_data.get("email"))
.and_then(Value::as_str)
.map(ToOwned::to_owned),
user_data
.get("email")
.and_then(Value::as_str)
.map(ToOwned::to_owned),
username_or_email.clone(),
username_or_email,
admin_provider_ops_sub2api_non_empty_string(user_data, "email"),
Some(balance + points),
Some(extra),
),
@@ -820,6 +811,17 @@ pub fn admin_provider_ops_sub2api_verify_payload(
)
}
fn admin_provider_ops_sub2api_non_empty_string(
map: &Map<String, Value>,
key: &str,
) -> Option<String> {
map.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
#[cfg(test)]
mod tests {
use super::{
@@ -913,6 +915,28 @@ mod tests {
);
}
#[test]
fn sub2api_verify_payload_falls_back_to_email_when_username_is_null() {
let payload = admin_provider_ops_sub2api_verify_payload(
StatusCode::OK,
&json!({
"code": 0,
"data": {
"username": null,
"email": "user@example.com",
"balance": 2.0,
"points": 0.0
}
}),
None,
);
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["data"]["username"], json!("user@example.com"));
assert_eq!(payload["data"]["display_name"], json!("user@example.com"));
assert_eq!(payload["data"]["email"], json!("user@example.com"));
}
#[test]
fn anyrouter_verify_payload_uses_cookie_auth_messages_and_usage_fields() {
let payload = admin_provider_ops_anyrouter_verify_payload(
+427 -2
View File
@@ -1191,6 +1191,317 @@ pub fn parse_kiro_usage_response(
Some(serde_json::Value::Object(result))
}
pub fn parse_windsurf_user_status_response(
value: &serde_json::Value,
updated_at_unix_secs: u64,
) -> Option<serde_json::Value> {
let user_status = value
.get("userStatus")
.or_else(|| value.get("user_status"))?;
let plan_status = user_status
.get("planStatus")
.or_else(|| user_status.get("plan_status"))?;
let plan_info = plan_status
.get("planInfo")
.or_else(|| plan_status.get("plan_info"));
let mut result = serde_json::Map::new();
result.insert("updated_at".to_string(), json!(updated_at_unix_secs));
if let Some(plan_name) = plan_info
.and_then(|value| {
coerce_json_string(value.get("planName").or_else(|| value.get("plan_name")))
})
.or_else(|| {
coerce_json_string(
plan_status
.get("planName")
.or_else(|| plan_status.get("plan_name")),
)
})
{
result.insert("plan_name".to_string(), json!(plan_name));
}
if let Some(email) = coerce_json_string(user_status.get("email")) {
result.insert("email".to_string(), json!(email));
}
if let Some(value) = plan_status
.get("dailyQuotaRemainingPercent")
.or_else(|| plan_status.get("daily_quota_remaining_percent"))
.and_then(coerce_json_f64)
{
result.insert("daily_remaining_percent".to_string(), json!(value));
}
if let Some(value) = plan_status
.get("weeklyQuotaRemainingPercent")
.or_else(|| plan_status.get("weekly_quota_remaining_percent"))
.and_then(coerce_json_f64)
{
result.insert("weekly_remaining_percent".to_string(), json!(value));
}
if let Some(value) = plan_status
.get("dailyQuotaResetAtUnix")
.or_else(|| plan_status.get("daily_quota_reset_at_unix"))
.and_then(coerce_json_u64)
{
result.insert("daily_reset_at".to_string(), json!(value));
}
if let Some(value) = plan_status
.get("weeklyQuotaResetAtUnix")
.or_else(|| plan_status.get("weekly_quota_reset_at_unix"))
.and_then(coerce_json_u64)
{
result.insert("weekly_reset_at".to_string(), json!(value));
}
if let Some(value) = plan_status
.get("overageBalanceMicros")
.or_else(|| plan_status.get("overage_balance_micros"))
.and_then(coerce_json_f64)
{
result.insert("overage_balance".to_string(), json!(value / 1_000_000.0));
}
let legacy_credit =
|value: Option<&serde_json::Value>| value.and_then(coerce_json_f64).map(|n| n / 100.0);
if let Some(value) = legacy_credit(
plan_status
.get("availablePromptCredits")
.or_else(|| plan_status.get("available_prompt_credits")),
) {
result.insert("prompt_remaining".to_string(), json!(value));
}
if let Some(value) = legacy_credit(
plan_status
.get("usedPromptCredits")
.or_else(|| plan_status.get("used_prompt_credits")),
) {
result.insert("prompt_used".to_string(), json!(value));
}
if let Some(value) = legacy_credit(plan_info.and_then(|plan_info| {
plan_info
.get("monthlyPromptCredits")
.or_else(|| plan_info.get("monthly_prompt_credits"))
})) {
result.insert("prompt_limit".to_string(), json!(value));
}
if let Some(value) = legacy_credit(
plan_status
.get("availableFlexCredits")
.or_else(|| plan_status.get("available_flex_credits")),
) {
result.insert("flex_remaining".to_string(), json!(value));
}
if let Some(value) = legacy_credit(
plan_status
.get("usedFlexCredits")
.or_else(|| plan_status.get("used_flex_credits")),
) {
result.insert("flex_used".to_string(), json!(value));
}
if let Some(value) = legacy_credit(plan_info.and_then(|plan_info| {
plan_info
.get("monthlyFlexCreditPurchaseAmount")
.or_else(|| plan_info.get("monthly_flex_credit_purchase_amount"))
})) {
result.insert("flex_limit".to_string(), json!(value));
}
let mut status_sources = vec![value, user_status, plan_status];
if let Some(plan_info) = plan_info {
status_sources.push(plan_info);
}
for (target, aliases) in [
(
"banned",
&[
"banned",
"isBanned",
"is_banned",
"accountBanned",
"account_banned",
][..],
),
(
"quarantined",
&[
"quarantined",
"isQuarantined",
"is_quarantined",
"accountQuarantined",
"account_quarantined",
][..],
),
(
"is_forbidden",
&[
"isForbidden",
"is_forbidden",
"forbidden",
"accountForbidden",
"account_forbidden",
][..],
),
] {
if let Some(found) = status_sources.iter().find_map(|source| {
aliases
.iter()
.find_map(|alias| source.get(*alias).and_then(coerce_json_bool))
}) {
result.insert(target.to_string(), json!(found));
}
}
for (target, aliases) in [
(
"ban_reason",
&[
"banReason",
"ban_reason",
"blockedReason",
"blocked_reason",
"reason",
"message",
][..],
),
(
"quarantine_reason",
&["quarantineReason", "quarantine_reason", "reason", "message"][..],
),
(
"forbidden_reason",
&["forbiddenReason", "forbidden_reason", "reason", "message"][..],
),
] {
if let Some(found) = status_sources.iter().find_map(|source| {
aliases
.iter()
.find_map(|alias| coerce_json_string(source.get(*alias)))
}) {
result.insert(target.to_string(), json!(found));
}
}
Some(serde_json::Value::Object(result))
}
pub fn parse_windsurf_model_configs_response(
value: &serde_json::Value,
updated_at_unix_secs: u64,
) -> Option<serde_json::Value> {
let configs = value
.get("clientModelConfigs")
.or_else(|| value.get("client_model_configs"))
.and_then(serde_json::Value::as_array)?;
let mut models = Vec::new();
for config in configs {
let Some(model_uid) = coerce_json_string(
config
.get("modelUid")
.or_else(|| config.get("model_uid"))
.or_else(|| config.get("id"))
.or_else(|| config.get("name")),
) else {
continue;
};
let mut model = serde_json::Map::new();
model.insert("model_uid".to_string(), json!(model_uid));
if let Some(label) = coerce_json_string(
config
.get("label")
.or_else(|| config.get("displayName"))
.or_else(|| config.get("display_name")),
) {
model.insert("label".to_string(), json!(label));
}
if let Some(provider) = coerce_json_string(config.get("provider")) {
model.insert("provider".to_string(), json!(provider));
}
if let Some(value) = config
.get("supportsImages")
.or_else(|| config.get("supports_images"))
.and_then(coerce_json_bool)
{
model.insert("supports_images".to_string(), json!(value));
}
if let Some(value) = config
.get("creditMultiplier")
.or_else(|| config.get("credit_multiplier"))
.and_then(coerce_json_f64)
{
model.insert("credit_multiplier".to_string(), json!(value));
}
models.push(serde_json::Value::Object(model));
}
let mut result = serde_json::Map::new();
result.insert("updated_at".to_string(), json!(updated_at_unix_secs));
result.insert(
"allowed_models_count".to_string(),
json!(models.len() as u64),
);
result.insert("models".to_string(), serde_json::Value::Array(models));
if let Some(default_model_uid) = value
.get("defaultOverrideModelConfig")
.or_else(|| value.get("default_override_model_config"))
.and_then(|default_config| {
coerce_json_string(
default_config
.get("modelUid")
.or_else(|| default_config.get("model_uid")),
)
})
{
result.insert("default_model_uid".to_string(), json!(default_model_uid));
}
Some(serde_json::Value::Object(result))
}
pub fn parse_windsurf_rate_limit_response(
value: &serde_json::Value,
updated_at_unix_secs: u64,
) -> Option<serde_json::Value> {
let root = value.as_object()?;
if root.is_empty() {
return None;
}
let has_capacity = value
.get("hasCapacity")
.or_else(|| value.get("has_capacity"))
.and_then(coerce_json_bool)
.unwrap_or(true);
let messages_remaining = value
.get("messagesRemaining")
.or_else(|| value.get("messages_remaining"))
.and_then(coerce_json_f64);
let max_messages = value
.get("maxMessages")
.or_else(|| value.get("max_messages"))
.and_then(coerce_json_f64);
let retry_after_ms = value
.get("retryAfterMs")
.or_else(|| value.get("retry_after_ms"))
.and_then(coerce_json_u64);
let limited = !has_capacity || messages_remaining.is_some_and(|value| value <= 0.0);
let mut rate_limit = serde_json::Map::new();
rate_limit.insert("limited".to_string(), json!(limited));
rate_limit.insert("has_capacity".to_string(), json!(has_capacity));
if let Some(value) = messages_remaining {
rate_limit.insert("messages_remaining".to_string(), json!(value));
}
if let Some(value) = max_messages {
rate_limit.insert("max_messages".to_string(), json!(value));
}
if let Some(value) = retry_after_ms {
rate_limit.insert("retry_after_ms".to_string(), json!(value));
}
Some(json!({
"updated_at": updated_at_unix_secs,
"rate_limit": rate_limit,
}))
}
fn chatgpt_web_quota_feature_name(value: &serde_json::Value) -> Option<String> {
coerce_json_string(
value
@@ -1411,8 +1722,10 @@ mod tests {
codex_build_invalid_state, codex_runtime_invalid_reason,
parse_chatgpt_web_conversation_init_response, parse_codex_backend_me_response,
parse_codex_wham_usage_response, parse_gemini_cli_retrieve_user_quota_response,
parse_gemini_cli_v1internal_credits_response, OAUTH_ACCOUNT_BLOCK_PREFIX,
OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
parse_gemini_cli_v1internal_credits_response, parse_windsurf_model_configs_response,
parse_windsurf_rate_limit_response, parse_windsurf_user_status_response,
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
OAUTH_REQUEST_FAILED_PREFIX,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::json;
@@ -1853,6 +2166,118 @@ mod tests {
assert_eq!(parsed.get("updated_at"), Some(&json!(1_777_000_123u64)));
}
#[test]
fn parses_windsurf_user_status_response() {
let parsed = parse_windsurf_user_status_response(
&json!({
"userStatus": {
"email": "windsurf@example.com",
"isQuarantined": true,
"quarantineReason": "quota review",
"planStatus": {
"dailyQuotaRemainingPercent": 45.5,
"weeklyQuotaRemainingPercent": 80,
"dailyQuotaResetAtUnix": "1775553285",
"weeklyQuotaResetAtUnix": 1776158085u64,
"availablePromptCredits": 900,
"usedPromptCredits": 100,
"availableFlexCredits": 250,
"usedFlexCredits": 50,
"overageBalanceMicros": 1250000,
"planInfo": {
"planName": "Pro",
"monthlyPromptCredits": 1000,
"monthlyFlexCreditPurchaseAmount": 300
}
}
}
}),
1_770_000_000,
)
.expect("windsurf user status should parse");
assert_eq!(parsed.get("plan_name"), Some(&json!("Pro")));
assert_eq!(parsed.get("daily_remaining_percent"), Some(&json!(45.5)));
assert_eq!(parsed.get("weekly_remaining_percent"), Some(&json!(80.0)));
assert_eq!(parsed.get("daily_reset_at"), Some(&json!(1_775_553_285u64)));
assert_eq!(
parsed.get("weekly_reset_at"),
Some(&json!(1_776_158_085u64))
);
assert_eq!(parsed.get("prompt_remaining"), Some(&json!(9.0)));
assert_eq!(parsed.get("prompt_used"), Some(&json!(1.0)));
assert_eq!(parsed.get("prompt_limit"), Some(&json!(10.0)));
assert_eq!(parsed.get("flex_remaining"), Some(&json!(2.5)));
assert_eq!(parsed.get("flex_used"), Some(&json!(0.5)));
assert_eq!(parsed.get("flex_limit"), Some(&json!(3.0)));
assert_eq!(parsed.get("overage_balance"), Some(&json!(1.25)));
assert_eq!(parsed.get("email"), Some(&json!("windsurf@example.com")));
assert_eq!(parsed.get("quarantined"), Some(&json!(true)));
assert_eq!(
parsed.get("quarantine_reason"),
Some(&json!("quota review"))
);
assert_eq!(parsed.get("updated_at"), Some(&json!(1_770_000_000u64)));
}
#[test]
fn parses_windsurf_model_configs_response() {
let parsed = parse_windsurf_model_configs_response(
&json!({
"clientModelConfigs": [
{
"modelUid": "claude-sonnet-4-5",
"label": "Claude Sonnet 4.5",
"provider": "anthropic",
"supportsImages": true,
"creditMultiplier": 2
},
{
"modelUid": "gpt-5-mini",
"label": "GPT-5 mini"
}
],
"defaultOverrideModelConfig": {
"modelUid": "claude-sonnet-4-5"
}
}),
1_770_000_100,
)
.expect("windsurf model configs should parse");
assert_eq!(parsed.get("allowed_models_count"), Some(&json!(2u64)));
assert_eq!(
parsed.get("default_model_uid"),
Some(&json!("claude-sonnet-4-5"))
);
assert_eq!(parsed.get("updated_at"), Some(&json!(1_770_000_100u64)));
}
#[test]
fn parses_windsurf_rate_limit_response() {
let parsed = parse_windsurf_rate_limit_response(
&json!({
"hasCapacity": false,
"messagesRemaining": 0,
"maxMessages": 25,
"retryAfterMs": 45000
}),
1_770_000_200,
)
.expect("windsurf rate limit should parse");
assert_eq!(parsed.get("updated_at"), Some(&json!(1_770_000_200u64)));
assert_eq!(parsed.pointer("/rate_limit/limited"), Some(&json!(true)));
assert_eq!(
parsed.pointer("/rate_limit/messages_remaining"),
Some(&json!(0.0))
);
assert_eq!(
parsed.pointer("/rate_limit/retry_after_ms"),
Some(&json!(45000u64))
);
}
#[test]
fn parses_chatgpt_web_image_quota_from_conversation_init() {
let parsed = parse_chatgpt_web_conversation_init_response(

Some files were not shown because too many files have changed in this diff Show More