mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
fix(gateway): route openai image streams with images surface
This commit is contained in:
@@ -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) {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user