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

feat(rust): 支持 image 同步转流式 SSE、free_team_first 调度预设及账户错误自动重试

Browse Source 
- openai:image 同步响应桥接为流式 SSE（image_generation.completed / image_edit.completed 事件）
- image 请求解析与前置校验移除模型白名单，支持任意自定义模型名
- 调度器新增 free_team_first 预设（mode: both / free_only / team_only）
- pool config 解析重构：新增 POOL_ALLOWED_SCHEDULING_PRESETS 白名单，规范化 mode 字段
- 错误分类器新增账户/账单错误模式集，升级为 RetryUpstreamFailure 而非 StopSemanticClientError
- 修复 stream execution 中上游 headers 与输出 headers 混用导致 content-type 判断错误的问题
- codex image 工具始终写入 action 字段（generate/edit），仅 generate 操作填充默认 size/quality
- 前端：pool 节点组只展示实际执行过的候选节点，全部 skipped 时折叠为最后一个节点
- Redis 测试就绪检测从 TCP 连接改为 PING/PONG 协议验证
This commit is contained in:
fawney19
2026-04-24 10:05:55 +08:00
parent a9d10163af
commit 5d710d9d10
16 changed files with 1388 additions and 204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

205
apps/aether-gateway/src/ai_pipeline/finalize/internal/sync_to_stream.rs
View File

@@ -1,6 +1,7 @@
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
use serde_json::{json, Value};
use crate::ai_pipeline::finalize::sse::encode_json_sse;
use crate::ai_pipeline::{
convert_claude_cli_response_to_openai_cli, convert_gemini_cli_response_to_openai_cli,
convert_openai_chat_response_to_openai_cli, ClaudeClientEmitter, GeminiClientEmitter,
@@ -21,6 +22,9 @@ pub(crate) fn maybe_bridge_standard_sync_json_to_stream(
) -> Result<Option<SyncToStreamBridgeOutcome>, GatewayError> {
let provider_api_format = normalize_api_format(provider_api_format);
let client_api_format = normalize_api_format(client_api_format);
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
return maybe_bridge_openai_image_sync_json_to_stream(provider_body_json, report_context);
}
if !is_standard_api_format(provider_api_format.as_str())
|| !is_standard_api_format(client_api_format.as_str())
{
@@ -51,6 +55,56 @@ pub(crate) fn maybe_bridge_standard_sync_json_to_stream(
}))
}
fn maybe_bridge_openai_image_sync_json_to_stream(
provider_body_json: &Value,
report_context: Option<&Value>,
) -> Result<Option<SyncToStreamBridgeOutcome>, GatewayError> {
let Some(response) = provider_body_json.as_object() else {
return Ok(None);
};
let Some(image) = response
.get("data")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
.find_map(extract_openai_image_sync_b64_json)
else {
return Ok(None);
};
let usage = response.get("usage").cloned().unwrap_or(Value::Null);
let event_name = openai_image_completed_event_name(report_context);
let sse_body = encode_json_sse(
Some(event_name),
&json!({
"type": event_name,
"b64_json": image,
"usage": usage,
}),
)?;
Ok(Some(SyncToStreamBridgeOutcome {
sse_body,
terminal_summary: Some(ExecutionStreamTerminalSummary {
standardized_usage: response
.get("usage")
.and_then(standardized_usage_from_openai_usage),
finish_reason: Some("stop".to_string()),
response_id: response
.get("id")
.and_then(Value::as_str)
.map(ToOwned::to_owned),
model: response
.get("model")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.or_else(|| image_bridge_model(report_context)),
observed_finish: true,
parser_error: None,
}),
}))
}
fn normalize_api_format(value: &str) -> String {
value.trim().to_ascii_lowercase()
}
@@ -68,6 +122,57 @@ fn is_standard_api_format(value: &str) -> bool {
)
}
fn extract_openai_image_sync_b64_json(item: &serde_json::Map<String, Value>) -> Option<String> {
item.get("b64_json")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| {
item.get("url")
.and_then(Value::as_str)
.and_then(extract_base64_from_data_url)
})
}
fn extract_base64_from_data_url(value: &str) -> Option<String> {
let trimmed = value.trim();
let (metadata, payload) = trimmed.split_once(',')?;
if !metadata.starts_with("data:") || !metadata.ends_with(";base64") {
return None;
}
(!payload.trim().is_empty()).then(|| payload.trim().to_string())
}
fn openai_image_completed_event_name(report_context: Option<&Value>) -> &'static str {
if openai_image_request_operation(report_context) == Some("edit") {
"image_edit.completed"
} else {
"image_generation.completed"
}
}
fn openai_image_request_operation(report_context: Option<&Value>) -> Option<&str> {
report_context
.and_then(|value| value.get("image_request"))
.and_then(|value| value.get("operation"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn image_bridge_model(report_context: Option<&Value>) -> Option<String> {
report_context.and_then(|context| {
context
.get("mapped_model")
.or_else(|| context.get("model"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn build_bridge_report_context(
report_context: Option<&Value>,
provider_api_format: &str,
@@ -333,3 +438,103 @@ fn standardized_usage_from_openai_usage(value: &Value) -> Option<StandardizedUsa
.insert("total_tokens".to_string(), json!(total_tokens));
Some(standardized_usage.normalize_cache_creation_breakdown())
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::maybe_bridge_standard_sync_json_to_stream;
fn utf8(bytes: Vec<u8>) -> String {
String::from_utf8(bytes).expect("utf8 should decode")
}
#[test]
fn bridges_openai_image_sync_json_to_generation_completed_sse() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"mapped_model": "gpt-image-1",
"image_request": {
"operation": "generate"
}
});
let outcome = maybe_bridge_standard_sync_json_to_stream(
&json!({
"created": 1776971267,
"data": [{
"b64_json": "aGVsbG8="
}],
"usage": {
"total_tokens": 100,
"input_tokens": 50,
"output_tokens": 50,
"input_tokens_details": {
"text_tokens": 10,
"image_tokens": 40
}
}
}),
"openai:image",
"openai:image",
Some(&report_context),
)
.expect("bridge should succeed")
.expect("bridge should produce sse");
let output = utf8(outcome.sse_body);
assert!(output.contains("event: image_generation.completed"));
assert!(output.contains("\"type\":\"image_generation.completed\""));
assert!(output.contains("\"b64_json\":\"aGVsbG8=\""));
assert!(output.contains("\"total_tokens\":100"));
let summary = outcome
.terminal_summary
.expect("terminal summary should exist");
assert_eq!(summary.model.as_deref(), Some("gpt-image-1"));
assert_eq!(summary.finish_reason.as_deref(), Some("stop"));
assert_eq!(
summary
.standardized_usage
.as_ref()
.and_then(|usage| usage.dimensions.get("total_tokens"))
.cloned(),
Some(json!(100))
);
}
#[test]
fn bridges_openai_image_sync_data_url_to_edit_completed_sse() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"image_request": {
"operation": "edit"
}
});
let outcome = maybe_bridge_standard_sync_json_to_stream(
&json!({
"created": 1776971267,
"data": [{
"url": "data:image/webp;base64,d29ybGQ="
}],
"usage": {
"total_tokens": 9,
"input_tokens": 4,
"output_tokens": 5
}
}),
"openai:image",
"openai:image",
Some(&report_context),
)
.expect("bridge should succeed")
.expect("bridge should produce sse");
let output = utf8(outcome.sse_body);
assert!(output.contains("event: image_edit.completed"));
assert!(output.contains("\"type\":\"image_edit.completed\""));
assert!(output.contains("\"b64_json\":\"d29ybGQ=\""));
assert!(output.contains("\"total_tokens\":9"));
}
}

123
apps/aether-gateway/src/ai_pipeline/planner/pool_scheduler.rs
View File

@@ -593,6 +593,7 @@ fn build_pool_sort_vectors(
"cache_affinity" => cache_affinity_ranks.clone(),
"priority_first" => priority_first_ranks(items, &lru_ranks),
"single_account" => single_account_ranks(items),
"free_team_first" => plan_ranks(items, &lru_ranks, preset.mode.as_deref()),
"plus_first" => plan_ranks(items, &lru_ranks, Some("plus_only")),
"free_first" => plan_ranks(items, &lru_ranks, Some("free_only")),
"team_first" => plan_ranks(items, &lru_ranks, Some("team_only")),
@@ -991,7 +992,7 @@ fn normalize_enabled_pool_presets(
fn pool_preset_supported_for_provider(preset: &str, provider_type: &str) -> bool {
match preset {
"free_first" | "plus_first" | "recent_refresh" | "team_first" => {
"free_first" | "free_team_first" | "plus_first" | "recent_refresh" | "team_first" => {
matches!(provider_type, "codex" | "kiro")
}
_ => true,
@@ -1392,6 +1393,126 @@ mod tests {
);
}
#[test]
fn pool_scheduler_supports_free_team_first_modes() {
let key_plus = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-plus",
10,
Some(json!({
"pool_advanced": {
"scheduling_presets": [{"preset": "free_team_first", "enabled": true, "mode": "team_only"}]
}
})),
);
let key_free = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-free",
10,
Some(json!({
"pool_advanced": {
"scheduling_presets": [{"preset": "free_team_first", "enabled": true, "mode": "team_only"}]
}
})),
);
let key_team = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-team",
10,
Some(json!({
"pool_advanced": {
"scheduling_presets": [{"preset": "free_team_first", "enabled": true, "mode": "team_only"}]
}
})),
);
let key_context_by_id = BTreeMap::from([
(
"key-plus".to_string(),
PoolCatalogKeyContext {
oauth_plan_type: Some("plus".to_string()),
..PoolCatalogKeyContext::default()
},
),
(
"key-free".to_string(),
PoolCatalogKeyContext {
oauth_plan_type: Some("free".to_string()),
..PoolCatalogKeyContext::default()
},
),
(
"key-team".to_string(),
PoolCatalogKeyContext {
oauth_plan_type: Some("team".to_string()),
..PoolCatalogKeyContext::default()
},
),
]);
let (reordered, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
vec![key_plus, key_free, key_team],
&BTreeMap::new(),
&key_context_by_id,
);
assert!(skipped.is_empty());
assert_eq!(
reordered
.iter()
.map(|item| item.candidate.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-team", "key-free", "key-plus"]
);
}
#[test]
fn pool_scheduler_defaults_empty_pool_advanced_to_cache_affinity() {
let key_a = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10,
Some(json!({ "pool_advanced": {} })),
);
let key_b = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-b",
10,
Some(json!({ "pool_advanced": {} })),
);
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AdminProviderPoolRuntimeState {
lru_score_by_key: BTreeMap::from([
("key-a".to_string(), 10.0),
("key-b".to_string(), 200.0),
]),
..AdminProviderPoolRuntimeState::default()
},
)]);
let (reordered, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
vec![key_a, key_b],
&runtime_by_provider,
&BTreeMap::new(),
);
assert!(skipped.is_empty());
assert_eq!(
reordered
.iter()
.map(|item| item.candidate.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-b", "key-a"]
);
}
#[test]
fn normalizes_distribution_mutex_group_to_first_enabled_member() {
let presets = normalize_enabled_pool_presets(

95
apps/aether-gateway/src/ai_pipeline/planner/specialized/image/request.rs
View File

@@ -118,10 +118,10 @@ pub(super) fn resolve_requested_image_model_for_request(
.iter()
.find(|field| field.name.trim() == "model")
.map(|field| String::from_utf8_lossy(&field.data).trim().to_string());
normalize_requested_image_model(model.as_deref())?
normalize_requested_image_model(model.as_deref())
.or_else(|| Some(default_model_for_operation(operation).to_string()))
} else {
normalize_requested_image_model(body_json.get("model").and_then(Value::as_str))?
normalize_requested_image_model(body_json.get("model").and_then(Value::as_str))
.or_else(|| Some(default_model_for_operation(operation).to_string()))
}
}
@@ -357,13 +357,7 @@ fn normalize_openai_image_json_request(
return None;
}
let requested_model =
normalize_requested_image_model(object.get("model").and_then(Value::as_str))?;
if requested_model
.as_deref()
.is_some_and(|model| !image_model_supported_for_operation(operation, model))
{
return None;
}
normalize_requested_image_model(object.get("model").and_then(Value::as_str));
let prompt = normalize_prompt(object.get("prompt"), operation)?;
let response_format =
normalize_image_response_format(object.get("response_format").and_then(Value::as_str))?;
@@ -425,13 +419,7 @@ fn normalize_openai_image_multipart_request(
let multipart_fields = parse_multipart_fields_from_base64(parts, body_base64)?;
let requested_model = normalize_requested_image_model(
find_multipart_text_field(&multipart_fields, "model").as_deref(),
)?;
if requested_model
.as_deref()
.is_some_and(|model| !image_model_supported_for_operation(operation, model))
{
return None;
}
);
if find_multipart_text_field(&multipart_fields, "style").is_some() {
return None;
}
@@ -528,11 +516,11 @@ fn normalize_openai_image_multipart_request(
})
}
fn normalize_requested_image_model(value: Option<&str>) -> Option<Option<String>> {
let Some(model) = value.map(str::trim).filter(|value| !value.is_empty()) else {
return Some(None);
};
canonicalize_image_model(model).map(|canonical| Some(canonical.to_string()))
fn normalize_requested_image_model(value: Option<&str>) -> Option<String> {
value
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn default_model_for_operation(operation: OpenAiImageOperation) -> &'static str {
@@ -544,27 +532,6 @@ fn default_model_for_operation(operation: OpenAiImageOperation) -> &'static str
}
}
fn canonicalize_image_model(model: &str) -> Option<&'static str> {
match model.trim().to_ascii_lowercase().as_str() {
"gpt-image-1" => Some("gpt-image-1"),
"gpt-image-1.5" => Some("gpt-image-1.5"),
"gpt-image-1-mini" => Some("gpt-image-1-mini"),
"gpt-image-2" => Some("gpt-image-2"),
"chatgpt-image-latest" => Some("chatgpt-image-latest"),
"dall-e-2" => Some("dall-e-2"),
"dall-e-3" => Some("dall-e-3"),
_ => None,
}
}
fn image_model_supported_for_operation(operation: OpenAiImageOperation, model: &str) -> bool {
match operation {
OpenAiImageOperation::Generate => true,
OpenAiImageOperation::Edit => !matches!(model, "dall-e-3"),
OpenAiImageOperation::Variation => model == "dall-e-2",
}
}
fn normalize_prompt(
value: Option<&Value>,
operation: OpenAiImageOperation,
@@ -696,9 +663,16 @@ fn build_tool_options(
"type".to_string(),
Value::String("image_generation".to_string()),
);
if operation != OpenAiImageOperation::Generate {
tool.insert("action".to_string(), Value::String("edit".to_string()));
}
tool.insert(
"action".to_string(),
Value::String(
match operation {
OpenAiImageOperation::Generate => "generate",
OpenAiImageOperation::Edit | OpenAiImageOperation::Variation => "edit",
}
.to_string(),
),
);
for (key, value) in raw_values {
let normalized = match key.as_str() {
"size" | "background" | "moderation" | "input_fidelity" => {
@@ -1097,6 +1071,25 @@ mod tests {
);
}
#[test]
fn normalize_generate_json_request_accepts_custom_model_name() {
let parts = request_parts("/v1/images/generations", Some("application/json"));
let request = normalize_openai_image_request(
&parts,
&json!({
"model": " Custom/Image-Model:V1 ",
"prompt": "generate image"
}),
None,
)
.expect("custom image model request should normalize");
assert_eq!(
request.requested_model.as_deref(),
Some("Custom/Image-Model:V1")
);
}
#[test]
fn build_generate_request_defaults_codex_image_tool_and_tool_choice() {
let parts = request_parts("/v1/images/generations", Some("application/json"));
@@ -1114,6 +1107,10 @@ mod tests {
assert!(request.tool.get("quality").is_none());
assert!(request.tool.get("background").is_none());
assert!(request.tool.get("output_format").is_none());
assert_eq!(
request.tool.get("action").and_then(|value| value.as_str()),
Some("generate")
);
let mut provider_request_body = build_provider_request_body(&request);
assert!(provider_request_body.get("model").is_none());
@@ -1135,6 +1132,14 @@ mod tests {
.and_then(|value| value.as_str()),
Some("image_generation")
);
assert_eq!(
provider_request_body
.get("tools")
.and_then(|value| value.get(0))
.and_then(|value| value.get("action"))
.and_then(|value| value.as_str()),
Some("generate")
);
assert_eq!(
provider_request_body
.get("tools")

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

@@ -309,6 +309,7 @@ pub(crate) fn resolve_core_stream_error_finalize_report_kind(
pub(crate) fn resolve_core_stream_direct_finalize_report_kind(plan_kind: &str) -> Option<String> {
let report_kind = match plan_kind {
"openai_chat_stream" => "openai_chat_sync_finalize",
"openai_image_stream" => "openai_image_sync_finalize",
"claude_chat_stream" => "claude_chat_sync_finalize",
"gemini_chat_stream" => "gemini_chat_sync_finalize",
"openai_cli_stream" => "openai_cli_sync_finalize",

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

@@ -895,6 +895,7 @@ async fn execute_stream_from_frame_stream(
let direct_stream_finalize_kind = resolve_core_stream_direct_finalize_report_kind(plan_kind);
let normalized_stream_report_context =
normalize_provider_private_report_context(report_context.as_ref());
let upstream_headers = headers.clone();
let mut private_stream_normalizer =
maybe_build_provider_private_stream_normalizer(report_context.as_ref());
let mut local_stream_rewriter =
@@ -904,10 +905,10 @@ async fn execute_stream_from_frame_stream(
headers.remove("content-length");
headers.insert("content-type".to_string(), "text/event-stream".to_string());
}
let content_type = headers.get("content-type").map(String::as_str);
let upstream_content_type = upstream_headers.get("content-type").map(String::as_str);
let skip_direct_finalize_prefetch = should_skip_direct_finalize_prefetch(
direct_stream_finalize_kind.as_deref(),
content_type,
upstream_content_type,
plan.provider_api_format.as_str(),
plan.client_api_format.as_str(),
private_stream_normalizer.is_some(),
@@ -933,7 +934,7 @@ async fn execute_stream_from_frame_stream(
key_id = %plan.key_id,
model_name,
candidate_index = candidate_index.as_str(),
content_type = content_type.unwrap_or("-"),
content_type = upstream_content_type.unwrap_or("-"),
provider_api_format = plan.provider_api_format.as_str(),
client_api_format = plan.client_api_format.as_str(),
"gateway skipped direct finalize prefetch for same-format passthrough stream"
@@ -1012,8 +1013,10 @@ async fn execute_stream_from_frame_stream(
provider_prefetched_body.extend_from_slice(&chunk);
prefetched_inspection_body.extend_from_slice(&chunk);
let inspection =
inspect_prefetched_stream_body(&headers, &prefetched_inspection_body);
let inspection = inspect_prefetched_stream_body(
&upstream_headers,
&prefetched_inspection_body,
);
match inspection {
StreamPrefetchInspection::EmbeddedError(body_json) => {
debug!(
@@ -1062,7 +1065,8 @@ async fn execute_stream_from_frame_stream(
StreamPrefetchInspection::NonError => {}
}
if !response_headers_indicate_sse(&headers) && (200..300).contains(&status_code)
if !response_headers_indicate_sse(&upstream_headers)
&& (200..300).contains(&status_code)
{
if let Some(body_json) =
parse_prefetched_sync_json_body(&prefetched_inspection_body)
@@ -1285,7 +1289,7 @@ async fn execute_stream_from_frame_stream(
let request_id_for_report_log = short_request_id(&request_id);
let candidate_id_for_report = candidate_id.clone();
let emit_passthrough_sse_terminal_error =
skip_direct_finalize_prefetch && response_headers_indicate_sse(&headers);
skip_direct_finalize_prefetch && response_headers_indicate_sse(&upstream_headers);
let body_capture_policy = match UsageRuntimeAccess::body_capture_policy(state.data.as_ref())
.await
{
@@ -2138,6 +2142,119 @@ mod tests {
server.abort();
}
#[tokio::test]
async fn execute_execution_runtime_stream_bridges_openai_image_sync_json_from_remote_runtime_to_image_sse(
) {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let server = tokio::spawn(async move {
let app = Router::new().route(
"/v1/execute/stream",
any(|_request: Request| async move {
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/json\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"created\\\":1776972364,\\\"data\\\":[{\\\"b64_json\\\":\\\"aGVsbG8=\\\"}],\\\"usage\\\":{\\\"total_tokens\\\":100,\\\"input_tokens\\\":50,\\\"output_tokens\\\":50,\\\"input_tokens_details\\\":{\\\"text_tokens\\\":10,\\\"image_tokens\\\":40}}}\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = axum::http::Response::new(Body::from(frames));
response.headers_mut().insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}),
);
axum::serve(listener, app)
.await
.expect("server should start");
});
let state = AppState::new()
.expect("app state should build")
.with_execution_runtime_override_base_url(format!("http://{addr}"));
let plan = ExecutionPlan {
request_id: "req-remote-runtime-image-sync-json-stream".into(),
candidate_id: Some("cand-remote-runtime-image-sync-json-stream".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://chatgpt.com/backend-api/codex/responses".into(),
headers: BTreeMap::from([
("content-type".into(), "application/json".into()),
("accept".into(), "text/event-stream".into()),
]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-image-1",
"prompt": "hello",
"stream": true
})),
stream: true,
client_api_format: "openai:image".into(),
provider_api_format: "openai:image".into(),
model_name: Some("gpt-image-1".into()),
proxy: None,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let decision = GatewayControlDecision::synthetic(
"/v1/images/generations",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("image".to_string()),
Some("openai:image".to_string()),
)
.with_execution_runtime_candidate(true);
let response = execute_execution_runtime_stream(
&state,
plan,
"trace-remote-runtime-image-sync-json-stream",
&decision,
"openai_image_stream",
None,
Some(json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"mapped_model": "gpt-image-1",
"image_request": {
"operation": "generate"
}
})),
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
assert_eq!(
response
.headers()
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("text/event-stream")
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let text = String::from_utf8(body.to_vec()).expect("response body should be utf8");
assert!(text.contains("event: image_generation.completed"));
assert!(text.contains("\"type\":\"image_generation.completed\""));
assert!(text.contains("\"b64_json\":\"aGVsbG8=\""));
assert!(text.contains("\"total_tokens\":100"));
server.abort();
}
#[tokio::test]
async fn execute_execution_runtime_stream_returns_client_error_with_local_tunnel_message_before_first_data(
) {

140
apps/aether-gateway/src/execution_runtime/stream_pump.rs
View File

@@ -1014,6 +1014,146 @@ mod tests {
);
}
#[tokio::test]
async fn direct_execution_frame_stream_bridges_openai_image_sync_json_to_image_sse() {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let server = tokio::spawn(async move {
let app = Router::new().route(
"/responses",
post(|| async {
let body = serde_json::json!({
"created": 1776971267_u64,
"data": [{
"b64_json": "aGVsbG8="
}],
"usage": {
"total_tokens": 100,
"input_tokens": 50,
"output_tokens": 50,
"input_tokens_details": {
"text_tokens": 10,
"image_tokens": 40
}
}
});
let mut response = axum::http::Response::new(Body::from(
serde_json::to_vec(&body).expect("json should encode"),
));
response.headers_mut().insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
response
}),
);
axum::serve(listener, app)
.await
.expect("server should start");
});
let runtime = DirectSyncExecutionRuntime::new();
let execution = runtime
.execute_stream(&ExecutionPlan {
request_id: "req-image-sync-bridge".to_string(),
candidate_id: Some("cand-image-sync-bridge".to_string()),
provider_name: Some("OpenAI".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: format!("http://{addr}/responses"),
headers: BTreeMap::new(),
content_type: None,
content_encoding: None,
body: RequestBody::from_json(serde_json::json!({
"model": "gpt-image-1",
"prompt": "poster",
"stream": true
})),
stream: true,
client_api_format: "openai:image".to_string(),
provider_api_format: "openai:image".to_string(),
model_name: Some("gpt-image-1".into()),
proxy: None,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
})
.await
.expect("stream execution should succeed");
let frames = build_direct_execution_frame_stream(execution)
.map(|item| item.expect("frame should encode"))
.collect::<Vec<_>>()
.await
.into_iter()
.map(|bytes| String::from_utf8(bytes.to_vec()).expect("frame should be utf8"))
.collect::<Vec<_>>();
server.abort();
let header_frame: Value =
serde_json::from_str(&frames[0]).expect("headers frame should parse");
assert_eq!(
header_frame
.get("payload")
.and_then(|payload| payload.get("headers"))
.and_then(|headers| headers.get("content-type"))
.and_then(Value::as_str),
Some("text/event-stream")
);
let data_frame = frames
.iter()
.map(|line| serde_json::from_str::<Value>(line).expect("frame should parse"))
.find(|frame| frame.get("type").and_then(Value::as_str) == Some("data"))
.expect("data frame should exist");
let bridged_body = base64::engine::general_purpose::STANDARD
.decode(
data_frame
.get("payload")
.and_then(|payload| payload.get("chunk_b64"))
.and_then(Value::as_str)
.expect("chunk_b64 should exist"),
)
.expect("data frame should decode");
let bridged_text = String::from_utf8(bridged_body).expect("bridged body should be utf8");
assert!(bridged_text.contains("event: image_generation.completed"));
assert!(bridged_text.contains("\"type\":\"image_generation.completed\""));
assert!(bridged_text.contains("\"b64_json\":\"aGVsbG8=\""));
assert!(bridged_text.contains("\"total_tokens\":100"));
let eof_frame = frames
.iter()
.map(|line| serde_json::from_str::<Value>(line).expect("frame should parse"))
.find(|frame| frame.get("type").and_then(Value::as_str) == Some("eof"))
.expect("eof frame should exist");
assert_eq!(
eof_frame
.get("payload")
.and_then(|payload| payload.get("summary"))
.and_then(|summary| summary.get("model"))
.and_then(Value::as_str),
Some("gpt-image-1")
);
assert_eq!(
eof_frame
.get("payload")
.and_then(|payload| payload.get("summary"))
.and_then(|summary| summary.get("standardized_usage"))
.and_then(|usage| usage.get("dimensions"))
.and_then(|dimensions| dimensions.get("total_tokens"))
.and_then(Value::as_i64),
Some(100)
);
}
#[tokio::test]
async fn direct_execution_frame_stream_preserves_local_tunnel_stream_error_message() {
let state = AppState::new().expect("app state should build");

258
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
View File

@@ -3,48 +3,67 @@ use crate::handlers::admin::provider::shared::support::{
};
use serde_json::{Map, Value};
const POOL_ALLOWED_SCHEDULING_PRESETS: &[&str] = &[
"lru",
"cache_affinity",
"load_balance",
"single_account",
"priority_first",
"free_team_first",
"free_first",
"team_first",
"plus_first",
"health_first",
"latency_first",
"cost_first",
"quota_balanced",
"recent_refresh",
];
fn json_u64(value: &Value) -> Option<u64> {
value
.as_u64()
.or_else(|| value.as_i64().and_then(|raw| u64::try_from(raw).ok()))
}
fn parse_pool_scheduling_presets(
raw_pool_advanced: &Map<String, Value>,
) -> Vec<AdminProviderPoolSchedulingPreset> {
let Some(presets) = raw_pool_advanced
.get("scheduling_presets")
.and_then(Value::as_array)
else {
return raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
.filter(|enabled| *enabled)
.map(|_| {
vec![AdminProviderPoolSchedulingPreset {
preset: "lru".to_string(),
enabled: true,
mode: None,
}]
})
.unwrap_or_default();
};
let mut normalized = Vec::new();
for item in presets {
if let Some(preset) = item
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
{
normalized.push(AdminProviderPoolSchedulingPreset {
preset: preset.to_ascii_lowercase(),
enabled: true,
mode: None,
});
continue;
fn normalize_pool_preset_mode(preset: &str, raw_mode: Option<&Value>) -> Option<String> {
match preset {
"free_team_first" | "free_first" | "team_first" | "plus_first" => {
let default_mode = match preset {
"free_team_first" => "both",
"free_first" => "free_only",
"team_first" => "team_only",
"plus_first" => "plus_only",
_ => unreachable!("preset covered by outer match"),
};
let normalized = raw_mode
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase())
.filter(|value| match preset {
"free_team_first" => {
matches!(value.as_str(), "both" | "free_only" | "team_only")
}
"free_first" => value == "free_only",
"team_first" => value == "team_only",
"plus_first" => value == "plus_only",
_ => false,
})
.unwrap_or_else(|| default_mode.to_string());
Some(normalized)
}
_ => None,
}
}
fn parse_object_style_pool_scheduling_presets(
presets: &[Value],
) -> Vec<AdminProviderPoolSchedulingPreset> {
let mut normalized = Vec::new();
let mut seen = std::collections::BTreeSet::new();
for item in presets {
let Some(object) = item.as_object() else {
continue;
};
@@ -53,40 +72,116 @@ fn parse_pool_scheduling_presets(
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase())
else {
continue;
};
if !POOL_ALLOWED_SCHEDULING_PRESETS.contains(&preset.as_str())
|| !seen.insert(preset.clone())
{
continue;
}
normalized.push(AdminProviderPoolSchedulingPreset {
preset: preset.to_ascii_lowercase(),
mode: normalize_pool_preset_mode(&preset, object.get("mode")),
preset,
enabled: object
.get("enabled")
.and_then(Value::as_bool)
.unwrap_or(true),
mode: object
.get("mode")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase()),
});
}
if normalized.is_empty()
&& raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
.unwrap_or(false)
{
normalized.push(AdminProviderPoolSchedulingPreset {
if normalized.is_empty() {
vec![AdminProviderPoolSchedulingPreset {
preset: "lru".to_string(),
enabled: true,
mode: None,
}]
} else {
normalized
}
}
fn parse_legacy_string_style_pool_scheduling_presets(
raw_pool_advanced: &Map<String, Value>,
presets: &[Value],
) -> Vec<AdminProviderPoolSchedulingPreset> {
let lru_enabled = raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
.unwrap_or(true);
let mut normalized = vec![AdminProviderPoolSchedulingPreset {
preset: "lru".to_string(),
enabled: lru_enabled,
mode: None,
}];
let mut seen = std::collections::BTreeSet::from(["lru".to_string()]);
for item in presets {
let Some(preset) = item
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase())
else {
continue;
};
if preset == "lru"
|| !POOL_ALLOWED_SCHEDULING_PRESETS.contains(&preset.as_str())
|| !seen.insert(preset.clone())
{
continue;
}
normalized.push(AdminProviderPoolSchedulingPreset {
preset,
enabled: true,
mode: None,
});
}
normalized
}
fn parse_pool_scheduling_presets_from_legacy_fields(
raw_pool_advanced: &Map<String, Value>,
) -> Vec<AdminProviderPoolSchedulingPreset> {
if raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
.unwrap_or(false)
{
vec![AdminProviderPoolSchedulingPreset {
preset: "lru".to_string(),
enabled: true,
mode: None,
}]
} else {
vec![AdminProviderPoolSchedulingPreset {
preset: "cache_affinity".to_string(),
enabled: true,
mode: None,
}]
}
}
fn parse_pool_scheduling_presets(
raw_pool_advanced: &Map<String, Value>,
) -> Vec<AdminProviderPoolSchedulingPreset> {
match raw_pool_advanced
.get("scheduling_presets")
.and_then(Value::as_array)
{
Some(presets) if !presets.is_empty() => match presets.first() {
Some(Value::Object(_)) => parse_object_style_pool_scheduling_presets(presets),
Some(Value::String(_)) => {
parse_legacy_string_style_pool_scheduling_presets(raw_pool_advanced, presets)
}
_ => parse_pool_scheduling_presets_from_legacy_fields(raw_pool_advanced),
},
_ => parse_pool_scheduling_presets_from_legacy_fields(raw_pool_advanced),
}
}
fn parse_pool_unschedulable_rules(
raw_pool_advanced: &Map<String, Value>,
) -> Vec<AdminProviderPoolUnschedulableRule> {
@@ -115,16 +210,8 @@ fn parse_pool_unschedulable_rules(
}
fn admin_provider_pool_lru_enabled(
raw_pool_advanced: &Map<String, Value>,
scheduling_presets: &[AdminProviderPoolSchedulingPreset],
) -> bool {
if let Some(explicit) = raw_pool_advanced
.get("lru_enabled")
.and_then(Value::as_bool)
{
return explicit;
}
scheduling_presets
.iter()
.any(|item| item.enabled && item.preset.eq_ignore_ascii_case("lru"))
@@ -145,7 +232,11 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
let Some(pool_advanced) = raw_pool_advanced.as_object() else {
return Some(AdminProviderPoolConfig {
scheduling_presets: Vec::new(),
scheduling_presets: vec![AdminProviderPoolSchedulingPreset {
preset: "cache_affinity".to_string(),
enabled: true,
mode: None,
}],
unschedulable_rules: Vec::new(),
lru_enabled: false,
skip_exhausted_accounts: false,
@@ -167,7 +258,7 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
let unschedulable_rules = parse_pool_unschedulable_rules(pool_advanced);
Some(AdminProviderPoolConfig {
lru_enabled: admin_provider_pool_lru_enabled(pool_advanced, &scheduling_presets),
lru_enabled: admin_provider_pool_lru_enabled(&scheduling_presets),
scheduling_presets,
unschedulable_rules,
skip_exhausted_accounts: pool_advanced
@@ -316,6 +407,19 @@ mod tests {
assert_eq!(config.cost_limit_per_key_tokens, Some(4096));
}
#[test]
fn defaults_empty_pool_advanced_to_cache_affinity_preset() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {}
})))
.expect("pool config should parse");
assert!(!config.lru_enabled);
assert_eq!(config.scheduling_presets.len(), 1);
assert_eq!(config.scheduling_presets[0].preset, "cache_affinity");
assert!(config.scheduling_presets[0].enabled);
}
#[test]
fn parses_object_style_scheduling_presets_with_modes() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
@@ -339,6 +443,44 @@ mod tests {
);
}
#[test]
fn parses_legacy_string_style_scheduling_presets_like_python() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"lru_enabled": false,
"scheduling_presets": [
"free_team_first",
"recent_refresh",
"free_team_first"
]
}
})))
.expect("pool config should parse");
assert!(!config.lru_enabled);
assert_eq!(config.scheduling_presets.len(), 3);
assert_eq!(config.scheduling_presets[0].preset, "lru");
assert!(!config.scheduling_presets[0].enabled);
assert_eq!(config.scheduling_presets[1].preset, "free_team_first");
assert_eq!(config.scheduling_presets[2].preset, "recent_refresh");
}
#[test]
fn invalid_free_team_first_mode_defaults_to_both() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scheduling_presets": [
{"preset": "free_team_first", "enabled": true, "mode": "invalid_mode"}
]
}
})))
.expect("pool config should parse");
assert_eq!(config.scheduling_presets.len(), 1);
assert_eq!(config.scheduling_presets[0].preset, "free_team_first");
assert_eq!(config.scheduling_presets[0].mode.as_deref(), Some("both"));
}
#[test]
fn parses_unschedulable_rules_from_pool_advanced() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({

96
apps/aether-gateway/src/handlers/public/ai_public.rs
View File

@@ -16,14 +16,11 @@ const CLAUDE_COUNT_TOKENS_MISSING_BODY_DETAIL: &str = "请求体不能为空";
const GEMINI_VIDEO_TASK_NOT_FOUND_DETAIL: &str = "Video task not found";
const AI_PUBLIC_METHOD_NOT_ALLOWED_DETAIL: &str = "Method not allowed";
const AI_PUBLIC_UNAUTHORIZED_DETAIL: &str = "Unauthorized";
const OPENAI_CHAT_IMAGE_MODEL_DETAIL: &str =
"图片模型仅支持通过 /v1/images/generations、/v1/images/edits 或 /v1/images/variations 调用";
const OPENAI_IMAGE_PROMPT_DETAIL: &str = "图片生成/编辑请求缺少 prompt";
const OPENAI_IMAGE_EDIT_INPUT_DETAIL: &str = "图片编辑请求至少需要 1 张输入图片";
const OPENAI_IMAGE_VARIATION_INPUT_DETAIL: &str = "图片变体请求需要 image 文件";
const OPENAI_IMAGE_N_DETAIL: &str = "当前 Codex 图片反代仅支持 n=1";
const OPENAI_IMAGE_STREAM_VARIATION_DETAIL: &str = "图片变体接口当前仅支持同步响应";
const OPENAI_IMAGE_STREAM_MODEL_DETAIL: &str = "stream/partial_images 仅支持 GPT Image 系列模型";
const OPENAI_IMAGE_PARTIAL_IMAGES_DETAIL: &str =
"partial_images 仅支持 0-3，且必须配合 stream=true";
const OPENAI_IMAGE_STYLE_DETAIL: &str = "当前 Codex 图片反代暂不支持 style 参数";
@@ -132,14 +129,6 @@ fn maybe_build_local_openai_request_validation_response(
if decision.route_kind.as_deref() == Some("chat")
&& request_context.request_path == "/v1/chat/completions"
{
let payload = serde_json::from_slice::<Value>(request_body).ok()?;
let model = payload.get("model").and_then(Value::as_str)?;
if is_openai_image_model(model) {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_CHAT_IMAGE_MODEL_DETAIL,
));
}
return None;
}
@@ -169,20 +158,6 @@ fn maybe_build_local_openai_request_validation_response(
}
};
if validation
.model
.as_deref()
.is_some_and(|model| !image_model_supported_for_operation(operation, model))
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
format!(
"该接口不支持模型 {}",
validation.model.as_deref().unwrap_or_default()
),
));
}
match operation {
OpenAiImageOperation::Generate | OpenAiImageOperation::Edit
if validation.prompt.is_none() =>
@@ -237,12 +212,6 @@ fn maybe_build_local_openai_request_validation_response(
OPENAI_IMAGE_STREAM_VARIATION_DETAIL,
));
}
if !image_model_supports_streaming(validation.model.as_deref()) {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_STREAM_MODEL_DETAIL,
));
}
}
if validation
@@ -335,35 +304,6 @@ fn image_request_count(value: &Value) -> Option<u64> {
})
}
fn is_openai_image_model(model: &str) -> bool {
canonicalize_openai_image_model(model).is_some()
}
fn canonicalize_openai_image_model(model: &str) -> Option<&'static str> {
match model.trim().to_ascii_lowercase().as_str() {
"gpt-image-1" => Some("gpt-image-1"),
"gpt-image-1.5" => Some("gpt-image-1.5"),
"gpt-image-1-mini" => Some("gpt-image-1-mini"),
"gpt-image-2" => Some("gpt-image-2"),
"chatgpt-image-latest" => Some("chatgpt-image-latest"),
"dall-e-2" => Some("dall-e-2"),
"dall-e-3" => Some("dall-e-3"),
_ => None,
}
}
fn image_model_supported_for_operation(operation: OpenAiImageOperation, model: &str) -> bool {
match operation {
OpenAiImageOperation::Generate => true,
OpenAiImageOperation::Edit => !matches!(model, "dall-e-3"),
OpenAiImageOperation::Variation => model == "dall-e-2",
}
}
fn image_model_supports_streaming(model: Option<&str>) -> bool {
!matches!(model, Some("dall-e-2" | "dall-e-3"))
}
fn parse_openai_image_validation_input(
operation: OpenAiImageOperation,
content_type: Option<&str>,
@@ -398,8 +338,7 @@ fn parse_openai_image_validation_input_from_json(
Ok(OpenAiImageValidationInput {
model: normalize_openai_image_model_for_operation(
object.get("model").and_then(Value::as_str),
)
.ok_or(OPENAI_IMAGE_INVALID_JSON_DETAIL)?,
),
prompt: object
.get("prompt")
.and_then(Value::as_str)
@@ -480,8 +419,7 @@ fn parse_openai_image_validation_input_from_multipart(
.map(|field| String::from_utf8_lossy(&field.data).trim().to_string());
Ok(OpenAiImageValidationInput {
model: normalize_openai_image_model_for_operation(model.as_deref())
.ok_or(OPENAI_IMAGE_INVALID_MULTIPART_DETAIL)?,
model: normalize_openai_image_model_for_operation(model.as_deref()),
prompt: multipart_text_field(&fields, "prompt"),
image_count: fields
.iter()
@@ -515,11 +453,11 @@ fn parse_openai_image_validation_input_from_multipart(
})
}
fn normalize_openai_image_model_for_operation(model: Option<&str>) -> Option<Option<String>> {
let Some(model) = model.map(str::trim).filter(|value| !value.is_empty()) else {
return Some(None);
};
canonicalize_openai_image_model(model).map(|canonical| Some(canonical.to_string()))
fn normalize_openai_image_model_for_operation(model: Option<&str>) -> Option<String> {
model
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn count_json_images(object: &serde_json::Map<String, Value>) -> usize {
@@ -1090,7 +1028,10 @@ fn estimate_text_tokens(text: &str) -> u64 {
#[cfg(test)]
mod tests {
use super::estimate_claude_count_tokens;
use super::{
estimate_claude_count_tokens, parse_openai_image_validation_input, OpenAiImageOperation,
};
use axum::body::Bytes;
use serde_json::json;
#[test]
@@ -1125,4 +1066,19 @@ mod tests {
assert_eq!(estimate_claude_count_tokens(&payload), Err(()));
}
#[test]
fn image_validation_accepts_custom_model_name() {
let body =
Bytes::from_static(br#"{"model":" Custom/Image-Model:V1 ","prompt":"draw an image"}"#);
let validation = parse_openai_image_validation_input(
OpenAiImageOperation::Generate,
Some("application/json"),
&body,
)
.expect("custom image model should validate");
assert_eq!(validation.model.as_deref(), Some("Custom/Image-Model:V1"));
}
}

72
apps/aether-gateway/src/orchestration/classifier.rs
View File

@@ -82,6 +82,30 @@ const RETRYABLE_RATE_LIMIT_PATTERNS: &[&str] = &[
"quota hit",
];
const RETRYABLE_ACCOUNT_OR_BILLING_PATTERNS: &[&str] = &[
"organization has been disabled",
"organization_disabled",
"account has been disabled",
"account_disabled",
"account has been deactivated",
"account_deactivated",
"account deactivated",
"account suspended",
"account banned",
"subscription inactive",
"payment_required",
"payment required",
"insufficient_quota",
"insufficient quota",
"quota exhausted",
"credits exhausted",
"credit balance",
"credit limit",
"verify your account",
"account verification",
"verification required",
];
#[derive(Debug, Clone, PartialEq, Eq, Default)]
struct ParsedLocalErrorResponse {
type_name: Option<String>,
@@ -169,6 +193,10 @@ pub(crate) fn classify_local_failover(
return LocalFailoverClassification::RetrySemanticRateLimit;
}
if is_semantic_account_or_billing_error(input.status_code, &parsed_error) {
return LocalFailoverClassification::RetryUpstreamFailure;
}
if is_semantic_client_error(input.status_code, &parsed_error) {
return LocalFailoverClassification::StopSemanticClientError;
}
@@ -379,6 +407,21 @@ fn is_semantic_rate_limit_error(status_code: u16, parsed: &ParsedLocalErrorRespo
.any(|pattern| search_text.contains(&pattern.to_ascii_lowercase()))
}
fn is_semantic_account_or_billing_error(
status_code: u16,
parsed: &ParsedLocalErrorResponse,
) -> bool {
if status_code < 400 {
return false;
}
let search_text = semantic_search_text(parsed);
!search_text.is_empty()
&& RETRYABLE_ACCOUNT_OR_BILLING_PATTERNS
.iter()
.any(|pattern| search_text.contains(&pattern.to_ascii_lowercase()))
}
fn local_failover_regex_rule_matches(
rule: &LocalFailoverRegexRule,
response_text: &str,
@@ -539,6 +582,35 @@ mod tests {
);
}
#[test]
fn classifier_retries_account_and_billing_errors_before_client_error_stop() {
assert_eq!(
classify_local_failover(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(
403,
Some(
"{\"error\":{\"type\":\"invalid_request_error\",\"message\":\"verify your account before continuing\"}}"
)
)
),
LocalFailoverClassification::RetryUpstreamFailure
);
assert_eq!(
classify_local_failover(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(
402,
Some(
"{\"error\":{\"type\":\"invalid_request_error\",\"message\":\"payment required: credit balance exhausted\"}}"
)
)
),
LocalFailoverClassification::RetryUpstreamFailure
);
}
#[test]
fn classifier_keeps_embedded_rate_limit_error_in_success_response_on_default_path() {
assert_eq!(

314
apps/aether-gateway/src/tests/ai_execute/stream/image.rs
View File

@@ -28,6 +28,7 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth
authorization: String,
x_client_request_id: String,
tool_type: String,
tool_action: String,
tool_partial_images: Option<u64>,
request_stream: bool,
plan_stream: bool,
@@ -267,6 +268,15 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
tool_action: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("tools"))
.and_then(|value| value.get(0))
.and_then(|value| value.get("action"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
tool_partial_images: payload
.get("body")
.and_then(|value| value.get("json_body"))
@@ -417,6 +427,7 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth
"trace-codex-image-stream-local-123"
);
assert_eq!(seen_execution_runtime_request.tool_type, "image_generation");
assert_eq!(seen_execution_runtime_request.tool_action, "generate");
assert_eq!(seen_execution_runtime_request.tool_partial_images, Some(1));
assert!(seen_execution_runtime_request.request_stream);
assert!(seen_execution_runtime_request.plan_stream);
@@ -425,3 +436,306 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth
execution_runtime_handle.abort();
refresh_handle.abort();
}
#[tokio::test]
async fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeStreamRequest {
trace_id: String,
request_stream: bool,
plan_stream: bool,
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai", "codex"])),
Some(serde_json::json!(["openai:image"])),
Some(serde_json::json!(["gpt-image-2"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800_i64),
Some(serde_json::json!(["openai", "codex"])),
Some(serde_json::json!(["openai:image"])),
Some(serde_json::json!(["gpt-image-2"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-codex-image-stream-local-1".to_string(),
provider_name: "codex".to_string(),
provider_type: "codex".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-codex-image-stream-local-1".to_string(),
endpoint_api_format: "openai:image".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("image".to_string()),
endpoint_is_active: true,
key_id: "key-codex-image-stream-local-1".to_string(),
key_name: "oauth".to_string(),
key_auth_type: "oauth".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:image".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})),
model_id: "model-codex-image-stream-local-1".to_string(),
global_model_id: "global-model-codex-image-stream-local-1".to_string(),
global_model_name: "gpt-image-2".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-image-2".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-image-2".to_string(),
priority: 1,
api_formats: Some(vec!["openai:image".to_string()]),
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-codex-image-stream-local-1".to_string(),
"codex".to_string(),
Some("https://chatgpt.com".to_string()),
"codex".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-codex-image-stream-local-1".to_string(),
"provider-codex-image-stream-local-1".to_string(),
"openai:image".to_string(),
Some("openai".to_string()),
Some("image".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://chatgpt.com/backend-api/codex".to_string(),
None,
None,
Some(2),
None,
Some(serde_json::json!({"upstream_stream_policy":"force_stream"})),
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"rt-codex-image-stream-local-123"}"#,
)
.expect("auth config should encrypt");
StoredProviderCatalogKey::new(
"key-codex-image-stream-local-1".to_string(),
"provider-codex-image-stream-local-1".to_string(),
"oauth".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:image"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
.expect("placeholder api key should encrypt"),
Some(encrypted_auth_config),
None,
Some(serde_json::json!({"openai:image": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeStreamRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let refresh = Router::new().route(
"/oauth/token",
any(move |_request: Request| async move {
Json(json!({
"access_token": "refreshed-codex-image-stream-access-token",
"refresh_token": "rt-codex-image-stream-local-456",
"token_type": "Bearer",
"expires_in": 3600
}))
}),
);
let execution_runtime = Router::new().route(
"/v1/execute/stream",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
.expect("execution runtime payload should parse");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeStreamRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
request_stream: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("stream"))
.and_then(|value| value.as_bool())
.unwrap_or(false),
plan_stream: payload
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(false),
});
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/json\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"created\\\":1776991097,\\\"data\\\":[{\\\"b64_json\\\":\\\"aGVsbG8=\\\"}],\\\"usage\\\":{\\\"total_tokens\\\":100,\\\"input_tokens\\\":50,\\\"output_tokens\\\":50,\\\"input_tokens_details\\\":{\\\"text_tokens\\\":10,\\\"image_tokens\\\":40}}}\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = http::Response::builder()
.status(StatusCode::OK)
.body(Body::from(frames))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
http::HeaderValue::from_static("application/x-ndjson"),
);
response
}
}),
);
let client_api_key = "sk-client-codex-image-stream-local";
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(client_api_key)),
sample_auth_snapshot(
"key-codex-image-stream-client-123",
"user-codex-image-stream-client-123",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let (refresh_url, refresh_handle) = start_server(refresh).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", format!("{refresh_url}/oauth/token")),
),
]);
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::new(InMemoryRequestCandidateRepository::default()),
DEVELOPMENT_ENCRYPTION_KEY,
),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/images/generations"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
format!("Bearer {client_api_key}"),
)
.header(TRACE_ID_HEADER, "trace-codex-image-stream-json-123")
.body("{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张海报\",\"stream\":true}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("text/event-stream")
);
let response_text = response.text().await.expect("body should read");
assert!(response_text.contains("event: image_generation.completed"));
assert!(response_text.contains("\"type\":\"image_generation.completed\""));
assert!(response_text.contains("\"b64_json\":\"aGVsbG8=\""));
assert!(response_text.contains("\"total_tokens\":100"));
assert!(!response_text.trim_start().starts_with('{'));
assert!(!response_text.contains("\"created\": 1776991097"));
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime stream should be captured");
assert_eq!(
seen_execution_runtime_request.trace_id,
"trace-codex-image-stream-json-123"
);
assert!(seen_execution_runtime_request.request_stream);
assert!(seen_execution_runtime_request.plan_stream);
gateway_handle.abort();
execution_runtime_handle.abort();
refresh_handle.abort();
}

13
apps/aether-gateway/src/tests/frontdoor/ai.rs
View File

@@ -7,7 +7,7 @@ use super::{
use crate::tests::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, AppState, Arc, Body, Json, Mutex, Request, Router, StatusCode, EXECUTION_PATH_HEADER,
EXECUTION_PATH_LOCAL_AI_PUBLIC,
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
};
use axum::response::IntoResponse;
@@ -522,7 +522,7 @@ async fn gateway_rejects_invalid_claude_count_tokens_payload_without_hitting_fal
}
#[tokio::test]
async fn gateway_rejects_gpt_image_2_on_chat_completions_without_hitting_fallback_probe() {
async fn gateway_does_not_locally_reject_image_model_name_on_chat_completions() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
@@ -567,18 +567,13 @@ async fn gateway_rejects_gpt_image_2_on_chat_completions_without_hitting_fallbac
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(
payload["detail"],
"图片模型仅支持通过 /v1/images/generations、/v1/images/edits 或 /v1/images/variations 调用"
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);