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

@@ -172,7 +172,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
} else if is_grok {
build_grok_upstream_url(transport, GROK_CHAT_PATH)
} else {
build_openai_image_upstream_url(transport, parts.uri.query())
build_openai_image_upstream_url(transport, Some(parts.uri.path()), parts.uri.query())
};
let mut provider_request_body = if is_chatgpt_web {
match build_chatgpt_web_image_request_body(parts, body_json, body_base64) {

View File

@@ -625,7 +625,8 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
};
let upstream_is_stream = true;
let upstream_url = build_openai_image_upstream_url(transport, None);
let upstream_url =
build_openai_image_upstream_url(transport, Some("/v1/images/generations"), None);
let effective_headers = input.effective_headers(&parts.headers);
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {

View File

@@ -739,7 +739,11 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
let upstream_url = if is_chatgpt_web {
chatgpt_web_image_internal_url(&transport.endpoint.base_url)
} else {
build_openai_image_upstream_url(transport, parts.uri.query())
build_openai_image_upstream_url(
transport,
Some("/v1/images/generations"),
parts.uri.query(),
)
};
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {

View File

@@ -1315,9 +1315,29 @@ fn build_sse_body_stream(
}
fn stream_chunk_contains_sse_done(chunk: &[u8]) -> bool {
std::str::from_utf8(chunk)
.ok()
.is_some_and(|text| text.lines().any(|line| line.trim() == "data: [DONE]"))
std::str::from_utf8(chunk).ok().is_some_and(|text| {
text.lines().any(|line| {
let line = line.trim();
if matches!(
line,
"data: [DONE]" | "event: message_stop" | "event: response.completed"
) {
return true;
}
let Some(data) = line.strip_prefix("data:").map(str::trim) else {
return false;
};
data == "[DONE]"
|| serde_json::from_str::<serde_json::Value>(data).is_ok_and(|value| {
value
.get("type")
.and_then(serde_json::Value::as_str)
.is_some_and(|event_type| {
matches!(event_type, "message_stop" | "response.completed")
})
})
})
})
}
async fn next_stream_frame<R>(
@@ -3364,7 +3384,7 @@ mod tests {
build_sse_body_stream, execute_execution_runtime_stream, execute_stream_from_frame_stream,
maybe_apply_kiro_prompt_cache_usage_to_stream_summary, merge_stream_terminal_summary,
should_limit_direct_finalize_prefetch, should_probe_success_failover_before_stream,
should_skip_direct_finalize_prefetch,
should_skip_direct_finalize_prefetch, stream_chunk_contains_sse_done,
};
use crate::control::GatewayControlDecision;
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
@@ -3381,6 +3401,20 @@ mod tests {
.with_execution_runtime_candidate(true)
}
#[test]
fn detects_client_visible_sse_terminal_events() {
assert!(stream_chunk_contains_sse_done(b"data: [DONE]\n\n"));
assert!(stream_chunk_contains_sse_done(
b"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"
));
assert!(stream_chunk_contains_sse_done(
b"event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{}}\n\n"
));
assert!(!stream_chunk_contains_sse_done(
b"event: content_block_delta\ndata: {\"type\":\"content_block_delta\"}\n\n"
));
}
fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {
aether_contracts::ProxySnapshot {
enabled: Some(true),

View File

@@ -1798,6 +1798,7 @@ fn provider_query_chatgpt_web_image_internal_url(base_url: &str) -> String {
fn provider_query_openai_image_test_upstream_url(
transport: &AdminGatewayProviderTransportSnapshot,
request_path: Option<&str>,
request_query: Option<&str>,
) -> String {
if transport
@@ -1818,7 +1819,11 @@ fn provider_query_openai_image_test_upstream_url(
crate::provider_transport::GROK_CHAT_PATH,
)
} else {
crate::provider_transport::build_openai_image_upstream_url(transport, request_query)
crate::provider_transport::build_openai_image_upstream_url(
transport,
request_path,
request_query,
)
}
}
@@ -2062,7 +2067,11 @@ async fn provider_query_execute_openai_image_test_candidate(
} else {
normalized_request.summary_json.clone()
};
let request_url = provider_query_openai_image_test_upstream_url(&transport, parts.uri.query());
let request_url = provider_query_openai_image_test_upstream_url(
&transport,
Some(parts.uri.path()),
parts.uri.query(),
);
let upstream_is_stream = provider_request_body
.get("stream")
.and_then(Value::as_bool)

View File

@@ -739,17 +739,39 @@ fn provider_query_grok_image_test_uses_grok_app_chat_upstream_url() {
let transport = sample_openai_image_transport("grok");
assert_eq!(
provider_query_openai_image_test_upstream_url(&transport, Some("trace=1")),
provider_query_openai_image_test_upstream_url(
&transport,
Some("/v1/images/generations"),
Some("trace=1"),
),
"https://grok.com/rest/app-chat/conversations/new"
);
}
#[test]
fn provider_query_custom_image_test_uses_images_upstream_url() {
let transport = sample_openai_image_transport("custom");
assert_eq!(
provider_query_openai_image_test_upstream_url(
&transport,
Some("/v1/images/generations"),
Some("trace=1"),
),
"https://grok.com/v1/images/generations?trace=1"
);
}
#[test]
fn provider_query_chatgpt_web_image_test_uses_internal_upstream_url() {
let transport = sample_openai_image_transport("chatgpt_web");
assert_eq!(
provider_query_openai_image_test_upstream_url(&transport, Some("trace=1")),
provider_query_openai_image_test_upstream_url(
&transport,
Some("/v1/images/generations"),
Some("trace=1"),
),
"https://grok.com/__aether/chatgpt-web-image"
);
}