fix(gateway): route openai image streams with images surface

This commit is contained in:
ZheFox
2026-05-20 18:13:34 +08:00
parent de7be4f15b
commit 2cacf66a37
9 changed files with 425 additions and 17 deletions

View File

@@ -16,7 +16,9 @@ use crate::formats::shared::sse::{encode_done_sse, encode_json_sse};
use crate::formats::shared::stream_core::common::{
build_openai_chat_chunk, build_openai_chat_finish_chunk, build_openai_chat_usage_chunk,
};
use crate::formats::shared::stream_core::CanonicalStreamFrame;
use crate::formats::shared::stream_core::{
CanonicalStreamFrame, StreamingStandardFormatMatrix, StreamingStandardTerminalObserver,
};
use crate::formats::shared::AiSurfaceFinalizeError;
pub struct SyncToStreamBridgeOutcome {
@@ -32,6 +34,14 @@ pub fn maybe_bridge_standard_sync_json_to_stream(
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
let provider_api_format = normalize_api_format(provider_api_format);
let client_api_format = normalize_api_format(client_api_format);
if let Some(outcome) = maybe_bridge_aether_sse_response_capture_to_stream(
provider_body_json,
provider_api_format.as_str(),
client_api_format.as_str(),
report_context,
)? {
return Ok(Some(outcome));
}
if provider_api_format == "openai:image" {
return match client_api_format.as_str() {
"openai:image" => {
@@ -586,6 +596,173 @@ fn is_standard_api_format(value: &str) -> bool {
)
}
fn maybe_bridge_aether_sse_response_capture_to_stream(
provider_body_json: &Value,
provider_api_format: &str,
client_api_format: &str,
report_context: Option<&Value>,
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
let Some(object) = provider_body_json.as_object() else {
return Ok(None);
};
let status_code = object
.get("status_code")
.or_else(|| object.get("statusCode"))
.and_then(Value::as_u64)
.unwrap_or(0);
if !(200..300).contains(&status_code) {
return Ok(None);
}
let Some(headers) = object.get("headers").and_then(Value::as_object) else {
return Ok(None);
};
let content_type = response_capture_header(headers, "content-type").unwrap_or_default();
if !content_type
.to_ascii_lowercase()
.contains("text/event-stream")
{
return Ok(None);
}
let Some(body_text) = object.get("body").and_then(Value::as_str) else {
return Ok(None);
};
if body_text.trim().is_empty() || body_text.contains("...[truncated]") {
return Ok(None);
}
let captured_api_format =
response_capture_header(headers, "x-aether-control-endpoint-signature")
.map(normalize_api_format)
.or_else(|| infer_sse_body_api_format(body_text))
.unwrap_or_else(|| provider_api_format.to_string());
if !is_standard_api_format(captured_api_format.as_str())
|| !is_standard_api_format(client_api_format)
{
return Ok(None);
}
let bridge_context = build_bridge_report_context(
report_context,
captured_api_format.as_str(),
client_api_format,
);
let sse_body = if captured_api_format == client_api_format {
body_text.as_bytes().to_vec()
} else {
rewrite_sse_body_between_formats(
body_text.as_bytes(),
captured_api_format.as_str(),
client_api_format,
&bridge_context,
)?
};
let terminal_summary = observe_sse_terminal_summary(
body_text.as_bytes(),
captured_api_format.as_str(),
&bridge_context,
)?;
Ok(Some(SyncToStreamBridgeOutcome {
sse_body,
terminal_summary,
}))
}
fn response_capture_header<'a>(headers: &'a Map<String, Value>, name: &str) -> Option<&'a str> {
headers
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case(name))
.and_then(|(_, value)| value.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn infer_sse_body_api_format(body_text: &str) -> Option<String> {
if body_text.contains("event: message_start")
|| body_text.contains("\"type\":\"message_start\"")
{
return Some("claude:messages".to_string());
}
if body_text.contains("event: response.") || body_text.contains("\"type\":\"response.") {
return Some("openai:responses".to_string());
}
if body_text.contains("data: [DONE]")
|| body_text.contains("\"object\":\"chat.completion.chunk\"")
{
return Some("openai:chat".to_string());
}
if body_text.contains("\"candidates\"") && body_text.contains("\"finishReason\"") {
return Some("gemini:generate_content".to_string());
}
None
}
fn rewrite_sse_body_between_formats(
body: &[u8],
provider_api_format: &str,
client_api_format: &str,
report_context: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut context =
build_bridge_report_context(Some(report_context), provider_api_format, client_api_format);
if let Some(object) = context.as_object_mut() {
object.insert(
"provider_stream_event_api_format".to_string(),
Value::String(provider_api_format.to_string()),
);
}
let mut matrix = StreamingStandardFormatMatrix::default();
let mut out = Vec::new();
for_each_sse_line(body, |line| {
out.extend(matrix.transform_line(&context, line)?);
Ok(())
})?;
out.extend(matrix.finish(&context)?);
Ok(out)
}
fn observe_sse_terminal_summary(
body: &[u8],
provider_api_format: &str,
report_context: &Value,
) -> Result<Option<ExecutionStreamTerminalSummary>, AiSurfaceFinalizeError> {
let mut context = build_bridge_report_context(
Some(report_context),
provider_api_format,
provider_api_format,
);
if let Some(object) = context.as_object_mut() {
object.insert(
"provider_stream_event_api_format".to_string(),
Value::String(provider_api_format.to_string()),
);
}
let mut observer = StreamingStandardTerminalObserver::default();
for_each_sse_line(body, |line| observer.push_line(&context, line))?;
observer.finish(&context)
}
fn for_each_sse_line<F>(body: &[u8], mut on_line: F) -> Result<(), AiSurfaceFinalizeError>
where
F: FnMut(Vec<u8>) -> Result<(), AiSurfaceFinalizeError>,
{
let mut start = 0usize;
for (index, byte) in body.iter().enumerate() {
if *byte == b'\n' {
on_line(body[start..=index].to_vec())?;
start = index + 1;
}
}
if start < body.len() {
on_line(body[start..].to_vec())?;
}
Ok(())
}
fn extract_openai_image_sync_b64_json(item: &serde_json::Map<String, Value>) -> Option<String> {
item.get("b64_json")
.and_then(Value::as_str)
@@ -1101,4 +1278,85 @@ mod tests {
assert!(output.contains("\"b64_json\":\"d29ybGQ=\""));
assert!(output.contains("\"total_tokens\":9"));
}
#[test]
fn bridges_aether_sse_response_capture_to_same_client_stream() {
let captured_body = concat!(
": aether-keepalive\n\n",
"event: message_start\n",
"data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"gpt-5.5\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":0,\"output_tokens\":0}}}\n\n",
"event: content_block_start\n",
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
"event: content_block_delta\n",
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n",
"event: message_delta\n",
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"input_tokens\":1,\"output_tokens\":2}}\n\n",
"event: message_stop\n",
"data: {\"type\":\"message_stop\"}\n\n",
);
let outcome = maybe_bridge_standard_sync_json_to_stream(
&json!({
"status_code": 200,
"headers": {
"content-type": "text/event-stream",
"x-aether-control-endpoint-signature": "claude:messages"
},
"body": captured_body
}),
"openai:responses",
"claude:messages",
None,
)
.expect("bridge should succeed")
.expect("capture should bridge");
let output = utf8(outcome.sse_body);
assert!(output.contains("event: message_start"));
assert!(output.contains("event: message_stop"));
assert!(!output.contains("status_code"));
assert_eq!(
outcome
.terminal_summary
.as_ref()
.and_then(|summary| summary.finish_reason.as_deref()),
Some("end_turn")
);
}
#[test]
fn rewrites_aether_sse_response_capture_to_requested_client_stream() {
let captured_body = concat!(
"event: message_start\n",
"data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"gpt-5.5\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":0,\"output_tokens\":0}}}\n\n",
"event: content_block_start\n",
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_1\",\"name\":\"Edit\",\"input\":{}}}\n\n",
"event: content_block_delta\n",
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"file_path\\\":\\\"/tmp/a.txt\\\"}\"}}\n\n",
"event: message_delta\n",
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"input_tokens\":1,\"output_tokens\":2}}\n\n",
"event: message_stop\n",
"data: {\"type\":\"message_stop\"}\n\n",
);
let outcome = maybe_bridge_standard_sync_json_to_stream(
&json!({
"status_code": 200,
"headers": {
"content-type": "text/event-stream",
"x-aether-control-endpoint-signature": "claude:messages"
},
"body": captured_body
}),
"claude:messages",
"openai:responses",
None,
)
.expect("bridge should succeed")
.expect("capture should bridge");
let output = utf8(outcome.sse_body);
assert!(output.contains("event: response.output_item.added"));
assert!(output.contains("event: response.function_call_arguments.delta"));
assert!(output.contains("event: response.completed"));
assert!(!output.contains("status_code"));
}
}

View File

@@ -7,7 +7,7 @@ use crate::grok::{is_grok_provider_transport, resolve_grok_session_auth};
use crate::policy::local_standard_transport_unsupported_reason_with_network;
use crate::rules::apply_local_header_rules_with_request_headers;
use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::url::build_openai_responses_url;
use crate::url::{build_openai_image_url, build_openai_responses_url};
#[derive(Debug, Clone, Copy)]
pub struct ProviderOpenAiImageHeadersInput<'a> {
@@ -56,9 +56,18 @@ pub fn resolve_openai_image_auth(
pub fn build_openai_image_upstream_url(
transport: &GatewayProviderTransportSnapshot,
request_path: Option<&str>,
request_query: Option<&str>,
) -> String {
build_openai_responses_url(&transport.endpoint.base_url, request_query, false)
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("codex")
{
return build_openai_responses_url(&transport.endpoint.base_url, request_query, false);
}
build_openai_image_url(&transport.endpoint.base_url, request_path, request_query)
}
pub fn build_openai_image_headers(
@@ -155,12 +164,30 @@ mod tests {
}
#[test]
fn builds_openai_image_url_on_responses_surface() {
let url = build_openai_image_upstream_url(&sample_transport(), Some("trace=1"));
fn codex_openai_image_url_stays_on_responses_surface() {
let url = build_openai_image_upstream_url(
&sample_transport(),
Some("/v1/images/generations"),
Some("trace=1"),
);
assert_eq!(url, "https://api.openai.com/v1/responses?trace=1");
}
#[test]
fn standard_openai_image_url_uses_images_surface() {
let mut transport = sample_transport();
transport.provider.provider_type = "openai".to_string();
let url = build_openai_image_upstream_url(
&transport,
Some("/v1/images/generations"),
Some("trace=1"),
);
assert_eq!(url, "https://api.openai.com/v1/images/generations?trace=1");
}
#[test]
fn chatgpt_web_is_supported_by_dedicated_openai_image_transport_policy() {
let mut transport = sample_transport();

View File

@@ -41,6 +41,43 @@ pub fn build_openai_responses_url(
url
}
pub fn build_openai_image_url(
upstream_base_url: &str,
request_path: Option<&str>,
query: Option<&str>,
) -> String {
let (trimmed, base_query) = split_base_url_query(upstream_base_url);
let trimmed = trimmed.trim_end_matches('/');
let suffix = openai_image_path_suffix(request_path);
let mut url = if openai_image_base_includes_operation_path(trimmed) {
trimmed.to_string()
} else if trimmed.ends_with("/v1") || google_openai_compat_base_includes_api_root(trimmed) {
format!("{trimmed}{suffix}")
} else {
format!("{trimmed}/v1{suffix}")
};
append_merged_query(&mut url, base_query, None, query, &[]);
url
}
fn openai_image_path_suffix(request_path: Option<&str>) -> &'static str {
match request_path
.map(str::trim)
.map(|value| value.trim_end_matches('/'))
{
Some("/v1/images/edits") | Some("/images/edits") => "/images/edits",
_ => "/images/generations",
}
}
fn openai_image_base_includes_operation_path(base_url: &str) -> bool {
let path = Url::parse(base_url)
.ok()
.map(|url| url.path().trim_end_matches('/').to_string())
.unwrap_or_else(|| base_url.trim_end_matches('/').to_string());
path.ends_with("/images/generations") || path.ends_with("/images/edits")
}
pub fn build_claude_messages_url(upstream_base_url: &str, query: Option<&str>) -> String {
let (trimmed, base_query) = split_base_url_query(upstream_base_url);
let trimmed = trimmed.trim_end_matches('/');
@@ -307,7 +344,7 @@ fn merge_query_string(
mod tests {
use super::{
build_gemini_content_url, build_gemini_files_passthrough_url,
build_gemini_video_predict_long_running_url, build_openai_chat_url,
build_gemini_video_predict_long_running_url, build_openai_chat_url, build_openai_image_url,
build_openai_responses_url, build_passthrough_path_url,
normalize_gemini_content_action_path,
};
@@ -353,6 +390,22 @@ mod tests {
);
}
#[test]
fn openai_image_url_uses_images_surface() {
assert_eq!(
build_openai_image_url(
"https://api.openai.example/v1?tenant=demo",
Some("/v1/images/generations"),
Some("trace=1")
),
"https://api.openai.example/v1/images/generations?tenant=demo&trace=1"
);
assert_eq!(
build_openai_image_url("https://api.openai.example", Some("/v1/images/edits"), None),
"https://api.openai.example/v1/images/edits"
);
}
#[test]
fn merges_base_url_query_for_dynamic_gemini_content_urls() {
assert_eq!(