i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 09:20:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(provider): 修复 Windsurf 原生工具桥接

Browse Source
This commit is contained in:
Entropy.Xu
2026-05-20 23:42:12 +08:00
parent 82637ad882
commit 129c7c90c0
30 changed files with 9090 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
apps/aether-gateway/src/ai_serving/api.rs
View File

@@ -44,6 +44,7 @@ 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,
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,

1
apps/aether-gateway/src/execution_runtime/mod.rs
View File

@@ -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::{

464
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -44,13 +44,14 @@ use super::error::{
#[path = "execution_failures.rs"]
mod execution_failures;
use self::execution_failures::{
build_stream_failure_from_execution_error, build_stream_failure_report,
build_stream_failure_from_execution_error, build_stream_failure_from_provider_error_body,
build_stream_failure_report, handle_prefetch_provider_private_stream_error,
handle_prefetch_stream_failure, submit_midstream_stream_failure, StreamFailureReport,
};
use crate::ai_serving::api::{
maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer,
maybe_build_stream_response_rewriter, normalize_provider_private_report_context,
StreamingStandardTerminalObserver,
extract_provider_private_stream_error_body, maybe_bridge_standard_sync_json_to_stream,
maybe_build_provider_private_stream_normalizer, maybe_build_stream_response_rewriter,
normalize_provider_private_report_context, StreamingStandardTerminalObserver,
};
use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
@@ -73,14 +74,15 @@ use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry;
#[cfg(test)]
use crate::execution_runtime::remote_compat::post_stream_plan_to_remote_execution_runtime;
use crate::execution_runtime::submission::{
resolve_core_error_background_report_kind, strip_utf8_bom_and_ws,
submit_local_core_error_or_sync_finalize,
resolve_core_error_background_report_kind, resolve_local_sync_error_status_code,
strip_utf8_bom_and_ws, submit_local_core_error_or_sync_finalize,
};
use crate::execution_runtime::transport::{
execute_stream_plan_via_local_tunnel, record_manual_proxy_request_failure,
record_manual_proxy_request_success, record_manual_proxy_stream_error,
DirectSyncExecutionRuntime, DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
};
use crate::execution_runtime::windsurf::maybe_execute_windsurf_stream;
use crate::execution_runtime::{
apply_endpoint_response_header_rules, attach_provider_response_headers_to_report_context,
local_failover_response_text, resolve_core_stream_direct_finalize_report_kind,
@@ -849,6 +851,58 @@ pub(crate) async fn execute_execution_runtime_stream(
return Ok(None);
}
}
match maybe_execute_windsurf_stream(state, &plan, report_context.as_ref()).await {
Ok(Some(windsurf_stream)) => {
return execute_stream_from_frame_stream(
state,
plan,
trace_id,
decision,
plan_kind,
report_kind,
windsurf_stream.report_context.or(report_context),
candidate_started_unix_secs,
stream_started_at,
windsurf_stream.frame_stream,
provider_pool_in_flight_guard.take(),
)
.await;
}
Ok(None) => {}
Err(err) => {
info!(
event_name = "windsurf_native_execution_unavailable",
log_type = "ops",
trace_id = %trace_id,
request_id = %plan_request_id_for_log,
candidate_id = ?plan.candidate_id,
provider_name = provider_name.as_str(),
endpoint_id = %endpoint_id,
key_id = %key_id,
model_name = model_name.as_str(),
candidate_index = candidate_index.as_str(),
error = %err,
"gateway native Windsurf stream execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("windsurf_native_execution_unavailable".to_string()),
error_message: Some(err.to_string()),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
}
match maybe_execute_kiro_web_search_stream(state, &plan, report_context.as_ref()).await {
Ok(Some(kiro_web_search)) => {
return execute_stream_from_frame_stream(
@@ -1197,13 +1251,13 @@ fn encode_terminal_sse_error_event(failure: &StreamFailureReport) -> Result<Byte
let payload = failure
.to_json_string()
.map_err(|err| IoError::other(err.to_string()))?;
let mut event = String::from("event: aether.error\n");
let mut event = String::new();
for line in payload.lines() {
event.push_str("data: ");
event.push_str(line);
event.push('\n');
}
event.push('\n');
event.push_str("\ndata: [DONE]\n\n");
Ok(Bytes::from(event))
}
@@ -1693,24 +1747,46 @@ async fn execute_stream_from_frame_stream(
if !(200..300).contains(&status_code) {
let provider_error_body = collect_error_body(&mut lines).await?;
let synthetic_body_json =
should_synthesize_non_success_stream_error_body(status_code, &provider_error_body)
.then(|| build_synthetic_non_success_stream_error_body(status_code, &headers));
let (provider_body_json, provider_body_base64) =
decode_stream_error_body(&headers, &provider_error_body);
let client_status_code = stream_client_error_status_code_for_upstream_status(status_code);
let wrapped_binary_body_json = wrap_non_json_binary_stream_error_for_client(
plan_kind,
&headers,
let private_error_body_json = extract_provider_private_stream_error_body(
report_context.as_ref(),
&provider_error_body,
)?;
let (client_body_json, client_error_body) =
);
let provider_private_error_decoded = private_error_body_json.is_some();
let synthetic_body_json = (!provider_private_error_decoded
&& should_synthesize_non_success_stream_error_body(status_code, &provider_error_body))
.then(|| build_synthetic_non_success_stream_error_body(status_code, &headers));
let (provider_body_json, provider_body_base64) =
if let Some(error_body_json) = private_error_body_json {
(Some(error_body_json), None)
} else {
decode_stream_error_body(&headers, &provider_error_body)
};
let client_status_code = stream_client_error_status_code_for_upstream_status(status_code);
let wrapped_binary_body_json = if provider_private_error_decoded {
None
} else {
wrap_non_json_binary_stream_error_for_client(plan_kind, &headers, &provider_error_body)?
};
let (client_body_json, client_error_body, payload_client_body_json) =
if let Some(body_json) = synthetic_body_json.or(wrapped_binary_body_json) {
let body_bytes = serde_json::to_vec(&body_json)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
(Some(body_json), body_bytes)
(Some(body_json.clone()), body_bytes, Some(body_json))
} else if provider_private_error_decoded {
let body_json = provider_body_json.clone().ok_or_else(|| {
GatewayError::Internal(
"decoded provider private stream error body is missing".to_string(),
)
})?;
let body_bytes = serde_json::to_vec(&body_json)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
(Some(body_json), body_bytes, None)
} else {
(provider_body_json.clone(), provider_error_body.clone())
(
provider_body_json.clone(),
provider_error_body.clone(),
provider_body_json.clone(),
)
};
let error_response_text =
local_failover_response_text(client_body_json.as_ref(), &client_error_body, None);
@@ -1897,6 +1973,11 @@ async fn execute_stream_from_frame_stream(
} else {
headers.clone()
};
if provider_private_error_decoded {
client_headers.remove("content-encoding");
client_headers.remove("content-length");
client_headers.insert("content-type".to_string(), "application/json".to_string());
}
apply_endpoint_response_header_rules(
state,
&plan,
@@ -1928,7 +2009,7 @@ async fn execute_stream_from_frame_stream(
provider_body_json,
provider_body_base64,
client_headers,
client_body_json,
payload_client_body_json,
None,
);
record_sync_terminal_usage(state, &plan, payload.report_context.as_ref(), &payload);
@@ -2127,6 +2208,30 @@ async fn execute_stream_from_frame_stream(
provider_prefetched_body.extend_from_slice(&chunk);
prefetched_inspection_body.extend_from_slice(&chunk);
if let Some(error_body_json) = extract_provider_private_stream_error_body(
report_context.as_ref(),
&prefetched_inspection_body,
) {
let error_status_code =
resolve_local_sync_error_status_code(status_code, &error_body_json);
return handle_prefetch_provider_private_stream_error(
state,
trace_id,
decision,
&plan,
report_context,
request_id,
candidate_id,
report_kind,
headers,
prefetched_telemetry,
&provider_prefetched_body,
error_status_code,
error_body_json,
)
.await;
}
let inspection = inspect_prefetched_stream_body(
&upstream_headers,
&prefetched_inspection_body,
@@ -2835,6 +2940,11 @@ async fn execute_stream_from_frame_stream(
} else {
chunk
};
let provider_private_error_body_json =
extract_provider_private_stream_error_body(
stream_usage_report_context.as_ref(),
&normalized_chunk,
);
if let (Some(observer), Some(report_context)) = (
stream_usage_observer.as_mut(),
stream_usage_report_context.as_ref(),
@@ -2873,6 +2983,18 @@ async fn execute_stream_from_frame_stream(
};
if rewritten_chunk.is_empty() {
if let Some(error_body_json) = provider_private_error_body_json {
let error_status_code = resolve_local_sync_error_status_code(
status_code,
&error_body_json,
);
terminal_failure =
Some(build_stream_failure_from_provider_error_body(
error_status_code,
&error_body_json,
));
break;
}
continue;
}
@@ -2935,6 +3057,15 @@ async fn execute_stream_from_frame_stream(
Ordering::Relaxed,
);
}
if let Some(error_body_json) = provider_private_error_body_json {
let error_status_code =
resolve_local_sync_error_status_code(status_code, &error_body_json);
terminal_failure = Some(build_stream_failure_from_provider_error_body(
error_status_code,
&error_body_json,
));
break;
}
}
StreamFramePayload::Telemetry {
telemetry: frame_telemetry,
@@ -2992,6 +3123,11 @@ async fn execute_stream_from_frame_stream(
if let Some(normalizer) = private_stream_normalizer.as_mut() {
match normalizer.finish() {
Ok(normalized_chunk) if !normalized_chunk.is_empty() => {
let provider_private_error_body_json =
extract_provider_private_stream_error_body(
stream_usage_report_context.as_ref(),
&normalized_chunk,
);
if let (Some(observer), Some(report_context)) = (
stream_usage_observer.as_mut(),
stream_usage_report_context.as_ref(),
@@ -3065,6 +3201,16 @@ async fn execute_stream_from_frame_stream(
);
}
}
if let Some(error_body_json) = provider_private_error_body_json {
let error_status_code =
resolve_local_sync_error_status_code(status_code, &error_body_json);
terminal_failure.get_or_insert_with(|| {
build_stream_failure_from_provider_error_body(
error_status_code,
&error_body_json,
)
});
}
}
Ok(_) => {}
Err(err) => {
@@ -3478,6 +3624,7 @@ mod tests {
use axum::extract::Request;
use axum::routing::any;
use axum::{http::header, http::HeaderValue, Router};
use base64::Engine as _;
use futures_util::StreamExt as _;
use serde_json::{json, Value};
use tokio::sync::{mpsc, watch, Notify};
@@ -3528,6 +3675,14 @@ mod tests {
}
}
fn connect_json_frame(flags: u8, payload: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(5 + payload.len());
out.push(flags);
out.extend_from_slice(&(payload.len() as u32).to_be_bytes());
out.extend_from_slice(payload);
out
}
#[test]
fn merge_stream_terminal_summary_prefers_more_complete_observed_usage() {
let mut runtime_usage = StandardizedUsage::new();
@@ -4412,6 +4567,268 @@ mod tests {
assert!(text.contains("\"type\":\"image_stream_total_timeout\""));
}
#[tokio::test]
async fn execute_stream_from_frame_stream_treats_windsurf_connect_trailer_error_as_failure() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let plan = ExecutionPlan {
request_id: "req-windsurf-connect-error".into(),
candidate_id: Some("cand-windsurf-connect-error".into()),
provider_name: Some("windsurf".into()),
provider_id: "provider-windsurf".into(),
endpoint_id: "endpoint-windsurf-chat".into(),
key_id: "key-windsurf".into(),
method: "POST".into(),
url: "https://server.codeium.com/exa.api_server_pb.ApiServerService/GetChatMessage?beta=true".into(),
headers: BTreeMap::from([
("content-type".into(), "application/connect+json".into()),
("accept".into(), "application/connect+json".into()),
]),
content_type: Some("application/connect+json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "claude-sonnet-4",
"messages": [],
"stream": true
})),
stream: true,
client_api_format: "claude:messages".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("claude-sonnet-4".into()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let trailer_error = connect_json_frame(
2,
br#"{"error":{"code":"resource_exhausted","message":"an internal error occurred"}}"#,
);
let trailer_error_b64 = base64::engine::general_purpose::STANDARD.encode(trailer_error);
let frame = format!(
"{{\"type\":\"data\",\"payload\":{{\"kind\":\"data\",\"chunk_b64\":\"{trailer_error_b64}\"}}}}\n"
);
let frame_stream = stream! {
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/connect+json\"}}}\n",
));
yield Ok::<Bytes, std::io::Error>(Bytes::from(frame));
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n",
));
}
.boxed();
let response = execute_stream_from_frame_stream(
&state,
plan,
"trace-windsurf-connect-error",
&test_decision(),
"claude_chat_stream",
Some("claude_chat_stream_success".to_string()),
Some(json!({
"request_id": "req-windsurf-connect-error",
"candidate_id": "cand-windsurf-connect-error",
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "openai:chat",
"client_api_format": "claude:messages",
"needs_conversion": true,
"has_envelope": true,
"envelope_name": "windsurf:GetChatMessage",
"local_failover_policy": {
"stop_status_codes": [429]
}
})),
crate::clock::current_unix_ms(),
Instant::now(),
frame_stream,
None,
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
let status = response.status();
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let body_json: Value =
serde_json::from_slice(&body).expect("response body should decode as json");
assert_eq!(status.as_u16(), 429);
assert_eq!(body_json["type"], json!("error"));
assert_eq!(body_json["error"]["type"], json!("rate_limit_error"));
assert_eq!(body_json["error"]["code"], json!("resource_exhausted"));
assert_eq!(
body_json["error"]["message"],
json!("an internal error occurred")
);
let candidates = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let candidates = request_candidate_repository
.list_by_request_id("req-windsurf-connect-error")
.await
.expect("request candidates should read");
if candidates
.first()
.is_some_and(|candidate| candidate.status == RequestCandidateStatus::Failed)
{
break candidates;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("candidate should be marked failed");
assert_eq!(candidates[0].status_code, Some(429));
assert_eq!(
candidates[0].error_type.as_deref(),
Some("resource_exhausted")
);
}
#[tokio::test]
async fn execute_stream_from_frame_stream_decodes_non_success_windsurf_connect_error_body() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let plan = ExecutionPlan {
request_id: "req-windsurf-connect-429".into(),
candidate_id: Some("cand-windsurf-connect-429".into()),
provider_name: Some("windsurf".into()),
provider_id: "provider-windsurf".into(),
endpoint_id: "endpoint-windsurf-chat".into(),
key_id: "key-windsurf".into(),
method: "POST".into(),
url: "https://server.codeium.com/exa.api_server_pb.ApiServerService/GetChatMessage?beta=true".into(),
headers: BTreeMap::from([
("content-type".into(), "application/connect+json".into()),
("accept".into(), "application/connect+json".into()),
]),
content_type: Some("application/connect+json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "claude-sonnet-4",
"messages": [],
"stream": true
})),
stream: true,
client_api_format: "claude:messages".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("claude-sonnet-4".into()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let connect_error = connect_json_frame(
2,
br#"{"error":{"code":"resource_exhausted","message":"quota exhausted"}}"#,
);
let connect_error_b64 = base64::engine::general_purpose::STANDARD.encode(connect_error);
let frame_stream = stream! {
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":429,\"headers\":{\"content-type\":\"application/connect+json\"}}}\n",
));
yield Ok::<Bytes, std::io::Error>(Bytes::from(format!(
"{{\"type\":\"data\",\"payload\":{{\"kind\":\"data\",\"chunk_b64\":\"{connect_error_b64}\"}}}}\n"
)));
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n",
));
}
.boxed();
let response = execute_stream_from_frame_stream(
&state,
plan,
"trace-windsurf-connect-429",
&test_decision(),
"claude_chat_stream",
Some("claude_chat_stream_success".to_string()),
Some(json!({
"request_id": "req-windsurf-connect-429",
"candidate_id": "cand-windsurf-connect-429",
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "openai:chat",
"client_api_format": "claude:messages",
"needs_conversion": true,
"has_envelope": true,
"envelope_name": "windsurf:GetChatMessage",
"local_failover_policy": {
"stop_status_codes": [429]
}
})),
crate::clock::current_unix_ms(),
Instant::now(),
frame_stream,
None,
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
assert_eq!(response.status().as_u16(), 429);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let body_json: Value =
serde_json::from_slice(&body).expect("response body should decode as json");
assert_eq!(body_json["type"], json!("error"));
assert_eq!(body_json["error"]["type"], json!("rate_limit_error"));
assert_eq!(body_json["error"]["code"], json!("resource_exhausted"));
assert_eq!(body_json["error"]["message"], json!("quota exhausted"));
let record = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if let Some(usage) = usage_repository
.find_by_request_id("req-windsurf-connect-429")
.await
.expect("usage should read")
.filter(|usage| usage.status == "failed")
{
break usage;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("usage should be written");
assert_eq!(record.status_code, Some(429));
assert_eq!(
record
.response_body
.as_ref()
.and_then(|body| body.get("error"))
.and_then(|error| error.get("code")),
Some(&json!("resource_exhausted"))
);
assert!(record.response_body_ref.is_none());
}
#[tokio::test]
async fn execute_stream_from_frame_stream_stops_upstream_when_client_drops_body() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
@@ -5560,8 +5977,9 @@ mod tests {
let body = body_task.await.expect("body task should complete");
assert!(body.contains("data: hello\n\n"));
assert!(body.contains("event: aether.error\n"));
assert!(body.contains("data: {\"error\":"));
assert!(body.contains(original_error));
assert!(body.contains("data: [DONE]\n\n"));
assert!(
!body.contains("unexpected EOF during chunk size line"),
"same-format SSE path should surface the original terminal error event"

91
apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs
View File

@@ -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,

5
apps/aether-gateway/src/execution_runtime/submission.rs
View File

@@ -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;
}

90
apps/aether-gateway/src/execution_runtime/sync/execution.rs
View File

@@ -20,7 +20,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;
@@ -30,7 +29,8 @@ 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,
maybe_build_sync_finalize_outcome, extract_provider_private_stream_error_body,
LocalCoreSyncErrorKind, LocalCoreSyncFinalizeOutcome,
};
use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
@@ -43,12 +43,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::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,
@@ -398,6 +402,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,
@@ -788,6 +829,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)
@@ -945,27 +992,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(),
@@ -1696,8 +1725,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 {
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,

121
apps/aether-gateway/src/execution_runtime/transport.rs
View File

@@ -27,6 +27,7 @@ use thiserror::Error;
#[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 +233,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 +330,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 +555,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(),
@@ -1486,6 +1456,50 @@ pub(crate) fn response_body_is_json(headers: &BTreeMap<String, String>, body_byt
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) =
aether_ai_formats::api::extract_provider_private_stream_error_body(None, decoded_body_bytes)
.or_else(|| {
aether_ai_formats::api::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;
@@ -1510,7 +1524,8 @@ mod tests {
use tokio::sync::watch;
use super::{
build_browser_wreq_client, build_client, build_request_headers, execute_sync_plan,
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,
@@ -2780,6 +2795,30 @@ mod tests {
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()

4947
apps/aether-gateway/src/execution_runtime/windsurf.rs Normal file
View File

File diff suppressed because it is too large Load Diff

88
apps/aether-gateway/src/handlers/admin/observability/monitoring/tests/trace.rs
View File

@@ -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(

17
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs
View File

@@ -3158,17 +3158,22 @@ async fn provider_query_execute_windsurf_test_candidate(
}
let incoming_request_headers = provider_query_extract_request_headers(payload);
let mut request_body = original_request_body.clone();
if let Some(object) = request_body.as_object_mut() {
object.insert("stream".to_string(), Value::Bool(false));
}
let request_body = original_request_body.clone();
let request_model =
provider_query_request_body_model(&request_body, &candidate.effective_model);
let upstream_is_stream = provider_query_resolve_standard_test_upstream_is_stream(
transport.endpoint.config.as_ref(),
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),

151
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/summary.rs
View File

@@ -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]"))
);
}
}

192
apps/aether-gateway/src/handlers/shared/catalog.rs
View File

@@ -1268,17 +1268,11 @@ fn build_windsurf_quota_status_snapshot(
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);
let limited = rate_limit_object
.get("limited")
.or_else(|| rate_limit_object.get("is_limited"))
.or_else(|| rate_limit_object.get("isLimited"))
.and_then(admin_provider_quota_pure::coerce_json_bool)
== Some(true);
if limited || retry_after_ms.is_some() {
.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 =
retry_after_ms.map(|value| value.saturating_add(999) / 1000);
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)
@@ -1674,6 +1668,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)
@@ -1699,6 +1699,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,
@@ -2671,6 +2726,123 @@ mod tests {
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();

116
apps/aether-gateway/src/tests/control/admin/provider_query.rs
View File

@@ -2383,6 +2383,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(