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"));
}
}