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

fix(gateway): support openai image accept negotiation

Browse Source
This commit is contained in:
zhefox
2026-06-08 20:40:30 +08:00
parent 82040bfc21
commit 04ba8cbe9e
8 changed files with 273 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/planner/specialized/image/request.rs
+5
View File
@@ -213,6 +213,11 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
headers: effective_headers,
auth_header: &auth_header,
auth_value: &auth_value,
accept: if is_codex || is_chatgpt_web {
"text/event-stream"
} else {
"application/json"
},
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs
+1
View File
@@ -1225,6 +1225,7 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
headers: effective_headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
accept: "text/event-stream",
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &converted.body_json,
original_request_body: body_json,
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs
+1
View File
@@ -1051,6 +1051,7 @@ async fn resolve_openai_chat_to_openai_image_payload_parts(
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
accept: "text/event-stream",
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs
+1
View File
@@ -1118,6 +1118,7 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
accept: "text/event-stream",
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs
+5
View File
@@ -2183,6 +2183,11 @@ async fn provider_query_execute_openai_image_test_candidate(
headers: &parts.headers,
auth_header: &auth_header,
auth_value: &auth_value,
accept: if is_codex || is_chatgpt_web {
"text/event-stream"
} else {
"application/json"
},
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: &request_body,
crates/aether-provider-transport/src/openai_image/mod.rs
+44 -1
View File
@@ -14,6 +14,7 @@ pub struct ProviderOpenAiImageHeadersInput<'a> {
pub headers: &'a http::HeaderMap,
pub auth_header: &'a str,
pub auth_value: &'a str,
pub accept: &'a str,
pub header_rules: Option<&'a Value>,
pub provider_request_body: &'a Value,
pub original_request_body: &'a Value,
@@ -80,7 +81,7 @@ pub fn build_openai_image_headers(
&BTreeMap::new(),
);
provider_request_headers.insert("content-type".to_string(), "application/json".to_string());
provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string());
provider_request_headers.insert("accept".to_string(), input.accept.to_string());
if !apply_local_header_rules_with_request_headers(
&mut provider_request_headers,
input.header_rules,
@@ -256,6 +257,7 @@ mod tests {
headers: &HeaderMap::new(),
auth_header: "authorization",
auth_value: "Bearer secret",
accept: "text/event-stream",
header_rules: Some(&json!([
{"action":"set","key":"x-image-route","value":"codex"}
])),
@@ -278,4 +280,45 @@ mod tests {
);
assert_eq!(headers.get("x-image-route"), Some(&"codex".to_string()));
}
#[test]
fn standard_openai_compatible_image_headers_can_request_json() {
let headers = build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &HeaderMap::new(),
auth_header: "authorization",
auth_value: "Bearer secret",
accept: "application/json",
header_rules: None,
provider_request_body: &json!({
"model": "upstream-image-model",
"prompt": "draw a city",
}),
original_request_body: &json!({"prompt":"draw a city"}),
})
.expect("headers should build");
assert_eq!(headers.get("accept"), Some(&"application/json".to_string()));
assert_eq!(
headers.get("content-type"),
Some(&"application/json".to_string())
);
}
#[test]
fn standard_openai_compatible_image_url_supports_aether_api_root() {
let mut transport = sample_transport();
transport.provider.provider_type = "custom".to_string();
transport.endpoint.base_url = "https://upstream-aether.example/v1".to_string();
let url = build_openai_image_upstream_url(
&transport,
Some("/v1/images/generations"),
Some("trace=1"),
);
assert_eq!(
url,
"https://upstream-aether.example/v1/images/generations?trace=1"
);
}
}
crates/aether-usage-runtime/src/body_capture.rs
+111 -9
View File
@@ -10,6 +10,7 @@ use crate::event::UsageEvent;
use crate::runtime::{UsageBodyCapturePolicy, UsageRequestRecordLevel};
const TRUNCATED_BODY_STRING_SUFFIX: &str = "...[truncated]";
const LARGE_JSON_STRING_CAPTURE_PREFIX_BYTES: usize = 256;
#[derive(Debug)]
struct LimitedUsageBodyCapture {
@@ -336,13 +337,18 @@ fn limit_usage_body_capture_value(
let truncated_value = match value {
Value::String(text) => Value::String(truncate_usage_body_string(&text, limit)),
other => json!({
"truncated": true,
"reason": "body_capture_limit_exceeded",
"max_bytes": limit,
"source_bytes": source_len,
"value_kind": usage_value_kind(&other),
}),
other => {
let value_kind = usage_value_kind(&other);
compact_usage_body_json_value_for_limit(other, limit).unwrap_or_else(|| {
json!({
"truncated": true,
"reason": "body_capture_limit_exceeded",
"max_bytes": limit,
"source_bytes": source_len,
"value_kind": value_kind,
})
})
}
};
let stored_bytes = json_serialized_len(&truncated_value);
LimitedUsageBodyCapture {
@@ -354,6 +360,51 @@ fn limit_usage_body_capture_value(
}
}
fn compact_usage_body_json_value_for_limit(value: Value, max_bytes: usize) -> Option<Value> {
let mut changed = false;
let compacted = compact_large_usage_body_strings(value, &mut changed);
if !changed {
return None;
}
json_serialized_len(&compacted)
.is_some_and(|bytes| bytes <= max_bytes as u64)
.then_some(compacted)
}
fn compact_large_usage_body_strings(value: Value, changed: &mut bool) -> Value {
match value {
Value::String(text) => compact_large_usage_body_string(text, changed),
Value::Array(values) => Value::Array(
values
.into_iter()
.map(|value| compact_large_usage_body_strings(value, changed))
.collect(),
),
Value::Object(object) => Value::Object(
object
.into_iter()
.map(|(key, value)| (key, compact_large_usage_body_strings(value, changed)))
.collect(),
),
other => other,
}
}
fn compact_large_usage_body_string(text: String, changed: &mut bool) -> Value {
let Some(source_bytes) = json_serialized_len(&text) else {
return Value::String(text);
};
if source_bytes <= LARGE_JSON_STRING_CAPTURE_PREFIX_BYTES as u64 {
return Value::String(text);
}
*changed = true;
Value::String(truncate_usage_body_string(
&text,
LARGE_JSON_STRING_CAPTURE_PREFIX_BYTES,
))
}
fn truncate_usage_body_string(value: &str, max_bytes: usize) -> String {
let mut end = value.len();
while end > 0 {
@@ -712,8 +763,8 @@ fn usage_value_kind(value: &Value) -> &'static str {
#[cfg(test)]
mod tests {
use super::{
build_plan_body_capture_metadata, sync_usage_body_ref_metadata,
trim_owned_non_empty_string, truncate_usage_body_string,
build_plan_body_capture_metadata, limit_usage_body_capture_value,
sync_usage_body_ref_metadata, trim_owned_non_empty_string, truncate_usage_body_string,
upsert_body_capture_metadata_value_entry,
};
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
@@ -826,4 +877,55 @@ mod tests {
.ok()
.is_some_and(|bytes| bytes.len() <= limit));
}
#[test]
fn body_capture_limit_preserves_image_response_shape_when_b64_is_large() {
let response = serde_json::json!({
"created": 1_779_273_523,
"data": [{
"b64_json": "a".repeat(8 * 1024),
"revised_prompt": "draw a small city"
}],
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
});
let limited = limit_usage_body_capture_value(response, Some(2048));
assert!(limited.truncated);
assert!(limited.stored_bytes.is_some_and(|bytes| bytes <= 2048));
assert_eq!(limited.value["created"], 1_779_273_523);
assert_eq!(limited.value["usage"]["total_tokens"], 3);
assert_eq!(
limited.value["data"][0]["revised_prompt"],
"draw a small city"
);
assert!(limited.value["data"][0]["b64_json"]
.as_str()
.is_some_and(|value| value.ends_with("...[truncated]")));
}
#[test]
fn body_capture_limit_preserves_image_request_shape_when_input_image_is_large() {
let request = serde_json::json!({
"model": "upstream-image-model",
"prompt": "edit this image",
"image": "data:image/png;base64,".to_string() + &"a".repeat(8 * 1024),
"size": "1024x1024"
});
let limited = limit_usage_body_capture_value(request, Some(2048));
assert!(limited.truncated);
assert!(limited.stored_bytes.is_some_and(|bytes| bytes <= 2048));
assert_eq!(limited.value["model"], "upstream-image-model");
assert_eq!(limited.value["prompt"], "edit this image");
assert_eq!(limited.value["size"], "1024x1024");
assert!(limited.value["image"]
.as_str()
.is_some_and(|value| value.ends_with("...[truncated]")));
}
}
crates/aether-usage-runtime/src/write.rs
+105
View File
@@ -5593,6 +5593,111 @@ mod tests {
);
}
#[test]
fn openai_image_sync_terminal_usage_captures_request_and_response_bodies() {
let plan = ExecutionPlan {
request_id: "req-openai-image-sync-body-1".to_string(),
candidate_id: Some("cand-openai-image-sync-body-1".to_string()),
provider_name: Some("Upstream Aether".to_string()),
provider_id: "provider-aether-1".to_string(),
endpoint_id: "endpoint-aether-1".to_string(),
key_id: "key-aether-1".to_string(),
method: "POST".to_string(),
url: "https://upstream-aether.example/v1/images/generations".to_string(),
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-image-2-upstream",
"prompt": "Draw a red kite",
"size": "1024x1024",
"n": 1,
"response_format": "b64_json"
})),
stream: false,
client_api_format: "openai:image".to_string(),
provider_api_format: "openai:image".to_string(),
model_name: Some("gpt-image-2".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let payload = GatewaySyncReportRequest {
trace_id: "trace-openai-image-sync-body-1".to_string(),
report_kind: "openai_image_sync_finalize".to_string(),
report_context: Some(json!({
"client_api_format": "openai:image",
"provider_api_format": "openai:image",
"needs_conversion": true,
"original_request_body": {
"model": "gpt-image-2",
"prompt": "Draw a red kite",
"size": "1024x1024",
"response_format": "b64_json"
},
"provider_request_body": {
"model": "gpt-image-2-upstream",
"prompt": "Draw a red kite",
"size": "1024x1024",
"n": 1,
"response_format": "b64_json"
}
})),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body_json: Some(json!({
"created": 1776839946,
"data": [{
"b64_json": "aGVsbG8=",
"revised_prompt": "red kite"
}],
"usage": {
"input_tokens": 11,
"output_tokens": 22,
"total_tokens": 33
}
})),
client_body_json: None,
body_base64: None,
telemetry: None,
};
let event =
build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
assert_eq!(
event.data.request_body,
payload
.report_context
.as_ref()
.and_then(|value| value.get("original_request_body"))
.cloned()
);
assert_eq!(
event.data.provider_request_body,
payload
.report_context
.as_ref()
.and_then(|value| value.get("provider_request_body"))
.cloned()
);
assert_eq!(event.data.response_body, payload.body_json);
assert!(event.data.client_response_body.is_none());
assert_eq!(
event.data.request_body_state,
Some(UsageBodyCaptureState::Inline)
);
assert_eq!(
event.data.provider_request_body_state,
Some(UsageBodyCaptureState::Inline)
);
assert_eq!(
event.data.response_body_state,
Some(UsageBodyCaptureState::Inline)
);
}
#[test]
fn sync_terminal_usage_applies_kiro_simulated_cache_context() {
let plan = ExecutionPlan {