i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-12 04:09:48 +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:07:24 +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
apps/aether-gateway/src/ai_pipeline/finalize/internal/sync_to_stream.rs
+205
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"));
}
}
apps/aether-gateway/src/ai_pipeline/planner/pool_scheduler.rs
+122 -1
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(
apps/aether-gateway/src/ai_pipeline/planner/specialized/image/request.rs
+50 -45
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")