Merge remote-tracking branch 'origin/pr-483' into merge-pr-483

# Conflicts:
#	apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs
#	apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs
#	apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs
#	apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs
#	apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs
This commit is contained in:
fawney19
2026-05-19 02:23:14 +08:00
53 changed files with 8081 additions and 305 deletions

View File

@@ -121,7 +121,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
planner_state,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
false,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
input.routing_policy.as_ref(),

View File

@@ -13,7 +13,9 @@ mod gemini;
mod normalize;
mod openai;
pub(crate) use self::codex::apply_codex_openai_responses_special_headers;
pub(crate) use self::codex::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
};
pub(crate) use self::family::{
build_local_stream_attempt_source, build_local_stream_plan_and_reports,
build_local_sync_attempt_source, build_local_sync_plan_and_reports,

View File

@@ -9,6 +9,7 @@ pub(super) use self::payload::maybe_build_local_openai_chat_decision_payload_for
pub(super) use self::support::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_image_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
};

View File

@@ -89,7 +89,32 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
&mut extra_fields,
resolved.transport.provider.provider_type.as_str(),
);
if let Some(image_request_summary) = resolved.image_request_summary.as_ref() {
extra_fields.insert("image_request".to_string(), image_request_summary.clone());
}
if resolved
.provider_api_format
.eq_ignore_ascii_case("openai:image")
&& resolved
.transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web")
{
extra_fields.insert("chatgpt_web_image".to_string(), serde_json::json!(true));
extra_fields.insert(
"local_failover_policy".to_string(),
serde_json::json!({
"stop_status_codes": [400, 401, 403, 429, 500, 502, 503, 504],
"error_stop_patterns": [
{ "pattern": ".*" }
]
}),
);
}
let super::request::LocalOpenAiChatCandidatePayloadParts {
client_api_format,
auth_header,
auth_value,
mapped_model,
@@ -104,6 +129,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
transport,
request_redacted,
transport_profile: _,
image_request_summary: _,
} = resolved;
let original_request_body_json = if request_redacted {
Some(&provider_request_body)
@@ -128,7 +154,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: &provider_api_format,
client_api_format: "openai:chat",
client_api_format: &client_api_format,
mapped_model: Some(&mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
@@ -160,7 +186,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
}),
execution_strategy,
conversion_mode,
"openai:chat",
client_api_format.as_str(),
candidate.endpoint_api_format.as_str(),
),
&transport,

View File

@@ -4,7 +4,7 @@ use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_contracts::ResolvedTransportProfile;
use serde_json::Value;
use serde_json::{json, Value};
use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
@@ -16,9 +16,10 @@ use crate::ai_serving::planner::common::{
request_requires_body_stream_field, OPENAI_CHAT_STREAM_PLAN_KIND,
};
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
build_local_openai_chat_upstream_url, request_body_build_failure_extra_data,
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_upstream_url,
build_local_openai_chat_request_body, build_local_openai_chat_upstream_url,
request_body_build_failure_extra_data,
};
use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth;
use crate::ai_serving::transport::kiro::{
@@ -29,8 +30,10 @@ use crate::ai_serving::transport::kiro::{
use crate::ai_serving::transport::local_openai_chat_transport_unsupported_reason;
use crate::ai_serving::transport::{
build_grok_browser_headers, build_grok_upstream_url, build_kiro_cross_format_upstream_url,
build_standard_provider_request_headers, GrokHeaderInput, StandardProviderRequestHeadersInput,
GROK_CHAT_PATH,
build_openai_image_headers, build_openai_image_upstream_url,
build_standard_provider_request_headers, openai_image_transport_unsupported_reason,
resolve_openai_image_auth, GrokHeaderInput, ProviderOpenAiImageHeadersInput,
StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
};
use crate::ai_serving::{
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
@@ -53,6 +56,7 @@ use super::support::{
};
pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
pub(super) client_api_format: String,
pub(super) auth_header: String,
pub(super) auth_value: String,
pub(super) mapped_model: String,
@@ -67,6 +71,7 @@ pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) request_redacted: bool,
pub(super) transport_profile: Option<ResolvedTransportProfile>,
pub(super) image_request_summary: Option<Value>,
}
fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
@@ -157,6 +162,7 @@ async fn resolve_chat_pii_redaction_feature_settings(
.await
.map_err(|err| {
warn!(
error = ?err,
"gateway failed to read api key chat pii redaction feature settings"
);
@@ -309,6 +315,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
);
return Ok(Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
@@ -323,6 +330,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
transport: Arc::clone(transport),
request_redacted: redaction.redacted,
transport_profile,
image_request_summary: None,
}));
}
@@ -473,6 +481,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
);
return Ok(Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
mapped_model: prepared_candidate.mapped_model,
@@ -487,10 +496,26 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
transport: Arc::clone(transport),
request_redacted: redaction.redacted,
transport_profile,
image_request_summary: None,
}));
};
let provider_api_format = provider_api_format.trim().to_ascii_lowercase();
if provider_api_format == "openai:image" {
return resolve_openai_chat_to_openai_image_payload_parts(
state,
parts,
trace_id,
body_json,
input,
eligible,
candidate_index,
candidate_id,
upstream_is_stream,
)
.await;
}
let Some(conversion_kind) =
request_conversion_kind("openai:chat", provider_api_format.as_str())
else {
@@ -762,6 +787,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
ai_local_execution_contract_for_formats("openai:chat", provider_api_format.as_str());
Ok(Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
mapped_model: prepared_candidate.mapped_model,
@@ -776,9 +802,496 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
transport: Arc::clone(transport),
request_redacted: redaction.redacted,
transport_profile: None,
image_request_summary: None,
}))
}
#[allow(clippy::too_many_arguments)]
async fn resolve_openai_chat_to_openai_image_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
upstream_is_stream: bool,
) -> Result<Option<LocalOpenAiChatCandidatePayloadParts>, GatewayError> {
let candidate = &eligible.candidate;
let transport = &eligible.transport;
let provider_api_format = "openai:image";
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, provider_api_format)
{
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return Ok(None);
}
let prepared_candidate = match prepare_header_authenticated_candidate(
crate::ai_serving::PlannerAppState::new(state),
transport,
candidate,
resolve_openai_image_auth(transport),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "openai_chat_image_bridge",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_chat_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return Ok(None);
}
};
let is_chatgpt_web = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web");
let Some((mut provider_request_body, image_request_summary)) = (if is_chatgpt_web {
build_chatgpt_web_image_provider_body_from_openai_chat_body(
body_json,
&input.requested_model,
)
} else {
build_openai_image_provider_body_from_openai_chat_body(
body_json,
&input.requested_model,
upstream_is_stream,
)
}) else {
mark_skipped_local_openai_chat_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(body_json, "openai:chat", provider_api_format),
)
.await;
return Ok(None);
};
if !is_chatgpt_web {
apply_codex_openai_responses_special_body_edits(
&mut provider_request_body,
transport.provider.provider_type.as_str(),
provider_api_format,
transport.endpoint.body_rules.as_ref(),
Some(candidate.key_id.as_str()),
);
}
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())
};
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
})
else {
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
"openai:chat",
provider_api_format,
"openai_chat_image_bridge_headers",
),
)
.await;
return Ok(None);
};
if is_chatgpt_web {
provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else {
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
}
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", provider_api_format);
Ok(Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
execution_strategy,
conversion_mode,
report_kind: "openai_chat_stream_success".to_string(),
envelope_name: None,
transport: Arc::clone(transport),
request_redacted: false,
transport_profile: None,
image_request_summary: Some(image_request_summary),
}))
}
fn build_openai_image_provider_body_from_openai_chat_body(
body_json: &Value,
requested_model: &str,
upstream_is_stream: bool,
) -> Option<(Value, Value)> {
let (prompt, images) = collect_openai_chat_image_prompt_and_images(body_json)?;
let operation = if images.is_empty() {
"generate"
} else {
"edit"
};
let mut tool = serde_json::Map::new();
tool.insert(
"type".to_string(),
Value::String("image_generation".to_string()),
);
tool.insert("action".to_string(), Value::String(operation.to_string()));
copy_openai_chat_image_tool_option(body_json, &mut tool, "size");
copy_openai_chat_image_tool_option(body_json, &mut tool, "quality");
copy_openai_chat_image_tool_option(body_json, &mut tool, "background");
copy_openai_chat_image_tool_option(body_json, &mut tool, "output_format");
copy_openai_chat_image_tool_option(body_json, &mut tool, "output_compression");
copy_openai_chat_image_tool_option(body_json, &mut tool, "moderation");
copy_openai_chat_image_tool_option(body_json, &mut tool, "input_fidelity");
copy_openai_chat_image_tool_option(body_json, &mut tool, "partial_images");
let input = if images.is_empty() {
serde_json::json!([{
"role": "user",
"content": prompt,
}])
} else {
let mut content = vec![serde_json::json!({
"type": "input_text",
"text": prompt,
})];
content.extend(images);
serde_json::json!([{
"role": "user",
"content": content,
}])
};
let mut body = serde_json::Map::new();
if let Some(model) = body_json
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
let requested_model = requested_model.trim();
(!requested_model.is_empty()).then_some(requested_model)
})
{
body.insert("model".to_string(), Value::String(model.to_string()));
}
body.insert("input".to_string(), input);
body.insert(
"tools".to_string(),
Value::Array(vec![Value::Object(tool.clone())]),
);
if upstream_is_stream {
body.insert("stream".to_string(), Value::Bool(true));
}
if let Some(user) = body_json
.get("user")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
body.insert("user".to_string(), Value::String(user.to_string()));
}
let mut summary = serde_json::Map::new();
summary.insert(
"operation".to_string(),
Value::String(operation.to_string()),
);
for key in ["output_format", "partial_images", "size", "quality"] {
if let Some(value) = tool.get(key) {
summary.insert(key.to_string(), value.clone());
}
}
Some((Value::Object(body), Value::Object(summary)))
}
fn build_chatgpt_web_image_provider_body_from_openai_chat_body(
body_json: &Value,
requested_model: &str,
) -> Option<(Value, Value)> {
let (prompt, images) = collect_openai_chat_image_prompt_and_images(body_json)?;
let operation = if images.is_empty() {
"generate"
} else {
"edit"
};
let size = body_json
.get("size")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("1024x1024");
let output_format = body_json
.get("output_format")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("png");
let quality = body_json
.get("quality")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("medium");
let model = body_json
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or_else(|| requested_model.trim());
let web_model = body_json
.get("web_model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("gpt-5-5-thinking");
let image_urls = openai_image_inputs_as_urls(&images);
let body = json!({
"operation": operation,
"model": if model.is_empty() { "gpt-image-2" } else { model },
"web_model": web_model,
"prompt": prompt,
"size": size,
"ratio": chatgpt_web_ratio_for_size(size),
"output_format": output_format,
"images": image_urls,
});
let summary = json!({
"operation": operation,
"output_format": output_format,
"size": size,
"quality": quality,
});
Some((body, summary))
}
fn copy_openai_chat_image_tool_option(
body_json: &Value,
tool: &mut serde_json::Map<String, Value>,
key: &str,
) {
if let Some(value) = body_json.get(key) {
tool.insert(key.to_string(), value.clone());
}
}
fn collect_openai_chat_image_prompt_and_images(body_json: &Value) -> Option<(String, Vec<Value>)> {
let messages = body_json.get("messages").and_then(Value::as_array)?;
let mut prompt_parts = Vec::new();
let mut images = Vec::new();
for message in messages.iter().filter_map(Value::as_object) {
let role = message
.get("role")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
let content = message.get("content");
if matches!(role, "system" | "developer" | "user") {
if let Some(text) = crate::ai_serving::extract_openai_text_content(content)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
{
prompt_parts.push(text);
}
}
if role == "user" {
collect_openai_chat_image_inputs(content, &mut images);
}
}
let prompt = prompt_parts.join("\n").trim().to_string();
(!prompt.is_empty()).then_some((prompt, images))
}
fn collect_openai_chat_image_inputs(content: Option<&Value>, images: &mut Vec<Value>) {
let Some(parts) = content.and_then(Value::as_array) else {
return;
};
for part in parts.iter().filter_map(Value::as_object) {
let part_type = part
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if matches!(part_type, "image_url" | "input_image") {
if let Some(url) = part
.get("image_url")
.and_then(|value| {
value
.as_str()
.or_else(|| value.get("url").and_then(Value::as_str))
})
.map(str::trim)
.filter(|value| !value.is_empty())
{
images.push(serde_json::json!({
"type": "input_image",
"image_url": url,
}));
} else if let Some(file_id) = part
.get("file_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
images.push(serde_json::json!({
"type": "input_image",
"file_id": file_id,
}));
}
}
}
}
fn openai_image_inputs_as_urls(images: &[Value]) -> Vec<Value> {
images
.iter()
.filter_map(|image| {
image
.get("image_url")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| Value::String(value.to_string()))
})
.collect()
}
fn chatgpt_web_ratio_for_size(size: &str) -> String {
let Some((width, height)) = size.split_once('x') else {
return "1:1".to_string();
};
let Ok(width) = width.trim().parse::<u64>() else {
return "1:1".to_string();
};
let Ok(height) = height.trim().parse::<u64>() else {
return "1:1".to_string();
};
if width == 0 || height == 0 {
return "1:1".to_string();
}
let divisor = gcd(width, height);
format!("{}:{}", width / divisor, height / divisor)
}
fn gcd(mut left: u64, mut right: u64) -> u64 {
while right != 0 {
let next = left % right;
left = right;
right = next;
}
left.max(1)
}
fn chatgpt_web_image_internal_url(base_url: &str) -> String {
let base_url = base_url.trim().trim_end_matches('/');
let base_url = if base_url.is_empty() {
"https://chatgpt.com"
} else {
base_url
};
format!("{base_url}/__aether/chatgpt-web-image")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn chatgpt_web_chat_image_bridge_body_uses_internal_web_shape() {
let body_json = json!({
"model": "gpt-image-2",
"messages": [
{"role": "system", "content": "Use crisp vector-like shapes."},
{
"role": "user",
"content": [
{"type": "text", "text": "Draw a glass city"},
{"type": "image_url", "image_url": {"url": "https://example.com/ref.png"}}
]
}
],
"size": "1536x1024",
"output_format": "webp",
"web_model": "gpt-5-image-test"
});
let (provider_body, summary) =
build_chatgpt_web_image_provider_body_from_openai_chat_body(&body_json, "gpt-image-2")
.expect("chat image body should convert");
assert_eq!(provider_body["operation"], "edit");
assert_eq!(provider_body["model"], "gpt-image-2");
assert_eq!(provider_body["web_model"], "gpt-5-image-test");
assert_eq!(
provider_body["prompt"],
"Use crisp vector-like shapes.\nDraw a glass city"
);
assert_eq!(provider_body["size"], "1536x1024");
assert_eq!(provider_body["ratio"], "3:2");
assert_eq!(provider_body["output_format"], "webp");
assert_eq!(provider_body["images"][0], "https://example.com/ref.png");
assert_eq!(summary["operation"], "edit");
assert_eq!(summary["output_format"], "webp");
}
}
#[allow(clippy::too_many_arguments)]
async fn build_kiro_openai_chat_cross_format_payload_parts(
state: &AppState,
@@ -900,6 +1413,7 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
ai_local_execution_contract_for_formats("openai:chat", provider_api_format);
Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:chat".to_string(),
auth_header,
auth_value,
mapped_model,
@@ -914,6 +1428,7 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
transport: Arc::clone(transport),
request_redacted,
transport_profile: None,
image_request_summary: None,
})
}

View File

@@ -14,7 +14,10 @@ use crate::ai_serving::planner::candidate_metadata::{
LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::planner::candidate_source::LocalCandidatePreselectionKeyMode;
use crate::ai_serving::planner::candidate_source::{
preselect_local_execution_candidates_for_api_formats_with_serving,
LocalCandidatePreselectionKeyMode,
};
use crate::ai_serving::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
@@ -23,7 +26,7 @@ use crate::ai_serving::{
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
ExecutionRuntimeAuthContext, PlannerAppState,
};
use crate::AppState;
use crate::{AppState, GatewayError};
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiChatCandidateAttempt;
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttemptSource as LocalOpenAiChatCandidateAttemptSource;
@@ -354,3 +357,95 @@ pub(crate) async fn build_lazy_local_openai_chat_candidate_attempt_source<'a>(
)
.await
}
pub(crate) async fn build_local_openai_chat_image_candidate_attempt_source<'a>(
state: &'a AppState,
trace_id: &str,
input: &LocalOpenAiChatDecisionInput,
body_json: &serde_json::Value,
) -> Result<(LocalOpenAiChatCandidateAttemptSource<'a>, usize), GatewayError> {
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiChatDecision,
);
let preselection = preselect_local_execution_candidates_for_api_formats_with_serving(
planner_state,
"openai:chat",
&input.requested_model,
false,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
input.routing_policy.as_ref(),
input.client_session_affinity.as_ref(),
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
vec!["openai:image".to_string()],
)
.await?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
"openai:chat",
Some(&input.requested_model),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
input.routing_policy.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
preselection.candidates,
preselection.skipped_candidates,
LocalCandidateResolutionMode::WithoutTransportPairGate,
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", &provider_api_format);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: "openai:chat",
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.trim(),
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", &provider_api_format);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
"openai:chat",
serde_json::Map::new(),
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
)
.await)
}

View File

@@ -11,6 +11,7 @@ mod plans;
use self::decision::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_image_candidate_attempt_source,
maybe_build_local_openai_chat_decision_payload_for_candidate, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
};

View File

@@ -1,8 +1,10 @@
use async_trait::async_trait;
use tracing::warn;
use super::super::super::openai_request_is_image_generation_intent;
use super::super::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_image_candidate_attempt_source,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
LocalOpenAiChatDecisionInput,
@@ -49,14 +51,53 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
};
let effective_body_json = input.effective_body_json(body_json).clone();
let (candidates, candidate_count) = build_lazy_local_openai_chat_candidate_attempt_source(
state,
trace_id,
&input,
&effective_body_json,
true,
)
.await;
let image_generation_intent =
openai_request_is_image_generation_intent(&input.requested_model, body_json);
let (mut candidates, mut candidate_count) = if image_generation_intent {
let (image_candidates, image_candidate_count) =
build_local_openai_chat_image_candidate_attempt_source(
state,
trace_id,
&input,
&effective_body_json,
)
.await?;
if image_candidate_count > 0 {
(image_candidates, image_candidate_count)
} else {
build_lazy_local_openai_chat_candidate_attempt_source(
state,
trace_id,
&input,
&effective_body_json,
true,
)
.await
}
} else {
build_lazy_local_openai_chat_candidate_attempt_source(
state,
trace_id,
&input,
&effective_body_json,
true,
)
.await
};
if !image_generation_intent && candidate_count == 0 {
let (image_candidates, image_candidate_count) =
build_local_openai_chat_image_candidate_attempt_source(
state,
trace_id,
&input,
&effective_body_json,
)
.await?;
if image_candidate_count > 0 {
candidates = image_candidates;
candidate_count = image_candidate_count;
}
}
if candidate_count == 0 {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,

View File

@@ -0,0 +1,75 @@
pub(crate) fn openai_request_is_image_generation_intent(
requested_model: &str,
body_json: &serde_json::Value,
) -> bool {
openai_model_is_image_generation(requested_model)
|| body_json
.get("model")
.and_then(serde_json::Value::as_str)
.is_some_and(openai_model_is_image_generation)
|| openai_tool_choice_selects_image_generation(body_json.get("tool_choice"))
}
fn openai_model_is_image_generation(model: &str) -> bool {
model.trim().to_ascii_lowercase().starts_with("gpt-image-")
}
fn openai_tool_choice_selects_image_generation(choice: Option<&serde_json::Value>) -> bool {
let Some(choice) = choice else {
return false;
};
if let Some(value) = choice.as_str() {
return value.trim().eq_ignore_ascii_case("image_generation");
}
let Some(object) = choice.as_object() else {
return false;
};
object
.get("type")
.and_then(serde_json::Value::as_str)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("image_generation"))
|| object
.get("tool")
.and_then(|value| value.get("type"))
.and_then(serde_json::Value::as_str)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("image_generation"))
|| object
.get("function")
.and_then(|value| value.get("name"))
.and_then(serde_json::Value::as_str)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("image_generation"))
}
#[cfg(test)]
mod tests {
use super::openai_request_is_image_generation_intent;
use serde_json::json;
#[test]
fn detects_openai_image_generation_intent_like_compat_proxies() {
assert!(openai_request_is_image_generation_intent(
"GPT-IMAGE-2",
&json!({})
));
assert!(openai_request_is_image_generation_intent(
"gpt-5",
&json!({"model":"gpt-image-2"})
));
assert!(openai_request_is_image_generation_intent(
"gpt-5",
&json!({"tool_choice":{"function":{"name":"image_generation"}}})
));
assert!(openai_request_is_image_generation_intent(
"gpt-5",
&json!({"tool_choice":{"type":"image_generation"}})
));
assert!(!openai_request_is_image_generation_intent(
"gpt-5",
&json!({"tools":[{"type":"image_generation"}]})
));
assert!(!openai_request_is_image_generation_intent(
"gpt-5",
&json!({"messages":[{"role":"user","content":"hello"}]})
));
}
}

View File

@@ -1,4 +1,5 @@
mod chat;
mod image_intent;
mod responses;
pub(crate) use crate::ai_serving::{
@@ -14,6 +15,7 @@ pub(crate) use chat::{
maybe_build_stream_local_decision_payload, maybe_build_sync_local_decision_payload,
set_local_openai_chat_execution_exhausted_diagnostic,
};
pub(super) use image_intent::openai_request_is_image_generation_intent;
pub(crate) use responses::{
build_local_openai_responses_stream_attempt_source_for_kind,
build_local_openai_responses_stream_plan_and_reports_for_kind,

View File

@@ -81,6 +81,30 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
if let Some(envelope_name) = resolved.envelope_name {
extra_fields.insert("envelope_name".to_string(), json!(envelope_name));
}
if let Some(image_request_summary) = resolved.image_request_summary.as_ref() {
extra_fields.insert("image_request".to_string(), image_request_summary.clone());
}
if resolved
.provider_api_format
.eq_ignore_ascii_case("openai:image")
&& resolved
.transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web")
{
extra_fields.insert("chatgpt_web_image".to_string(), json!(true));
extra_fields.insert(
"local_failover_policy".to_string(),
json!({
"stop_status_codes": [400, 401, 403, 429, 500, 502, 503, 504],
"error_stop_patterns": [
{ "pattern": ".*" }
]
}),
);
}
insert_provider_stream_event_api_format(
&mut extra_fields,
resolved.transport.provider.provider_type.as_str(),
@@ -179,6 +203,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
upstream_is_stream,
transport,
transport_profile: _,
image_request_summary: _,
} = resolved;
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {

View File

@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
use std::sync::Arc;
use aether_contracts::ResolvedTransportProfile;
use serde_json::Value;
use serde_json::{json, Value};
use tracing::debug;
use crate::ai_serving::planner::candidate_preparation::{
@@ -16,7 +16,8 @@ use crate::ai_serving::planner::common::{
};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_responses_request_body,
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
build_cross_format_openai_responses_request_body,
build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body,
build_local_openai_responses_upstream_url, request_body_build_failure_extra_data,
};
@@ -37,9 +38,11 @@ use crate::ai_serving::transport::kiro::{
};
use crate::ai_serving::transport::{
build_grok_browser_headers, build_grok_upstream_url, build_kiro_cross_format_upstream_url,
build_openai_image_headers, build_openai_image_upstream_url,
build_standard_provider_request_headers,
local_standard_transport_unsupported_reason_with_network, GrokHeaderInput,
StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
local_standard_transport_unsupported_reason_with_network,
openai_image_transport_unsupported_reason, resolve_openai_image_auth, GrokHeaderInput,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
};
use crate::ai_serving::{
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
@@ -81,6 +84,7 @@ pub(crate) struct LocalOpenAiResponsesCandidatePayloadParts {
pub(super) upstream_is_stream: bool,
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) transport_profile: Option<ResolvedTransportProfile>,
pub(super) image_request_summary: Option<Value>,
}
#[allow(clippy::too_many_arguments)]
@@ -110,6 +114,21 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
.trim()
.eq_ignore_ascii_case("grok");
if provider_api_format.eq_ignore_ascii_case("openai:image") {
return resolve_openai_responses_to_openai_image_payload_parts(
state,
parts,
trace_id,
body_json,
input,
eligible,
candidate_index,
candidate_id,
spec,
)
.await;
}
let same_format = api_format_alias_matches(provider_api_format, &client_api_format);
let conversion_kind = request_conversion_kind(spec_metadata.api_format, provider_api_format);
let transport_unsupported_reason = if is_grok
@@ -596,6 +615,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
upstream_is_stream,
transport: Arc::clone(transport),
transport_profile,
image_request_summary: None,
})
}
@@ -603,6 +623,463 @@ fn api_format_alias_matches(left: &str, right: &str) -> bool {
crate::ai_serving::api_format_alias_matches(left, right)
}
#[allow(clippy::too_many_arguments)]
async fn resolve_openai_responses_to_openai_image_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiResponsesDecisionInput,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
spec: LocalOpenAiResponsesSpec,
) -> Option<LocalOpenAiResponsesCandidatePayloadParts> {
let spec_metadata = local_openai_responses_spec_metadata(spec);
let candidate = &eligible.candidate;
let transport = &eligible.transport;
let provider_api_format = "openai:image";
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, provider_api_format)
{
mark_skipped_local_openai_responses_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
PlannerAppState::new(state),
transport,
candidate,
resolve_openai_image_auth(transport),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "openai_responses_image_bridge",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_responses_candidate(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
skip_reason,
)
.await;
return None;
}
};
let is_chatgpt_web = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web");
let upstream_is_stream = resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
provider_api_format,
spec_metadata.require_streaming,
false,
);
let Some((mut provider_request_body, image_request_summary)) = (if is_chatgpt_web {
build_chatgpt_web_image_provider_body_from_openai_responses_body(
body_json,
&input.requested_model,
)
} else {
build_openai_image_provider_body_from_openai_responses_body(
body_json,
&input.requested_model,
upstream_is_stream,
)
}) else {
mark_skipped_local_openai_responses_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
body_json,
spec_metadata.api_format,
provider_api_format,
),
)
.await;
return None;
};
if !is_chatgpt_web {
apply_codex_openai_responses_special_body_edits(
&mut provider_request_body,
transport.provider.provider_type.as_str(),
provider_api_format,
transport.endpoint.body_rules.as_ref(),
Some(candidate.key_id.as_str()),
);
}
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())
};
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
})
else {
mark_skipped_local_openai_responses_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
candidate_index,
candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
spec_metadata.api_format,
provider_api_format,
"openai_responses_image_bridge_headers",
),
)
.await;
return None;
};
if is_chatgpt_web {
provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else {
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&provider_request_body,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
}
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, provider_api_format);
Some(LocalOpenAiResponsesCandidatePayloadParts {
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body,
provider_request_headers,
upstream_url,
execution_strategy,
conversion_mode,
is_antigravity: false,
envelope_name: None,
upstream_is_stream,
transport: Arc::clone(transport),
transport_profile: None,
image_request_summary: Some(image_request_summary),
})
}
fn build_openai_image_provider_body_from_openai_responses_body(
body_json: &Value,
requested_model: &str,
upstream_is_stream: bool,
) -> Option<(Value, Value)> {
let object = body_json.as_object()?;
let input = object.get("input")?.clone();
let mut tool = openai_responses_image_generation_tool(object).unwrap_or_else(|| {
serde_json::Map::from_iter([("type".to_string(), json!("image_generation"))])
});
tool.entry("type".to_string())
.or_insert_with(|| json!("image_generation"));
tool.entry("action".to_string())
.or_insert_with(|| json!("generate"));
let mut body = serde_json::Map::new();
body.insert("input".to_string(), input);
body.insert(
"tools".to_string(),
Value::Array(vec![Value::Object(tool.clone())]),
);
if let Some(model) = object
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
let requested_model = requested_model.trim();
(!requested_model.is_empty()).then_some(requested_model)
})
{
body.insert("model".to_string(), Value::String(model.to_string()));
}
for key in [
"user",
"metadata",
"include",
"parallel_tool_calls",
"store",
] {
if let Some(value) = object.get(key) {
body.insert(key.to_string(), value.clone());
}
}
if upstream_is_stream {
body.insert("stream".to_string(), Value::Bool(true));
} else if let Some(value) = object.get("stream") {
body.insert("stream".to_string(), value.clone());
}
let mut summary = serde_json::Map::new();
summary.insert(
"operation".to_string(),
tool.get("action")
.cloned()
.unwrap_or_else(|| json!("generate")),
);
for key in ["output_format", "partial_images", "size", "quality"] {
if let Some(value) = tool.get(key).or_else(|| object.get(key)) {
summary.insert(key.to_string(), value.clone());
}
}
Some((Value::Object(body), Value::Object(summary)))
}
fn openai_responses_image_generation_tool(
object: &serde_json::Map<String, Value>,
) -> Option<serde_json::Map<String, Value>> {
object
.get("tools")
.and_then(Value::as_array)?
.iter()
.filter_map(Value::as_object)
.find(|tool| {
tool.get("type")
.and_then(Value::as_str)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("image_generation"))
})
.cloned()
}
fn build_chatgpt_web_image_provider_body_from_openai_responses_body(
body_json: &Value,
requested_model: &str,
) -> Option<(Value, Value)> {
let object = body_json.as_object()?;
let (prompt, images) = collect_openai_responses_image_prompt_and_images(object.get("input"))?;
let operation = if images.is_empty() {
"generate"
} else {
"edit"
};
let tool = openai_responses_image_generation_tool(object);
let size = image_option_string(tool.as_ref(), object, "size").unwrap_or("1024x1024");
let output_format =
image_option_string(tool.as_ref(), object, "output_format").unwrap_or("png");
let quality = image_option_string(tool.as_ref(), object, "quality").unwrap_or("medium");
let model = object
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or_else(|| requested_model.trim());
let web_model = image_option_string(tool.as_ref(), object, "web_model")
.or_else(|| image_option_string(tool.as_ref(), object, "model"))
.unwrap_or("gpt-5-5-thinking");
let image_urls = openai_image_inputs_as_urls(&images);
let body = json!({
"operation": operation,
"model": if model.is_empty() { "gpt-image-2" } else { model },
"web_model": web_model,
"prompt": prompt,
"size": size,
"ratio": chatgpt_web_ratio_for_size(size),
"output_format": output_format,
"images": image_urls,
});
let summary = json!({
"operation": operation,
"output_format": output_format,
"size": size,
"quality": quality,
});
Some((body, summary))
}
fn image_option_string<'a>(
tool: Option<&'a serde_json::Map<String, Value>>,
object: &'a serde_json::Map<String, Value>,
key: &str,
) -> Option<&'a str> {
tool.and_then(|tool| tool.get(key))
.or_else(|| object.get(key))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn collect_openai_responses_image_prompt_and_images(
input: Option<&Value>,
) -> Option<(String, Vec<Value>)> {
let input = input?;
let mut prompt_parts = Vec::new();
let mut images = Vec::new();
collect_openai_responses_image_input(input, &mut prompt_parts, &mut images);
let prompt = prompt_parts.join("\n").trim().to_string();
(!prompt.is_empty()).then_some((prompt, images))
}
fn collect_openai_responses_image_input(
value: &Value,
prompt_parts: &mut Vec<String>,
images: &mut Vec<Value>,
) {
match value {
Value::String(text) => {
let text = text.trim();
if !text.is_empty() {
prompt_parts.push(text.to_string());
}
}
Value::Array(items) => {
for item in items {
collect_openai_responses_image_input(item, prompt_parts, images);
}
}
Value::Object(object) => {
let item_type = object
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if matches!(item_type, "input_text" | "text") {
if let Some(text) = object
.get("text")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
prompt_parts.push(text.to_string());
}
} else if matches!(item_type, "input_image" | "image_url") {
collect_openai_image_input_object(object, images);
}
if let Some(content) = object.get("content") {
collect_openai_responses_image_input(content, prompt_parts, images);
}
}
_ => {}
}
}
fn collect_openai_image_input_object(
object: &serde_json::Map<String, Value>,
images: &mut Vec<Value>,
) {
if let Some(url) = object
.get("image_url")
.and_then(|value| {
value
.as_str()
.or_else(|| value.get("url").and_then(Value::as_str))
})
.or_else(|| object.get("url").and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
{
images.push(json!({
"type": "input_image",
"image_url": url,
}));
} else if let Some(file_id) = object
.get("file_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
images.push(json!({
"type": "input_image",
"file_id": file_id,
}));
}
}
fn openai_image_inputs_as_urls(images: &[Value]) -> Vec<Value> {
images
.iter()
.filter_map(|image| {
image
.get("image_url")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| Value::String(value.to_string()))
})
.collect()
}
fn chatgpt_web_ratio_for_size(size: &str) -> String {
let Some((width, height)) = size.split_once('x') else {
return "1:1".to_string();
};
let Ok(width) = width.trim().parse::<u64>() else {
return "1:1".to_string();
};
let Ok(height) = height.trim().parse::<u64>() else {
return "1:1".to_string();
};
if width == 0 || height == 0 {
return "1:1".to_string();
}
let divisor = gcd(width, height);
format!("{}:{}", width / divisor, height / divisor)
}
fn gcd(mut left: u64, mut right: u64) -> u64 {
while right != 0 {
let next = left % right;
left = right;
right = next;
}
left.max(1)
}
fn chatgpt_web_image_internal_url(base_url: &str) -> String {
let base_url = base_url.trim().trim_end_matches('/');
let base_url = if base_url.is_empty() {
"https://chatgpt.com"
} else {
base_url
};
format!("{base_url}/__aether/chatgpt-web-image")
}
#[allow(clippy::too_many_arguments)]
async fn build_kiro_openai_responses_payload_parts(
state: &AppState,
@@ -748,5 +1225,6 @@ async fn build_kiro_openai_responses_payload_parts(
upstream_is_stream,
transport: Arc::clone(transport),
transport_profile: None,
image_request_summary: None,
})
}

View File

@@ -15,6 +15,7 @@ use crate::ai_serving::planner::candidate_metadata::{
LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_source::{
preselect_local_execution_candidates_for_api_formats_with_serving,
preselect_local_execution_candidates_with_serving, LocalCandidatePreselectionKeyMode,
};
use crate::ai_serving::planner::common::extract_standard_requested_model;
@@ -36,6 +37,7 @@ use crate::ai_serving::{
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::{AppState, GatewayError};
use super::super::super::openai_request_is_image_generation_intent;
use super::LocalOpenAiResponsesSpec;
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalOpenAiResponsesCandidateAttempt;
@@ -265,6 +267,16 @@ pub(crate) async fn build_local_openai_responses_candidate_attempt_source<'a>(
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision,
);
if openai_request_is_image_generation_intent(&input.requested_model, body_json) {
let (image_candidates, image_candidate_count) =
build_local_openai_responses_image_candidate_attempt_source(
state, trace_id, input, body_json, spec,
)
.await?;
if image_candidate_count > 0 {
return Ok((image_candidates, image_candidate_count));
}
}
Ok(
build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state,
@@ -335,6 +347,104 @@ pub(crate) async fn build_local_openai_responses_candidate_attempt_source<'a>(
)
}
pub(crate) async fn build_local_openai_responses_image_candidate_attempt_source<'a>(
state: &'a AppState,
trace_id: &str,
input: &LocalOpenAiResponsesDecisionInput,
body_json: &serde_json::Value,
spec: LocalOpenAiResponsesSpec,
) -> Result<(LocalOpenAiResponsesCandidateAttemptSource<'a>, usize), GatewayError> {
let spec_metadata = local_openai_responses_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision,
);
let preselection = preselect_local_execution_candidates_for_api_formats_with_serving(
planner_state,
spec_metadata.api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
input.routing_policy.as_ref(),
input.client_session_affinity.as_ref(),
true,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
vec!["openai:image".to_string()],
)
.await?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
spec_metadata.api_format,
Some(&input.requested_model),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
input.routing_policy.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
preselection.candidates,
preselection.skipped_candidates,
LocalCandidateResolutionMode::WithoutTransportPairGate,
move |eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.as_str(),
))
},
move |mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
spec_metadata.api_format,
serde_json::Map::new(),
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
)
.await)
}
pub(crate) async fn mark_skipped_local_openai_responses_candidate(
state: &AppState,
input: &LocalOpenAiResponsesDecisionInput,

View File

@@ -3,20 +3,22 @@ use std::io::Error as IoError;
use std::time::Instant;
use aether_contracts::{
ExecutionPlan, ExecutionResult, ExecutionTelemetry, RequestBody, ResolvedTransportProfile,
ResponseBody, StreamFrame, StreamFramePayload, StreamFrameType,
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_HTTP_MODE_AUTO, TRANSPORT_POOL_SCOPE_KEY,
ExecutionPlan, ExecutionResult, ExecutionStreamTerminalSummary, ExecutionTelemetry,
RequestBody, ResolvedTransportProfile, ResponseBody, StreamFrame, StreamFramePayload,
StreamFrameType, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ,
TRANSPORT_HTTP_MODE_AUTO, TRANSPORT_POOL_SCOPE_KEY,
};
use axum::body::Bytes;
use base64::Engine as _;
use chrono::{FixedOffset, Utc};
use futures_util::stream::{self, BoxStream};
use futures_util::StreamExt;
use serde_json::{json, Value};
use serde_json::{json, Map, Value};
use tracing::debug;
use uuid::Uuid;
use crate::ai_serving::api::StreamingStandardTerminalObserver;
use crate::clock::current_unix_secs;
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
use crate::execution_runtime::transport::{
@@ -111,16 +113,15 @@ pub(crate) async fn maybe_execute_chatgpt_web_image_stream(
Err(err) => chatgpt_web_transport_error_execution_result(plan, started_at, &err),
};
Ok(Some(ChatGptWebImageStream {
frame_stream: execution_result_frame_stream(&result),
frame_stream: execution_result_frame_stream(plan, &result, report_context),
report_context: report_context.cloned(),
}))
}
fn is_chatgpt_web_image_plan(plan: &ExecutionPlan, report_context: Option<&Value>) -> bool {
if !plan.client_api_format.eq_ignore_ascii_case("openai:image")
|| !plan
.provider_api_format
.eq_ignore_ascii_case("openai:image")
if !plan
.provider_api_format
.eq_ignore_ascii_case("openai:image")
{
return false;
}
@@ -1510,9 +1511,12 @@ fn bytes_execution_result(
}
fn execution_result_frame_stream(
plan: &ExecutionPlan,
result: &ExecutionResult,
report_context: Option<&Value>,
) -> BoxStream<'static, Result<Bytes, IoError>> {
let body = execution_result_body_bytes_lossy(result);
let terminal_summary = chatgpt_web_stream_terminal_summary(plan, result, report_context, &body);
let mut frames = vec![
StreamFrame {
frame_type: StreamFrameType::Headers,
@@ -1551,7 +1555,7 @@ fn execution_result_frame_stream(
}),
},
});
frames.push(StreamFrame::eof());
frames.push(StreamFrame::eof_with_summary(terminal_summary));
stream::iter(
frames
.into_iter()
@@ -1560,6 +1564,81 @@ fn execution_result_frame_stream(
.boxed()
}
fn chatgpt_web_stream_terminal_summary(
plan: &ExecutionPlan,
result: &ExecutionResult,
report_context: Option<&Value>,
body: &[u8],
) -> Option<ExecutionStreamTerminalSummary> {
if !(200..300).contains(&result.status_code) || body.is_empty() {
return None;
}
let observer_context = chatgpt_web_stream_observer_context(plan, report_context);
let mut observer = StreamingStandardTerminalObserver::default();
let mut line_start = 0usize;
for (index, byte) in body.iter().enumerate() {
if *byte != b'\n' {
continue;
}
observer
.push_line(&observer_context, body[line_start..=index].to_vec())
.ok()?;
line_start = index.saturating_add(1);
}
if line_start < body.len() {
observer
.push_line(&observer_context, body[line_start..].to_vec())
.ok()?;
}
observer.finish(&observer_context).ok().flatten()
}
fn chatgpt_web_stream_observer_context(
plan: &ExecutionPlan,
report_context: Option<&Value>,
) -> Value {
let mut context = report_context
.cloned()
.filter(Value::is_object)
.unwrap_or_else(|| json!({}));
let object = context
.as_object_mut()
.expect("observer context should be an object");
object
.entry("provider_api_format".to_string())
.or_insert_with(|| Value::String(plan.provider_api_format.clone()));
object
.entry("client_api_format".to_string())
.or_insert_with(|| Value::String(plan.client_api_format.clone()));
object
.entry("model".to_string())
.or_insert_with(|| Value::String(plan.model_name.clone().unwrap_or_default()));
if !object.contains_key("image_request") {
if let Some(image_request) = chatgpt_web_image_request_context(plan) {
object.insert("image_request".to_string(), image_request);
}
}
context
}
fn chatgpt_web_image_request_context(plan: &ExecutionPlan) -> Option<Value> {
let body = plan.body.json_body.as_ref()?.as_object()?;
let mut image_request = Map::new();
image_request.insert(
"operation".to_string(),
Value::String("generate".to_string()),
);
for key in ["model", "size", "quality", "output_format"] {
if let Some(value) = body.get(key).and_then(Value::as_str).map(str::trim) {
if !value.is_empty() {
image_request.insert(key.to_string(), Value::String(value.to_string()));
}
}
}
Some(Value::Object(image_request))
}
fn telemetry(started_at: Instant, upstream_bytes: u64) -> ExecutionTelemetry {
let elapsed_ms = started_at.elapsed().as_millis() as u64;
ExecutionTelemetry {
@@ -2280,6 +2359,31 @@ data: [DONE]
assert!(decoded_data.contains("\"width\":2"));
assert!(decoded_data.contains("\"height\":3"));
assert!(text.contains("\"type\":\"eof\""));
let eof_frame = text
.lines()
.filter_map(|line| serde_json::from_str::<Value>(line).ok())
.find(|frame| frame.get("type").and_then(Value::as_str) == Some("eof"))
.expect("eof frame should exist");
assert_eq!(
eof_frame
.get("payload")
.and_then(|payload| payload.get("summary"))
.and_then(|summary| summary.get("standardized_usage"))
.and_then(|usage| usage.get("dimensions"))
.and_then(|dimensions| dimensions.get("image_count"))
.and_then(Value::as_u64),
Some(1)
);
assert_eq!(
eof_frame
.get("payload")
.and_then(|payload| payload.get("summary"))
.and_then(|summary| summary.get("standardized_usage"))
.and_then(|usage| usage.get("dimensions"))
.and_then(|dimensions| dimensions.get("image_size"))
.and_then(Value::as_str),
Some("1024x1024")
);
handle.abort();
}
@@ -2314,6 +2418,32 @@ data: [DONE]
assert_eq!(body["error"]["code"], "chatgpt_web_image_unsupported");
}
#[tokio::test]
async fn chatgpt_web_image_executor_accepts_marked_responses_client_plan() {
let state = crate::AppState::new().expect("state should build");
let mut plan = sample_plan(
CHATGPT_WEB_DEFAULT_BASE_URL,
json!({
"error": {
"message": "ChatGPT-Web 不支持该分辨率",
"type": "invalid_request_error",
"code": "chatgpt_web_image_unsupported"
}
}),
false,
);
plan.client_api_format = "openai:responses".to_string();
let result = maybe_execute_chatgpt_web_image_sync(&state, &plan, None)
.await
.expect("executor should run")
.expect("marked image provider plan should be intercepted");
assert_eq!(result.status_code, 400);
let body = execution_result_json(&result).expect("error should be json");
assert_eq!(body["error"]["code"], "chatgpt_web_image_unsupported");
}
#[tokio::test]
async fn chatgpt_web_image_stream_path_wraps_executor_result_as_ndjson_frames() {
let state = crate::AppState::new().expect("state should build");

View File

@@ -13,6 +13,16 @@ fn sync_plan_kind_disables_local_candidate_failover(plan_kind: &str) -> bool {
)
}
fn openai_image_success_disables_local_success_failover(
plan: &ExecutionPlan,
status_code: u16,
) -> bool {
status_code == 200
&& plan
.provider_api_format
.eq_ignore_ascii_case("openai:image")
}
pub(crate) async fn should_retry_next_local_candidate_sync(
state: &AppState,
plan: &ExecutionPlan,
@@ -48,6 +58,10 @@ pub(crate) async fn analyze_local_candidate_failover_sync(
return LocalFailoverAnalysis::use_default();
}
if openai_image_success_disables_local_success_failover(plan, result.status_code) {
return LocalFailoverAnalysis::use_default();
}
resolve_local_failover_analysis_for_attempt(
state,
plan,
@@ -218,6 +232,10 @@ pub(crate) async fn resolve_local_candidate_failover_analysis_stream(
status_code: u16,
response_text: Option<&str>,
) -> LocalFailoverAnalysis {
if openai_image_success_disables_local_success_failover(plan, status_code) {
return LocalFailoverAnalysis::use_default();
}
resolve_local_failover_analysis_for_attempt(
state,
plan,
@@ -756,6 +774,75 @@ mod tests {
);
}
#[tokio::test]
async fn stream_success_failover_does_not_retry_openai_image_success() {
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"success_failover_patterns": [
{"pattern": ".*"}
]
}
})));
let mut plan = sample_plan();
plan.provider_api_format = "openai:image".to_string();
assert!(
!should_retry_next_local_candidate_stream(
&state,
&plan,
"openai_image_stream",
Some(&local_report_context),
200,
Some("{\"data\":[{\"b64_json\":\"aGVsbG8=\"}]}"),
)
.await,
"successful OpenAI image responses should not be retried by success failover rules"
);
}
#[tokio::test]
async fn sync_success_failover_does_not_retry_openai_image_success() {
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"success_failover_patterns": [
{"pattern": ".*"}
]
}
})));
let mut plan = sample_plan();
plan.provider_api_format = "openai:image".to_string();
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 200,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
assert!(
!should_retry_next_local_candidate_sync(
&state,
&plan,
"openai_image_sync",
Some(&local_report_context),
&result,
Some("{\"data\":[{\"b64_json\":\"aGVsbG8=\"}]}")
)
.await,
"successful OpenAI image responses should not be retried by success failover rules"
);
}
#[test]
fn resolve_local_failover_policy_reads_provider_rules() {
let state = build_state_with_provider_config(Some(serde_json::json!({

View File

@@ -1044,6 +1044,12 @@ 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]"))
}
async fn next_stream_frame<R>(
buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>,
@@ -1088,6 +1094,33 @@ fn should_refresh_stream_usage_telemetry(
|| (next_elapsed.is_some() && next_elapsed != previous_elapsed)
}
fn build_terminal_stream_telemetry(
stream_started_at: Instant,
telemetry: Option<&ExecutionTelemetry>,
usage_stream_telemetry: Option<&ExecutionTelemetry>,
upstream_bytes: u64,
) -> ExecutionTelemetry {
let current_elapsed_ms = stream_started_at
.elapsed()
.as_millis()
.min(u128::from(u64::MAX)) as u64;
let ttfb_ms = telemetry
.and_then(|telemetry| telemetry.ttfb_ms)
.or_else(|| usage_stream_telemetry.and_then(|telemetry| telemetry.ttfb_ms));
let prior_elapsed_ms = telemetry
.and_then(|telemetry| telemetry.elapsed_ms)
.or_else(|| usage_stream_telemetry.and_then(|telemetry| telemetry.elapsed_ms))
.unwrap_or(0);
let elapsed_ms = current_elapsed_ms
.max(prior_elapsed_ms)
.max(ttfb_ms.unwrap_or(0));
ExecutionTelemetry {
ttfb_ms,
elapsed_ms: Some(elapsed_ms),
upstream_bytes: Some(upstream_bytes),
}
}
fn should_skip_direct_finalize_prefetch(
direct_stream_finalize_kind: Option<&str>,
content_type: Option<&str>,
@@ -2052,6 +2085,8 @@ async fn execute_stream_from_frame_stream(
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let mut client_visible_stream_completed =
stream_chunk_contains_sse_done(&prefetched_body_for_report);
let mut usage_stream_telemetry: Option<ExecutionTelemetry> = initial_telemetry.clone();
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry;
let reached_eof = initial_reached_eof;
@@ -2478,6 +2513,8 @@ async fn execute_stream_from_frame_stream(
);
let rewritten_chunk_len =
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream =
stream_chunk_contains_sse_done(&rewritten_chunk);
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_disconnected",
@@ -2490,6 +2527,7 @@ async fn execute_stream_from_frame_stream(
downstream_dropped = true;
break;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
@@ -2604,6 +2642,8 @@ async fn execute_stream_from_frame_stream(
);
let rewritten_chunk_len =
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream =
stream_chunk_contains_sse_done(&rewritten_chunk);
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_flush_disconnected",
@@ -2615,6 +2655,7 @@ async fn execute_stream_from_frame_stream(
);
downstream_dropped = true;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes
.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
@@ -2661,6 +2702,8 @@ async fn execute_stream_from_frame_stream(
);
let flushed_chunk_len =
u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream =
stream_chunk_contains_sse_done(&flushed_chunk);
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_rewrite_flush_disconnected",
@@ -2672,6 +2715,7 @@ async fn execute_stream_from_frame_stream(
);
downstream_dropped = true;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
@@ -2781,6 +2825,18 @@ async fn execute_stream_from_frame_stream(
),
);
if downstream_dropped && client_visible_stream_completed && terminal_failure.is_none() {
debug!(
event_name = "execution_runtime_stream_downstream_closed_after_done",
log_type = "debug",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway treats downstream close after client-visible SSE DONE as completed"
);
downstream_dropped = false;
}
if downstream_dropped {
debug!(
event_name = "execution_runtime_stream_report_skipped",
@@ -2791,6 +2847,12 @@ async fn execute_stream_from_frame_stream(
trace_id = %trace_id_owned,
"gateway skipped stream report because downstream disconnected before completion"
);
let terminal_telemetry = Some(build_terminal_stream_telemetry(
stream_started_at_for_report,
telemetry.as_ref(),
usage_stream_telemetry.as_ref(),
provider_stream_bytes.load(Ordering::Relaxed),
));
let usage_payload = build_stream_usage_payload(
trace_id_owned,
report_kind_owned.unwrap_or_default(),
@@ -2802,7 +2864,7 @@ async fn execute_stream_from_frame_stream(
&buffered_body,
client_body_truncated,
stream_terminal_summary,
telemetry,
terminal_telemetry,
);
record_stream_terminal_usage(
&state_for_report,
@@ -2834,6 +2896,12 @@ async fn execute_stream_from_frame_stream(
if let Some(failure) = terminal_failure {
record_manual_proxy_stream_error(&state_for_report, &plan_for_report).await;
let terminal_telemetry = Some(build_terminal_stream_telemetry(
stream_started_at_for_report,
telemetry.as_ref(),
usage_stream_telemetry.as_ref(),
provider_stream_bytes.load(Ordering::Relaxed),
));
submit_midstream_stream_failure(
&state_for_report,
&trace_id_owned,
@@ -2841,7 +2909,7 @@ async fn execute_stream_from_frame_stream(
direct_stream_finalize_kind_owned.as_deref(),
report_context_owned,
headers_for_report,
telemetry,
terminal_telemetry,
&provider_buffered_body,
candidate_started_unix_secs_for_report,
failure,
@@ -2851,6 +2919,12 @@ async fn execute_stream_from_frame_stream(
}
let should_submit_report = report_kind_owned.is_some();
let terminal_telemetry = Some(build_terminal_stream_telemetry(
stream_started_at_for_report,
telemetry.as_ref(),
usage_stream_telemetry.as_ref(),
provider_stream_bytes.load(Ordering::Relaxed),
));
let usage_payload = build_stream_usage_payload(
trace_id_owned.clone(),
report_kind_owned.unwrap_or_default(),
@@ -2862,7 +2936,7 @@ async fn execute_stream_from_frame_stream(
&buffered_body,
client_body_truncated,
stream_terminal_summary,
telemetry,
terminal_telemetry,
);
apply_local_execution_effect(
&state_for_report,
@@ -3365,6 +3439,7 @@ mod tests {
.await
.expect("first business chunk should arrive");
assert_eq!(first.as_ref(), b"data: {\"id\":\"first\"}\n\n");
tokio::time::sleep(Duration::from_millis(30)).await;
drop(body_stream);
tokio::time::timeout(Duration::from_secs(1), frame_stream_dropped.notified())
@@ -3392,6 +3467,172 @@ mod tests {
candidates[0].error_type.as_deref(),
Some("downstream_disconnect")
);
let stored_usage = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let usage = usage_repository
.find_by_request_id("req-client-drop-cancels-upstream")
.await
.expect("usage should read");
if usage
.as_ref()
.is_some_and(|usage| usage.status == "cancelled")
{
break usage.expect("cancelled usage should exist");
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("usage should be marked cancelled");
assert_eq!(stored_usage.billing_status, "pending");
assert_eq!(stored_usage.status_code, Some(499));
let first_byte_time_ms = stored_usage
.first_byte_time_ms
.expect("cancelled stream should retain first byte time");
let response_time_ms = stored_usage
.response_time_ms
.expect("cancelled stream should record terminal duration");
assert!(
response_time_ms > first_byte_time_ms,
"terminal duration should include time after the first byte"
);
}
#[tokio::test]
async fn image_stream_downstream_close_after_done_is_recorded_success() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let plan = ExecutionPlan {
request_id: "req-image-done-close-success".into(),
candidate_id: Some("cand-image-done-close-success".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://example.com/v1/images/generations".into(),
headers: BTreeMap::from([("accept".into(), "text/event-stream".into())]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-image-2",
"prompt": "draw a small image",
"stream": true
})),
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:image".into(),
model_name: Some("gpt-image-2".into()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let frame_stream = stream! {
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
));
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_1\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\"}}\\n\\nevent: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_1\\\",\\\"model\\\":\\\"gpt-image-2\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":null}}\\n\\n\"}}\n",
));
std::future::pending::<()>().await;
}
.boxed();
let response = execute_stream_from_frame_stream(
&state,
plan,
"trace-image-done-close-success",
&test_decision(),
"openai_chat_stream",
Some("openai_chat_stream_success".to_string()),
Some(json!({
"request_id": "req-image-done-close-success",
"candidate_id": "cand-image-done-close-success",
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "openai:image",
"client_api_format": "openai:chat",
"image_request": {
"size": "1024x1024",
"quality": "medium"
}
})),
crate::clock::current_unix_ms(),
Instant::now(),
frame_stream,
None,
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
let mut body_stream = response.into_body().into_data_stream();
let mut body = Vec::new();
tokio::time::timeout(Duration::from_secs(1), async {
while !String::from_utf8_lossy(&body).contains("data: [DONE]") {
let chunk = body_stream
.next()
.await
.expect("body should yield until done")
.expect("chunk should be ok");
body.extend_from_slice(&chunk);
}
})
.await
.expect("final DONE should arrive");
drop(body_stream);
let candidates = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let candidates = request_candidate_repository
.list_by_request_id("req-image-done-close-success")
.await
.expect("request candidates should read");
if candidates
.first()
.is_some_and(|candidate| candidate.status == RequestCandidateStatus::Success)
{
break candidates;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("candidate should be marked success");
assert_eq!(candidates[0].status_code, Some(200));
let stored_usage = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let usage = usage_repository
.find_by_request_id("req-image-done-close-success")
.await
.expect("usage should read");
if usage
.as_ref()
.is_some_and(|usage| usage.status == "completed")
{
break usage.expect("completed usage should exist");
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("usage should be marked completed");
assert_eq!(stored_usage.status_code, Some(200));
assert!(stored_usage.total_tokens > 0);
}
#[tokio::test]

View File

@@ -107,6 +107,13 @@ pub(crate) async fn build_admin_global_model_routing_payload(
let Some(provider) = providers.get(&model.provider_id) else {
continue;
};
let provider_model_mapping_names =
provider_model_mapping_names_for_routing(model.provider_model_mappings.as_ref());
let key_match_model_names = key_match_model_names_for_routing(
&global_model.name,
&model.provider_model_name,
&provider_model_mapping_names,
);
let mut endpoint_payloads = Vec::new();
let mut active_endpoints = 0usize;
for endpoint in endpoints_by_provider
@@ -132,7 +139,7 @@ pub(crate) async fn build_admin_global_model_routing_payload(
.filter(|key| {
key_allowed_models_match_global_model_for_routing(
key.allowed_models.as_ref(),
&global_model.name,
&key_match_model_names,
&global_model_mappings,
)
})
@@ -313,7 +320,7 @@ pub(crate) async fn build_admin_global_model_routing_payload(
fn key_allowed_models_match_global_model_for_routing(
raw_allowed_models: Option<&serde_json::Value>,
global_model_name: &str,
model_names: &[String],
global_model_mappings: &[String],
) -> bool {
// 兼容 Python 预览逻辑None/[] 视为“不限制”,在链路预览中保留该 Key。
@@ -322,14 +329,16 @@ fn key_allowed_models_match_global_model_for_routing(
return true;
}
if allowed_models
.iter()
.any(|value| value == global_model_name)
{
return true;
}
for allowed_model in &allowed_models {
for allowed_model in allowed_models.iter().map(String::as_str).map(str::trim) {
if allowed_model.is_empty() {
continue;
}
if model_names
.iter()
.any(|model_name| model_name.eq_ignore_ascii_case(allowed_model))
{
return true;
}
for pattern in global_model_mappings {
if matches_model_mapping(pattern, allowed_model) {
return true;
@@ -340,6 +349,54 @@ fn key_allowed_models_match_global_model_for_routing(
false
}
fn provider_model_mapping_names_for_routing(
raw_mappings: Option<&serde_json::Value>,
) -> Vec<String> {
raw_mappings
.and_then(serde_json::Value::as_array)
.map(|items| {
items
.iter()
.filter_map(|item| {
item.as_str()
.or_else(|| item.get("name").and_then(serde_json::Value::as_str))
})
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<Vec<_>>()
})
.unwrap_or_default()
}
fn key_match_model_names_for_routing(
global_model_name: &str,
provider_model_name: &str,
provider_model_mapping_names: &[String],
) -> Vec<String> {
let mut names = Vec::new();
push_unique_model_name(&mut names, global_model_name);
push_unique_model_name(&mut names, provider_model_name);
for mapping_name in provider_model_mapping_names {
push_unique_model_name(&mut names, mapping_name);
}
names
}
fn push_unique_model_name(names: &mut Vec<String>, value: &str) {
let value = value.trim();
if value.is_empty() {
return;
}
if names
.iter()
.any(|existing| existing.eq_ignore_ascii_case(value))
{
return;
}
names.push(value.to_string());
}
pub(crate) async fn build_admin_assign_global_model_to_providers_payload(
state: &AdminAppState<'_>,
global_model_id: &str,

View File

@@ -13,6 +13,7 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::RequestCandidateReadRepository;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -741,3 +742,451 @@ async fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse() {
execution_runtime_handle.abort();
refresh_handle.abort();
}
#[derive(Debug, Clone)]
struct SeenImageBridgeExecutionPlan {
trace_id: String,
client_api_format: String,
provider_api_format: String,
url: String,
plan_stream: bool,
auth_header: String,
chatgpt_web_marker: String,
body_json: serde_json::Value,
}
fn image_bridge_hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn image_bridge_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
None,
Some(serde_json::json!([
"openai:chat",
"openai:responses",
"openai:image"
])),
Some(serde_json::json!(["gpt-image-2"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800_i64),
None,
Some(serde_json::json!([
"openai:chat",
"openai:responses",
"openai:image"
])),
Some(serde_json::json!(["gpt-image-2"])),
)
.expect("auth snapshot should build")
}
fn image_bridge_candidate_row(
prefix: &str,
provider_name: &str,
provider_type: &str,
) -> StoredMinimalCandidateSelectionRow {
let key_auth_type = if provider_type == "chatgpt_web" {
"bearer"
} else {
"api_key"
};
StoredMinimalCandidateSelectionRow {
provider_id: format!("provider-{prefix}"),
provider_name: provider_name.to_string(),
provider_type: provider_type.to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: format!("endpoint-{prefix}"),
endpoint_api_format: "openai:image".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("image".to_string()),
endpoint_is_active: true,
key_id: format!("key-{prefix}"),
key_name: "prod".to_string(),
key_auth_type: key_auth_type.to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:image".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})),
model_id: format!("model-{prefix}"),
global_model_id: format!("global-model-{prefix}"),
global_model_name: "gpt-image-2".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(false),
model_provider_model_name: "gpt-image-2".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-image-2".to_string(),
priority: 1,
api_formats: Some(vec!["openai:image".to_string()]),
endpoint_ids: None,
}]),
model_supports_streaming: Some(false),
model_is_active: true,
model_is_available: true,
}
}
fn image_bridge_provider_catalog_provider(
prefix: &str,
provider_name: &str,
provider_type: &str,
base_url: &str,
) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
format!("provider-{prefix}"),
provider_name.to_string(),
Some(base_url.to_string()),
provider_type.to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn image_bridge_provider_catalog_endpoint(
prefix: &str,
base_url: &str,
) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
format!("endpoint-{prefix}"),
format!("provider-{prefix}"),
"openai:image".to_string(),
Some("openai".to_string()),
Some("image".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
base_url.to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn image_bridge_provider_catalog_key(
prefix: &str,
provider_type: &str,
) -> StoredProviderCatalogKey {
let auth_type = if provider_type == "chatgpt_web" {
"bearer"
} else {
"api_key"
};
StoredProviderCatalogKey::new(
format!("key-{prefix}"),
format!("provider-{prefix}"),
"prod".to_string(),
auth_type.to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:image"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-image-bridge")
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"openai:image": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
async fn start_image_bridge_gateway(
prefix: &str,
provider_name: &str,
provider_type: &str,
base_url: &str,
execution_runtime_url: String,
) -> (
String,
tokio::task::JoinHandle<()>,
String,
Arc<InMemoryRequestCandidateRepository>,
) {
let client_api_key = format!("sk-client-{prefix}");
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(image_bridge_hash_api_key(&client_api_key)),
image_bridge_auth_snapshot(&format!("api-key-{prefix}"), &format!("user-{prefix}")),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
image_bridge_candidate_row(prefix, provider_name, provider_type),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![image_bridge_provider_catalog_provider(
prefix,
provider_name,
provider_type,
base_url,
)],
vec![image_bridge_provider_catalog_endpoint(prefix, base_url)],
vec![image_bridge_provider_catalog_key(prefix, provider_type)],
));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
(
gateway_url,
gateway_handle,
client_api_key,
request_candidate_repository,
)
}
fn capture_image_bridge_execution_plan(
parts: http::request::Parts,
payload: serde_json::Value,
) -> SeenImageBridgeExecutionPlan {
SeenImageBridgeExecutionPlan {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
client_api_format: payload
.get("client_api_format")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
provider_api_format: payload
.get("provider_api_format")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
plan_stream: payload
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(false),
auth_header: payload
.get("headers")
.and_then(|value| value.get("authorization"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
chatgpt_web_marker: payload
.get("headers")
.and_then(|value| value.get("x-aether-chatgpt-web-image"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
body_json: payload
.get("body")
.and_then(|value| value.get("json_body"))
.cloned()
.unwrap_or(serde_json::Value::Null),
}
}
fn image_bridge_execution_runtime(
seen_execution_plan: Arc<Mutex<Option<SeenImageBridgeExecutionPlan>>>,
) -> Router {
Router::new().route(
"/v1/execute/stream",
any(move |request: Request| {
let seen_execution_plan_inner = Arc::clone(&seen_execution_plan);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
.expect("execution runtime payload should parse");
*seen_execution_plan_inner.lock().expect("mutex should lock") =
Some(capture_image_bridge_execution_plan(parts, payload));
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_bridge_123\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\",\\\"output_format\\\":\\\"png\\\"}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_bridge_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-image-2\\\",\\\"status\\\":\\\"completed\\\",\\\"output\\\":[]}}\\n\\n\"}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = http::Response::builder()
.status(StatusCode::OK)
.body(Body::from(frames))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
http::HeaderValue::from_static("application/x-ndjson"),
);
response
}
}),
)
}
#[tokio::test]
async fn gateway_routes_openai_chat_stream_image_intent_to_openai_image_plan_without_streaming_support(
) {
let seen_execution_plan = Arc::new(Mutex::new(None::<SeenImageBridgeExecutionPlan>));
let execution_runtime = image_bridge_execution_runtime(Arc::clone(&seen_execution_plan));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let (gateway_url, gateway_handle, client_api_key, request_candidate_repository) =
start_image_bridge_gateway(
"chat-stream-image-bridge",
"image-provider",
"custom",
"https://images.example.com",
execution_runtime_url,
)
.await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(http::header::AUTHORIZATION, format!("Bearer {client_api_key}"))
.header(TRACE_ID_HEADER, "trace-chat-stream-image-bridge-123")
.body(
r#"{"model":"gpt-image-2","messages":[{"role":"user","content":"Draw a city made of glass"}],"stream":true,"size":"1024x1024"}"#,
)
.send()
.await
.expect("request should succeed");
let status = response.status();
let response_text = response.text().await.expect("body should read");
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-chat-stream-image-bridge-123")
.await
.expect("request candidates should read");
assert_eq!(
status,
StatusCode::OK,
"{response_text}\n{stored_candidates:#?}"
);
assert!(response_text.contains("\"object\":\"chat.completion.chunk\""));
assert!(response_text.contains("![generated image](data:image/png;base64,aGVsbG8=)"));
assert!(response_text.contains("data: [DONE]"));
assert!(!response_text.contains("image_generation.completed"));
let seen_plan = seen_execution_plan
.lock()
.expect("mutex should lock")
.clone()
.expect("execution plan should be captured");
assert_eq!(seen_plan.trace_id, "trace-chat-stream-image-bridge-123");
assert_eq!(seen_plan.client_api_format, "openai:chat");
assert_eq!(seen_plan.provider_api_format, "openai:image");
assert_eq!(seen_plan.url, "https://images.example.com/v1/responses");
assert!(seen_plan.plan_stream);
assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge");
assert_eq!(seen_plan.chatgpt_web_marker, "");
assert_eq!(seen_plan.body_json["model"], "gpt-image-2");
assert_eq!(seen_plan.body_json["stream"], true);
assert_eq!(
seen_plan.body_json["input"][0]["content"],
"Draw a city made of glass"
);
assert_eq!(seen_plan.body_json["tools"][0]["type"], "image_generation");
assert_eq!(seen_plan.body_json["tools"][0]["size"], "1024x1024");
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_routes_openai_responses_stream_image_intent_to_openai_image_plan_without_streaming_support(
) {
let seen_execution_plan = Arc::new(Mutex::new(None::<SeenImageBridgeExecutionPlan>));
let execution_runtime = image_bridge_execution_runtime(Arc::clone(&seen_execution_plan));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let (gateway_url, gateway_handle, client_api_key, _request_candidate_repository) =
start_image_bridge_gateway(
"responses-stream-image-bridge",
"image-provider",
"custom",
"https://images.example.com",
execution_runtime_url,
)
.await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/responses"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(http::header::AUTHORIZATION, format!("Bearer {client_api_key}"))
.header(TRACE_ID_HEADER, "trace-responses-stream-image-bridge-123")
.body(
r#"{"model":"gpt-image-2","input":"Draw a mountain observatory","tools":[{"type":"image_generation","size":"1024x1024"}],"stream":true}"#,
)
.send()
.await
.expect("request should succeed");
let status = response.status();
let response_text = response.text().await.expect("body should read");
assert_eq!(status, StatusCode::OK, "{response_text}");
assert!(response_text.contains("response.output_item.done"));
assert!(response_text.contains("image_generation_call"));
let seen_plan = seen_execution_plan
.lock()
.expect("mutex should lock")
.clone()
.expect("execution plan should be captured");
assert_eq!(
seen_plan.trace_id,
"trace-responses-stream-image-bridge-123"
);
assert_eq!(seen_plan.client_api_format, "openai:responses");
assert_eq!(seen_plan.provider_api_format, "openai:image");
assert_eq!(seen_plan.url, "https://images.example.com/v1/responses");
assert!(seen_plan.plan_stream);
assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge");
assert_eq!(seen_plan.body_json["stream"], true);
assert_eq!(seen_plan.body_json["input"], "Draw a mountain observatory");
assert_eq!(seen_plan.body_json["tools"][0]["type"], "image_generation");
assert_eq!(seen_plan.body_json["tools"][0]["size"], "1024x1024");
gateway_handle.abort();
execution_runtime_handle.abort();
}

View File

@@ -148,6 +148,120 @@ async fn gateway_handles_admin_provider_endpoints_locally_with_trusted_admin_pri
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_counts_fixed_provider_oauth_keys_for_inherited_endpoint_formats() {
let mut codex_provider = sample_provider("provider-codex", "codex", 10);
codex_provider.provider_type = "codex".to_string();
let mut chatgpt_web_provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 20);
chatgpt_web_provider.provider_type = "chatgpt_web".to_string();
let mut codex_key = sample_key(
"key-codex-oauth",
"provider-codex",
"openai:responses:compact",
"oauth-token",
);
codex_key.auth_type = "oauth".to_string();
codex_key.api_formats = Some(json!(["legacy:mismatch"]));
let mut chatgpt_web_key = sample_key(
"key-chatgpt-web-oauth",
"provider-chatgpt-web",
"openai:image",
"oauth-token",
);
chatgpt_web_key.auth_type = "oauth".to_string();
chatgpt_web_key.api_formats = Some(json!(["legacy:mismatch"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![codex_provider, chatgpt_web_provider],
vec![
sample_endpoint(
"endpoint-codex-compact",
"provider-codex",
"openai:responses:compact",
"https://chatgpt.com/backend-api/codex",
),
sample_endpoint(
"endpoint-codex-image",
"provider-codex",
"openai:image",
"https://chatgpt.com/backend-api/codex",
),
sample_endpoint(
"endpoint-chatgpt-web-image",
"provider-chatgpt-web",
"openai:image",
"https://chatgpt.com",
),
],
vec![codex_key, chatgpt_web_key],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let codex_response = client
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-codex/endpoints?skip=0&limit=50"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(codex_response.status(), StatusCode::OK);
let codex_payload: serde_json::Value = codex_response.json().await.expect("json should parse");
let codex_items = codex_payload
.as_array()
.expect("payload should be an array");
for api_format in ["openai:responses:compact", "openai:image"] {
let endpoint = codex_items
.iter()
.find(|item| item["api_format"] == api_format)
.expect("endpoint should exist");
assert_eq!(endpoint["total_keys"], 1);
assert_eq!(endpoint["active_keys"], 1);
}
let chatgpt_web_response = client
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-chatgpt-web/endpoints?skip=0&limit=50"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(chatgpt_web_response.status(), StatusCode::OK);
let chatgpt_web_payload: serde_json::Value = chatgpt_web_response
.json()
.await
.expect("json should parse");
let chatgpt_web_items = chatgpt_web_payload
.as_array()
.expect("payload should be an array");
let chatgpt_web_image = chatgpt_web_items
.iter()
.find(|item| item["api_format"] == "openai:image")
.expect("image endpoint should exist");
assert_eq!(chatgpt_web_image["total_keys"], 1);
assert_eq!(chatgpt_web_image["active_keys"], 1);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpoint() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -195,7 +309,7 @@ async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpo
vec![inherited_key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (_, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")

View File

@@ -936,6 +936,142 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_global_model_routing_counts_image_provider_keys_by_provider_model_name() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/models/global/global-gpt-image/routing",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut image_provider = sample_provider("provider-image", "image", 10);
image_provider.provider_type = "chatgpt_web".to_string();
let grok_provider = sample_provider("provider-grok", "grok2api", 20);
let mut image_key = sample_key(
"key-image-routing",
"provider-image",
"legacy:mismatch",
"sk-image-routing-1234",
);
image_key.name = "image-account".to_string();
image_key.auth_type = "oauth".to_string();
image_key.allowed_models = Some(json!(["gpt-image-2"]));
let mut grok_key = sample_key(
"key-grok-routing",
"provider-grok",
"openai:chat",
"sk-grok-routing-5678",
);
grok_key.name = "all".to_string();
grok_key.allowed_models = Some(json!(["gpt-image-2"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![image_provider, grok_provider],
vec![
sample_endpoint(
"endpoint-image",
"provider-image",
"openai:image",
"https://chatgpt.example",
),
sample_endpoint(
"endpoint-grok-chat",
"provider-grok",
"openai:chat",
"https://grok.example",
),
],
vec![image_key, grok_key],
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![sample_admin_global_model(
"global-gpt-image",
"GPT-Image-2",
"GPT-Image-2",
)])
.with_admin_provider_models(vec![
sample_admin_provider_model(
"model-image-gpt-image",
"provider-image",
"global-gpt-image",
"gpt-image-2",
),
sample_admin_provider_model(
"model-grok-gpt-image",
"provider-grok",
"global-gpt-image",
"gpt-image-2",
),
]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/models/global/global-gpt-image/routing"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["global_model_name"], "GPT-Image-2");
assert_eq!(payload["total_providers"], 2);
assert_eq!(payload["active_providers"], 2);
let providers = payload["providers"].as_array().expect("providers array");
assert_eq!(providers.len(), 2);
let image_endpoints = providers[0]["endpoints"]
.as_array()
.expect("image endpoints array");
assert_eq!(providers[0]["id"], "provider-image");
assert_eq!(image_endpoints[0]["api_format"], "openai:image");
assert_eq!(image_endpoints[0]["total_keys"], 1);
assert_eq!(image_endpoints[0]["active_keys"], 1);
assert_eq!(image_endpoints[0]["keys"][0]["name"], "image-account");
let grok_endpoints = providers[1]["endpoints"]
.as_array()
.expect("grok endpoints array");
assert_eq!(providers[1]["id"], "provider-grok");
assert_eq!(grok_endpoints[0]["api_format"], "openai:chat");
assert_eq!(grok_endpoints[0]["total_keys"], 1);
assert_eq!(grok_endpoints[0]["active_keys"], 1);
assert_eq!(grok_endpoints[0]["keys"][0]["name"], "all");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_creates_admin_global_model_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));

View File

@@ -113,10 +113,7 @@ struct ExpectedUsagePricing {
impl ExpectedUsagePricing {
fn total_tokens(self) -> u64 {
self.input_tokens
.saturating_add(self.output_tokens)
.saturating_add(self.cache_creation_tokens)
.saturating_add(self.cache_read_tokens)
self.input_tokens.saturating_add(self.output_tokens)
}
fn cache_creation_uncategorized_tokens(self) -> u64 {