feat(grok): add runtime image surfaces

This commit is contained in:
mayrain
2026-05-16 19:23:23 +08:00
parent 936e1ae37b
commit 5bf236957e
43 changed files with 7044 additions and 400 deletions

View File

@@ -41,25 +41,29 @@ pub(crate) use crate::ai_serving::{
AiExecutionDecision, AiExecutionPlanPayload, AiStreamAttempt, AiSyncAttempt,
};
pub(crate) use aether_ai_formats::api::{
build_core_error_body_for_client_format, core_error_background_report_kind,
core_error_default_client_api_format, core_success_background_report_kind,
encode_kiro_sse_events, implicit_sync_finalize_report_kind, is_core_error_finalize_kind,
build_core_error_body_for_client_format, convert_standard_chat_response,
core_error_background_report_kind, core_error_default_client_api_format,
core_success_background_report_kind, encode_kiro_sse_events,
implicit_sync_finalize_report_kind, is_core_error_finalize_kind,
normalize_provider_private_report_context, normalize_provider_private_response_value,
provider_private_response_allows_sync_finalize, resolve_claude_stream_spec,
resolve_claude_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
AiControlPlanRequest, ExecutionRuntimeAuthContext, LocalCoreSyncErrorKind,
LocalOpenAiImageSpec, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
AiControlPlanRequest, CanonicalContentPart, CanonicalStreamEvent, CanonicalStreamFrame,
ClaudeClientEmitter, ExecutionRuntimeAuthContext, LocalCoreSyncErrorKind, LocalOpenAiImageSpec,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
LocalStandardSourceMode, LocalStandardSpec, OpenAIChatClientEmitter,
OpenAIResponsesClientEmitter, StreamingStandardTerminalObserver,
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
pub(crate) use aether_ai_formats::protocol::stream::CanonicalUsage as StreamingCanonicalUsage;
pub(crate) fn parse_direct_request_body(
parts: &http::request::Parts,

View File

@@ -66,7 +66,10 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
.await;
let transport_profile = resolve_transport_profile(&resolved.transport);
let transport_profile = resolved
.transport_profile
.clone()
.or_else(|| resolve_transport_profile(&resolved.transport));
let mut extra_fields = serde_json::Map::new();
if let Some(proxy_value) =
build_request_trace_proxy_value(Some(&resolved.transport), proxy.as_ref())
@@ -149,6 +152,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
upstream_url,
provider_request_headers,
provider_request_body,
transport_profile: _,
} = resolved;
Some(build_ai_execution_decision_response(

View File

@@ -1,6 +1,7 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use aether_contracts::ResolvedTransportProfile;
use serde_json::Value;
use crate::ai_serving::planner::common::{
@@ -12,7 +13,8 @@ use crate::ai_serving::transport::antigravity::{
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
};
use crate::ai_serving::transport::{
build_same_format_provider_headers, SameFormatProviderHeadersInput,
build_grok_browser_headers, build_grok_upstream_url, build_same_format_provider_headers,
GrokHeaderInput, SameFormatProviderHeadersInput, GROK_CHAT_PATH,
};
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
use crate::AppState;
@@ -96,6 +98,7 @@ pub(crate) struct LocalSameFormatProviderCandidatePayloadParts {
pub(super) upstream_url: String,
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) provider_request_body: Value,
pub(super) transport_profile: Option<ResolvedTransportProfile>,
}
pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
@@ -246,15 +249,27 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
base_provider_request_body
};
let Some(upstream_url) = super::super::request::build_same_format_upstream_url(
parts,
&prepared.transport,
&prepared.mapped_model,
prepared.provider_api_format.as_str(),
spec,
prepared.upstream_is_stream,
prepared.kiro_auth.as_ref(),
) else {
let is_grok = prepared
.transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok");
let transport_profile =
crate::ai_serving::transport::resolve_transport_profile(&prepared.transport);
let Some(upstream_url) = (if is_grok {
Some(build_grok_upstream_url(&prepared.transport, GROK_CHAT_PATH))
} else {
super::super::request::build_same_format_upstream_url(
parts,
&prepared.transport,
&prepared.mapped_model,
prepared.provider_api_format.as_str(),
spec,
prepared.upstream_is_stream,
prepared.kiro_auth.as_ref(),
)
}) else {
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
state,
input,
@@ -277,7 +292,18 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
.as_ref()
.map(build_antigravity_static_identity_headers)
.unwrap_or_default();
let Some(provider_request_headers) =
let Some(provider_request_headers) = (if is_grok {
build_grok_browser_headers(GrokHeaderInput {
transport: &prepared.transport,
transport_profile: transport_profile.as_ref(),
request_headers: Some(&parts.headers),
content_type: "application/json",
accept: "text/event-stream",
header_rules: prepared.transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
})
} else {
build_same_format_provider_headers(SameFormatProviderHeadersInput {
headers: &parts.headers,
provider_request_body: &provider_request_body,
@@ -294,7 +320,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
.as_ref()
.map(|auth| auth.machine_id.as_str()),
})
else {
}) else {
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
state,
input,
@@ -326,5 +352,6 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
upstream_url,
provider_request_headers,
provider_request_body,
transport_profile,
})
}

View File

@@ -57,7 +57,10 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
.app()
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
.await;
let transport_profile = resolve_transport_profile(&transport);
let transport_profile = resolved
.transport_profile
.clone()
.or_else(|| resolve_transport_profile(&transport));
let mut extra_fields = serde_json::Map::new();
if let Some(proxy_value) = build_request_trace_proxy_value(Some(&transport), proxy.as_ref()) {
extra_fields.insert("proxy".to_string(), proxy_value);

View File

@@ -1,17 +1,19 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use aether_contracts::ResolvedTransportProfile;
use serde_json::Value;
use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, OauthPreparationContext,
};
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::pure::normalize_openai_image_request_with_options;
use crate::ai_serving::transport::{
build_openai_image_headers, build_openai_image_upstream_url,
build_standard_provider_request_headers, openai_image_transport_unsupported_reason,
resolve_openai_image_auth, ProviderOpenAiImageHeadersInput,
StandardProviderRequestHeadersInput,
build_grok_browser_headers, build_grok_upstream_url, 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::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
@@ -21,6 +23,7 @@ use crate::ai_serving::{
normalize_openai_image_request, request_conversion_direct_auth, CandidateFailureDiagnostic,
GatewayProviderTransportSnapshot, PlannerAppState, RequestConversionKind,
};
use crate::image_capabilities::openai_image_normalize_options_for_provider;
use crate::AppState;
use super::support::{
@@ -43,6 +46,7 @@ pub(super) struct LocalOpenAiImageCandidatePayloadParts {
pub(super) provider_request_body: Value,
pub(super) upstream_url: String,
pub(super) input_summary: Value,
pub(super) transport_profile: Option<ResolvedTransportProfile>,
}
pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
@@ -120,8 +124,13 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
let auth_header = prepared_candidate.auth_header;
let auth_value = prepared_candidate.auth_value;
let Some(normalized_request) = normalize_openai_image_request(parts, body_json, body_base64)
else {
let normalized_request = normalize_openai_image_request_with_options(
parts,
body_json,
body_base64,
openai_image_normalize_options_for_provider(&transport.provider.provider_type),
);
let Some(normalized_request) = normalized_request else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
@@ -145,8 +154,16 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web");
let is_grok = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok");
let transport_profile = crate::ai_serving::transport::resolve_transport_profile(transport);
let upstream_url = if is_chatgpt_web {
chatgpt_web_image_internal_url(&transport.endpoint.base_url)
} else if is_grok {
build_grok_upstream_url(transport, GROK_CHAT_PATH)
} else {
build_openai_image_upstream_url(transport, parts.uri.query())
};
@@ -168,7 +185,18 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
);
}
let Some(mut provider_request_headers) =
let Some(mut provider_request_headers) = (if is_grok {
build_grok_browser_headers(GrokHeaderInput {
transport,
transport_profile: transport_profile.as_ref(),
request_headers: Some(&parts.headers),
content_type: "application/json",
accept: "*/*",
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
})
} else {
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &auth_header,
@@ -177,7 +205,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
provider_request_body: &provider_request_body,
original_request_body: body_json,
})
else {
}) else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
@@ -197,6 +225,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
};
if is_chatgpt_web {
provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else if is_grok {
} else {
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
@@ -222,7 +251,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
.unwrap_or_default()
.to_string();
let input_summary = if is_chatgpt_web {
let input_summary = if is_chatgpt_web || is_grok {
provider_request_body.clone()
} else {
normalized_request.summary_json
@@ -239,6 +268,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
provider_request_body,
upstream_url,
input_summary,
transport_profile,
})
}
@@ -423,6 +453,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
provider_request_body: converted.body_json,
upstream_url,
input_summary: converted.summary_json,
transport_profile: None,
})
}

View File

@@ -144,7 +144,10 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
),
&resolved.transport,
);
let transport_profile = resolve_transport_profile(&resolved.transport);
let transport_profile = resolved
.transport_profile
.clone()
.or_else(|| resolve_transport_profile(&resolved.transport));
let timeouts = resolve_transport_execution_timeouts(&resolved.transport);
let super::request::LocalStandardCandidatePayloadParts {
auth_header,
@@ -157,6 +160,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
upstream_is_stream,
envelope_name: _,
transport,
transport_profile: _,
} = resolved;
Some(build_ai_execution_decision_response(

View File

@@ -1,6 +1,7 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use aether_contracts::ResolvedTransportProfile;
use serde_json::Value;
use crate::ai_serving::planner::candidate_preparation::{
@@ -21,10 +22,11 @@ use crate::ai_serving::transport::kiro::{
KIRO_ENVELOPE_NAME,
};
use crate::ai_serving::transport::{
build_kiro_cross_format_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_standard_provider_request_headers,
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput,
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, openai_image_transport_unsupported_reason,
resolve_grok_session_auth, resolve_openai_image_auth, GrokHeaderInput,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
};
use crate::ai_serving::{
build_openai_image_request_body_from_gemini_image_request, gemini_request_is_image_generation,
@@ -49,6 +51,14 @@ pub(crate) struct LocalStandardCandidatePayloadParts {
pub(super) upstream_is_stream: bool,
pub(super) envelope_name: Option<&'static str>,
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) transport_profile: Option<ResolvedTransportProfile>,
}
fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
matches!(
crate::ai_serving::normalize_api_format_alias(provider_api_format).as_str(),
"openai:chat" | "openai:responses" | "openai:responses:compact" | "claude:messages"
)
}
pub(crate) async fn resolve_local_standard_candidate_payload_parts(
@@ -64,7 +74,13 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
let planner_state = crate::ai_serving::PlannerAppState::new(state);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let transport_profile = crate::ai_serving::transport::resolve_transport_profile(transport);
let provider_api_format = attempt.eligible.provider_api_format.as_str();
let is_grok = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok");
if spec_metadata.api_format == "gemini:generate_content"
&& provider_api_format == "openai:image"
&& gemini_request_is_image_generation(body_json)
@@ -75,6 +91,104 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
.await;
}
let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
if is_grok && is_grok_text_provider_api_format(provider_api_format) {
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
transport,
candidate,
resolve_grok_session_auth(transport),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "standard_family_grok_text_request",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
};
let mut provider_request_body = body_json.clone();
if let Some(object) = provider_request_body.as_object_mut() {
object.insert(
"model".to_string(),
serde_json::Value::String(prepared_candidate.mapped_model.clone()),
);
}
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 force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
enforce_provider_body_stream_policy(
&mut provider_request_body,
provider_api_format,
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
let upstream_url = build_grok_upstream_url(transport, GROK_CHAT_PATH);
let Some(provider_request_headers) = build_grok_browser_headers(GrokHeaderInput {
transport,
transport_profile: transport_profile.as_ref(),
request_headers: Some(&parts.headers),
content_type: "application/json",
accept: "text/event-stream",
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
}) else {
mark_skipped_local_standard_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
spec_metadata.api_format,
provider_api_format,
"grok_standard_family_headers",
),
)
.await;
return None;
};
return Some(LocalStandardCandidatePayloadParts {
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,
upstream_is_stream,
envelope_name: None,
transport: Arc::clone(transport),
transport_profile,
});
}
let Some(conversion_kind) =
crate::ai_serving::request_conversion_kind(spec_metadata.api_format, provider_api_format)
else {
@@ -361,6 +475,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
upstream_is_stream,
envelope_name: None,
transport: Arc::clone(transport),
transport_profile: None,
})
}
@@ -496,6 +611,7 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
upstream_is_stream,
envelope_name: None,
transport: Arc::clone(transport),
transport_profile: None,
})
}
@@ -614,5 +730,6 @@ async fn build_kiro_cross_format_payload_parts(
upstream_is_stream,
envelope_name: Some(KIRO_ENVELOPE_NAME),
transport: Arc::clone(transport),
transport_profile: None,
})
}

View File

@@ -67,7 +67,10 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
.await;
let transport_profile = resolve_transport_profile(&resolved.transport);
let transport_profile = resolved
.transport_profile
.clone()
.or_else(|| resolve_transport_profile(&resolved.transport));
let timeouts = resolve_transport_execution_timeouts(&resolved.transport);
let mut extra_fields = serde_json::Map::new();
if let Some(proxy_value) =
@@ -99,6 +102,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
envelope_name,
transport,
request_redacted,
transport_profile: _,
} = resolved;
let original_request_body_json = if request_redacted {
Some(&provider_request_body)

View File

@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_contracts::ResolvedTransportProfile;
use serde_json::Value;
use crate::ai_serving::planner::candidate_preparation::{
@@ -27,8 +28,9 @@ use crate::ai_serving::transport::kiro::{
};
use crate::ai_serving::transport::local_openai_chat_transport_unsupported_reason;
use crate::ai_serving::transport::{
build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
StandardProviderRequestHeadersInput,
build_grok_browser_headers, build_grok_upstream_url, build_kiro_cross_format_upstream_url,
build_standard_provider_request_headers, GrokHeaderInput, StandardProviderRequestHeadersInput,
GROK_CHAT_PATH,
};
use crate::ai_serving::{
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
@@ -64,6 +66,14 @@ pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
pub(super) envelope_name: Option<&'static str>,
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) request_redacted: bool,
pub(super) transport_profile: Option<ResolvedTransportProfile>,
}
fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
matches!(
crate::ai_serving::normalize_api_format_alias(provider_api_format).as_str(),
"openai:chat" | "openai:responses" | "openai:responses:compact" | "claude:messages"
)
}
fn request_identity_response_encoding_when_redacted(
@@ -177,6 +187,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
let candidate = &eligible.candidate;
let provider_api_format = eligible.provider_api_format.as_str();
let transport = &eligible.transport;
let transport_profile = crate::ai_serving::transport::resolve_transport_profile(transport);
let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let enable_model_directives =
@@ -190,6 +201,129 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
resolve_provider_chat_request_redaction(state, parts, body_json, input, candidate_id)
.await?;
let body_json = redaction.body_json.as_ref();
let is_grok = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok");
if is_grok && is_grok_text_provider_api_format(provider_api_format) {
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
transport,
candidate,
crate::ai_serving::transport::resolve_grok_session_auth(transport),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "openai_chat_same_format",
},
)
.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 Some(provider_request_body) = build_local_openai_chat_request_body(
body_json,
&prepared_candidate.mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.endpoint.body_rules.as_ref(),
&parts.headers,
enable_model_directives,
) 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);
};
let upstream_url = build_grok_upstream_url(transport, GROK_CHAT_PATH);
let Some(mut provider_request_headers) = build_grok_browser_headers(GrokHeaderInput {
transport,
transport_profile: transport_profile.as_ref(),
request_headers: Some(&parts.headers),
content_type: "application/json",
accept: "text/event-stream",
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,
"grok_openai_chat_headers",
),
)
.await;
return Ok(None);
};
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", provider_api_format);
let resolved_report_kind =
if decision_kind == OPENAI_CHAT_STREAM_PLAN_KIND || !upstream_is_stream {
report_kind.to_string()
} else {
"openai_chat_sync_finalize".to_string()
};
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
redaction.redacted,
);
return Ok(Some(LocalOpenAiChatCandidatePayloadParts {
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: resolved_report_kind,
envelope_name: None,
transport: Arc::clone(transport),
request_redacted: redaction.redacted,
transport_profile,
}));
}
if provider_api_format == "openai:chat" {
if let Some(skip_reason) = local_openai_chat_transport_unsupported_reason(transport) {
@@ -351,6 +485,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
envelope_name: None,
transport: Arc::clone(transport),
request_redacted: redaction.redacted,
transport_profile,
}));
};
@@ -639,6 +774,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
envelope_name: None,
transport: Arc::clone(transport),
request_redacted: redaction.redacted,
transport_profile: None,
}))
}
@@ -775,6 +911,7 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
envelope_name: Some(KIRO_ENVELOPE_NAME),
transport: Arc::clone(transport),
request_redacted,
transport_profile: None,
})
}

View File

@@ -63,7 +63,10 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
.await;
let transport_profile = resolve_transport_profile(&resolved.transport);
let transport_profile = resolved
.transport_profile
.clone()
.or_else(|| resolve_transport_profile(&resolved.transport));
let timeouts = resolve_transport_execution_timeouts(&resolved.transport);
let mut extra_fields = serde_json::Map::new();
if let Some(proxy_value) =
@@ -170,6 +173,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
envelope_name: _,
upstream_is_stream,
transport,
transport_profile: _,
} = resolved;
Some(build_ai_execution_decision_response(

View File

@@ -1,6 +1,7 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use aether_contracts::ResolvedTransportProfile;
use serde_json::Value;
use tracing::debug;
@@ -35,8 +36,10 @@ use crate::ai_serving::transport::kiro::{
KiroRequestAuth, KIRO_ENVELOPE_NAME,
};
use crate::ai_serving::transport::{
build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
local_standard_transport_unsupported_reason_with_network, StandardProviderRequestHeadersInput,
build_grok_browser_headers, build_grok_upstream_url, build_kiro_cross_format_upstream_url,
build_standard_provider_request_headers,
local_standard_transport_unsupported_reason_with_network, GrokHeaderInput,
StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
};
use crate::ai_serving::{
ai_local_execution_contract_for_formats, request_conversion_direct_auth,
@@ -56,6 +59,13 @@ use super::LocalOpenAiResponsesSpec;
const ANTIGRAVITY_ENVELOPE_NAME: &str = "antigravity:v1internal";
fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
matches!(
crate::ai_serving::normalize_api_format_alias(provider_api_format).as_str(),
"openai:chat" | "openai:responses" | "openai:responses:compact" | "claude:messages"
)
}
pub(crate) struct LocalOpenAiResponsesCandidatePayloadParts {
pub(super) auth_header: String,
pub(super) auth_value: String,
@@ -70,6 +80,7 @@ pub(crate) struct LocalOpenAiResponsesCandidatePayloadParts {
pub(super) envelope_name: Option<&'static str>,
pub(super) upstream_is_stream: bool,
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) transport_profile: Option<ResolvedTransportProfile>,
}
#[allow(clippy::too_many_arguments)]
@@ -90,12 +101,22 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
let candidate = &eligible.candidate;
let provider_api_format = eligible.provider_api_format.as_str();
let transport = &eligible.transport;
let transport_profile = crate::ai_serving::transport::resolve_transport_profile(transport);
let is_antigravity = is_antigravity_provider_transport(transport);
let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
let is_grok = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok");
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 same_format && is_kiro_claude_cli {
let transport_unsupported_reason = if is_grok
&& is_grok_text_provider_api_format(provider_api_format)
{
None
} else if same_format && is_kiro_claude_cli {
local_kiro_request_transport_unsupported_reason_with_network(transport)
} else if same_format {
local_standard_transport_unsupported_reason_with_network(transport, provider_api_format)
@@ -154,7 +175,9 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
None
};
let direct_auth = if kiro_auth.is_some() {
let direct_auth = if is_grok && is_grok_text_provider_api_format(provider_api_format) {
crate::ai_serving::transport::resolve_grok_session_auth(transport)
} else if kiro_auth.is_some() {
None
} else if same_format {
match crate::ai_serving::normalize_api_format_alias(provider_api_format).as_str() {
@@ -236,42 +259,57 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
);
let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let Some(mut base_provider_request_body) = (if needs_bidirectional_conversion {
build_cross_format_openai_responses_request_body(
body_json,
&mapped_model,
spec_metadata.api_format,
provider_api_format,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
if is_kiro_claude_cli {
None
} else {
transport.endpoint.body_rules.as_ref()
},
Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives,
)
} else {
build_local_openai_responses_request_body(
body_json,
&mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
provider_api_format,
if is_kiro_claude_cli {
None
} else {
transport.endpoint.body_rules.as_ref()
},
Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives,
)
}) else {
let Some(mut base_provider_request_body) =
(if is_grok && is_grok_text_provider_api_format(provider_api_format) {
build_local_openai_responses_request_body(
body_json,
&mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
spec_metadata.api_format,
transport.endpoint.body_rules.as_ref(),
Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives,
)
} else if needs_bidirectional_conversion {
build_cross_format_openai_responses_request_body(
body_json,
&mapped_model,
spec_metadata.api_format,
provider_api_format,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
if is_kiro_claude_cli {
None
} else {
transport.endpoint.body_rules.as_ref()
},
Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives,
)
} else {
build_local_openai_responses_request_body(
body_json,
&mapped_model,
upstream_is_stream,
force_body_stream_field,
transport.provider.provider_type.as_str(),
provider_api_format,
if is_kiro_claude_cli {
None
} else {
transport.endpoint.body_rules.as_ref()
},
Some(input.auth_context.api_key_id.as_str()),
&parts.headers,
enable_model_directives,
)
})
else {
mark_skipped_local_openai_responses_candidate_with_extra_data(
state,
input,
@@ -390,7 +428,9 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
.await;
}
let Some(upstream_url) = (if needs_bidirectional_conversion {
let Some(upstream_url) = (if is_grok && is_grok_text_provider_api_format(provider_api_format) {
Some(build_grok_upstream_url(transport, GROK_CHAT_PATH))
} else if needs_bidirectional_conversion {
build_cross_format_openai_responses_upstream_url(
parts,
transport,
@@ -427,48 +467,86 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
.as_ref()
.map(build_antigravity_static_identity_headers)
.unwrap_or_default();
let Some(resolved_headers) =
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
let resolved_headers = if is_grok && is_grok_text_provider_api_format(provider_api_format) {
let Some(headers) = build_grok_browser_headers(GrokHeaderInput {
transport,
provider_api_format,
same_format,
headers: &parts.headers,
auth_header: &auth_header,
auth_value: &auth_value,
extra_headers: &extra_headers,
transport_profile: transport_profile.as_ref(),
request_headers: Some(&parts.headers),
content_type: "application/json",
accept: "text/event-stream",
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
upstream_is_stream,
})
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,
}) 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,
"grok_openai_responses_headers",
),
)
.await;
return None;
};
crate::ai_serving::transport::StandardProviderRequestHeaders {
headers,
auth_header: auth_header.clone(),
auth_value: auth_value.clone(),
}
} else {
let Some(resolved_headers) =
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
transport,
provider_api_format,
"openai_responses_headers",
),
)
.await;
return None;
same_format,
headers: &parts.headers,
auth_header: &auth_header,
auth_value: &auth_value,
extra_headers: &extra_headers,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
upstream_is_stream,
})
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_headers",
),
)
.await;
return None;
};
resolved_headers
};
let mut provider_request_headers = resolved_headers.headers;
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(),
);
if !is_grok {
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);
@@ -516,6 +594,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
},
upstream_is_stream,
transport: Arc::clone(transport),
transport_profile,
})
}
@@ -666,5 +745,6 @@ async fn build_kiro_openai_responses_payload_parts(
envelope_name: Some(KIRO_ENVELOPE_NAME),
upstream_is_stream,
transport: Arc::clone(transport),
transport_profile: None,
})
}

View File

@@ -61,6 +61,7 @@ pub(crate) use aether_ai_formats::api::{
maybe_build_standard_sync_finalize_product_from_normalized_payload, model_directive_base_model,
normalize_api_format_alias, normalize_claude_request_to_openai_chat_request,
normalize_gemini_request_to_openai_chat_request, normalize_openai_image_request,
normalize_openai_image_request_with_options,
normalize_openai_responses_request_to_openai_chat_request,
normalize_provider_private_report_context, normalize_provider_private_response_value,
normalize_standard_request_to_openai_chat_request, openai_image_operation_from_path,
@@ -97,10 +98,10 @@ pub(crate) use aether_ai_formats::api::{
LocalStandardSourceMode, LocalStandardSpec, LocalSyncReportParts, LocalVideoCreateFamily,
LocalVideoCreateSpec, NormalizedOpenAiImageRequest, OpenAIChatClientEmitter,
OpenAIChatProviderState, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState,
OpenAiImageOperation, OpenAiImageRequestForGemini, OpenAiImageResponseFormat,
OpenAiImageStreamState, OpenAiImageSyncFinalizeProduct, ProviderAdaptationDescriptor,
ProviderAdaptationSurface, ProviderPrivateStreamNormalizer, RequestConversionKind,
StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
OpenAiImageNormalizeOptions, OpenAiImageOperation, OpenAiImageRequestForGemini,
OpenAiImageResponseFormat, OpenAiImageStreamState, OpenAiImageSyncFinalizeProduct,
ProviderAdaptationDescriptor, ProviderAdaptationSurface, ProviderPrivateStreamNormalizer,
RequestConversionKind, StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
StreamingStandardFormatMatrix, SyncChatResponseConversionKind, SyncCliResponseConversionKind,
SyncToStreamBridgeOutcome, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME, CLAUDE_CHAT_STREAM_PLAN_KIND,
CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND,

View File

@@ -14,6 +14,10 @@ pub(crate) mod kiro {
pub(crate) use aether_provider_transport::kiro::*;
}
pub(crate) mod grok {
pub(crate) use aether_provider_transport::grok::*;
}
pub(crate) mod oauth_refresh {
pub(crate) use aether_provider_transport::oauth_refresh::*;
}
@@ -54,6 +58,7 @@ pub(crate) use aether_provider_transport::{
body_rules_are_locally_supported, body_rules_handle_path, body_rules_have_enabled_rules,
build_cross_format_openai_chat_upstream_url, build_cross_format_openai_responses_upstream_url,
build_gemini_files_headers, build_gemini_files_request_body, build_gemini_files_upstream_url,
build_grok_app_chat_body, build_grok_browser_headers, build_grok_upstream_url,
build_kiro_cross_format_upstream_url, build_local_openai_chat_upstream_url,
build_local_openai_responses_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_passthrough_headers, build_request_trace_proxy_value,
@@ -72,21 +77,23 @@ pub(crate) use aether_provider_transport::{
openai_image_transport_unsupported_reason, request_conversion_direct_auth,
request_conversion_enabled_for_transport, request_conversion_transport_supported,
request_conversion_transport_unsupported_reason, request_pair_allowed_for_transport,
resolve_gemini_files_auth, resolve_openai_image_auth, resolve_same_format_provider_direct_auth,
resolve_transport_execution_timeouts, resolve_transport_profile,
resolve_transport_proxy_snapshot, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_video_create_auth, same_format_provider_transport_supported,
same_format_provider_transport_unsupported_reason, should_skip_upstream_passthrough_header,
should_try_same_format_provider_oauth_auth, supports_local_gemini_transport_with_network,
resolve_gemini_files_auth, resolve_grok_session_auth, resolve_openai_image_auth,
resolve_same_format_provider_direct_auth, resolve_transport_execution_timeouts,
resolve_transport_profile, resolve_transport_proxy_snapshot,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_video_create_auth,
same_format_provider_transport_supported, same_format_provider_transport_unsupported_reason,
should_skip_upstream_passthrough_header, should_try_same_format_provider_oauth_auth,
supports_local_gemini_transport_with_network,
supports_local_generic_oauth_request_auth_resolution,
supports_local_oauth_request_auth_resolution, transport_proxy_is_locally_supported,
video_create_transport_unsupported_reason, CandidateTransportPolicyFacts,
GatewayProviderTransportSnapshot, GeminiFilesHeadersInput, GeminiFilesRequestBodyError,
GeminiFilesRequestBodyParts, LocalResolvedOAuthRequestAuth, ProviderOpenAiImageHeadersInput,
ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput, SameFormatProviderFamily,
SameFormatProviderHeadersInput, SameFormatProviderRequestBehavior,
GeminiFilesRequestBodyParts, GrokHeaderInput, LocalResolvedOAuthRequestAuth,
ProviderOpenAiImageHeadersInput, ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput,
SameFormatProviderFamily, SameFormatProviderHeadersInput, SameFormatProviderRequestBehavior,
SameFormatProviderRequestBehaviorParams, SameFormatProviderRequestBodyInput,
SameFormatProviderUpstreamUrlParams, StandardPlanFallbackAcceptPolicy,
StandardPlanFallbackHeadersInput, StandardProviderRequestHeaders,
StandardProviderRequestHeadersInput, TransportRequestUrlParams,
StandardProviderRequestHeadersInput, TransportRequestUrlParams, GROK_CHAT_PATH,
GROK_INTERNAL_HEADER, GROK_RATE_LIMITS_PATH,
};

View File

@@ -17,7 +17,6 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[
"/v1/responses/compact",
"/v1/images/generations",
"/v1/images/edits",
"/v1/images/variations",
];
const AI_ANY_ROUTE_PATTERNS: &[&str] = &[

View File

@@ -115,7 +115,6 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[
"/v1/rerank",
"/v1/images/generations",
"/v1/images/edits",
"/v1/images/variations",
"/v1/messages",
"/v1/messages/count_tokens",
"/v1/responses",

View File

@@ -55,7 +55,7 @@ pub(super) fn classify_ai_public_route(
} else if method == http::Method::POST
&& matches!(
normalized_path,
"/v1/images/generations" | "/v1/images/edits" | "/v1/images/variations"
"/v1/images/generations" | "/v1/images/edits"
)
{
Some(classified(

View File

@@ -80,6 +80,28 @@ fn classifies_openai_chat_and_responses_separately_from_embedding() {
assert_ne!(responses.route_kind.as_deref(), Some("embedding"));
}
#[test]
fn classifies_openai_image_generation_and_edit_but_not_variation() {
let headers = headers(&[("authorization", "Bearer sk-test")]);
for path in ["/v1/images/generations", "/v1/images/edits"] {
let uri: Uri = path.parse().expect("uri should parse");
let decision = classify_control_route(&http::Method::POST, &uri, &headers)
.expect("image route should classify");
assert_eq!(decision.route_family.as_deref(), Some("openai"));
assert_eq!(decision.route_kind.as_deref(), Some("image"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("openai:image")
);
assert!(decision.is_execution_runtime_candidate());
}
let variation_uri: Uri = "/v1/images/variations".parse().expect("uri should parse");
assert!(classify_control_route(&http::Method::POST, &variation_uri, &headers).is_none());
}
#[test]
fn classifies_models_list_as_claude_when_headers_match() {
let headers = headers(&[

View File

@@ -3,9 +3,10 @@ use std::io::Error as IoError;
use std::time::Instant;
use aether_contracts::{
ExecutionPlan, ExecutionResult, ExecutionTelemetry, RequestBody, ResponseBody, StreamFrame,
StreamFramePayload, StreamFrameType, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
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,
};
use axum::body::Bytes;
use base64::Engine as _;
@@ -30,6 +31,7 @@ const CHATGPT_WEB_CLIENT_VERSION: &str = "prod-be885abbfcfe7b1f511e88b3003d9ee44
const CHATGPT_WEB_BUILD_NUMBER: &str = "5955942";
const CHATGPT_WEB_SEC_CH_UA: &str =
r#""Microsoft Edge";v="143", "Chromium";v="143", "Not A(Brand";v="24""#;
const CHATGPT_WEB_BROWSER_PROFILE: &str = "chrome143";
pub(crate) struct ChatGptWebImageStream {
pub(crate) frame_stream: BoxStream<'static, Result<Bytes, IoError>>,
@@ -921,7 +923,7 @@ async fn execute_subrequest(
provider_api_format: plan.provider_api_format.clone(),
model_name: plan.model_name.clone(),
proxy: plan.proxy.clone(),
transport_profile: plan.transport_profile.clone(),
transport_profile: chatgpt_web_image_transport_profile(plan),
timeouts: plan.timeouts.clone(),
};
DirectSyncExecutionRuntime::new()
@@ -929,6 +931,34 @@ async fn execute_subrequest(
.await
}
fn chatgpt_web_image_transport_profile(plan: &ExecutionPlan) -> Option<ResolvedTransportProfile> {
match plan.transport_profile.as_ref() {
Some(profile)
if profile
.backend
.trim()
.eq_ignore_ascii_case(TRANSPORT_BACKEND_BROWSER_WREQ) =>
{
Some(profile.clone())
}
_ => Some(default_chatgpt_web_image_transport_profile()),
}
}
fn default_chatgpt_web_image_transport_profile() -> ResolvedTransportProfile {
ResolvedTransportProfile {
profile_id: CHATGPT_WEB_BROWSER_PROFILE.to_string(),
backend: TRANSPORT_BACKEND_BROWSER_WREQ.to_string(),
http_mode: TRANSPORT_HTTP_MODE_AUTO.to_string(),
pool_scope: TRANSPORT_POOL_SCOPE_KEY.to_string(),
header_fingerprint: None,
extra: Some(json!({
"browser_profile": CHATGPT_WEB_BROWSER_PROFILE,
"source": "chatgpt_web_image_default",
})),
}
}
fn web_base_headers(fp: &WebFingerprint, token: &str, path: &str) -> BTreeMap<String, String> {
let mut headers = BTreeMap::from([
("user-agent".to_string(), fp.user_agent.to_string()),
@@ -1962,6 +1992,30 @@ mod tests {
}
}
#[test]
fn chatgpt_web_image_subrequests_default_to_browser_wreq_transport() {
let plan = sample_plan(
CHATGPT_WEB_DEFAULT_BASE_URL,
json!({"prompt": "draw a small test image"}),
false,
);
let profile = chatgpt_web_image_transport_profile(&plan).expect("transport profile");
assert_eq!(profile.backend, TRANSPORT_BACKEND_BROWSER_WREQ);
assert_eq!(profile.profile_id, CHATGPT_WEB_BROWSER_PROFILE);
assert_eq!(profile.http_mode, TRANSPORT_HTTP_MODE_AUTO);
assert_eq!(profile.pool_scope, TRANSPORT_POOL_SCOPE_KEY);
assert_eq!(
profile
.extra
.as_ref()
.and_then(|value| value.get("source"))
.and_then(Value::as_str),
Some("chatgpt_web_image_default")
);
}
async fn start_mock_chatgpt_web() -> (String, tokio::task::JoinHandle<()>) {
let app = Router::new().fallback(any(|request: Request| async move {
let path = request.uri().path().to_string();

File diff suppressed because it is too large Load Diff

View File

@@ -6,6 +6,7 @@ use serde_json::{Map, Value};
mod chatgpt_web_image;
mod constants;
mod fallback;
mod grok;
mod kiro_web_search;
pub(crate) mod ndjson;
mod oauth_retry;
@@ -54,6 +55,7 @@ pub(crate) use sync::{
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessBuild,
LocalVideoSyncSuccessOutcome,
};
pub(crate) use transport::execute_sync_plan_with_report_context as execute_execution_runtime_sync_plan_with_report_context;
pub(crate) use transport::{
execute_sync_plan as execute_execution_runtime_sync_plan, DirectSyncExecutionRuntime,
DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,

View File

@@ -370,6 +370,8 @@ impl IntoResponse for ExecutionRuntimeAppError {
) => StatusCode::BAD_REQUEST,
ExecutionRuntimeServerError::Transport(
ExecutionRuntimeTransportError::ClientBuild(_)
| ExecutionRuntimeTransportError::BrowserClientBuild(_)
| ExecutionRuntimeTransportError::BrowserBody(_)
| ExecutionRuntimeTransportError::UpstreamRequest(_)
| ExecutionRuntimeTransportError::RelayError(_)
| ExecutionRuntimeTransportError::InvalidJson(_),

View File

@@ -60,6 +60,7 @@ use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER};
use crate::control::GatewayControlDecision;
use crate::execution_runtime::build_direct_execution_frame_stream;
use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_stream;
use crate::execution_runtime::grok::maybe_execute_grok_stream;
use crate::execution_runtime::kiro_web_search::maybe_execute_kiro_web_search_stream;
use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry;
#[cfg(test)]
@@ -525,6 +526,58 @@ pub(crate) async fn execute_execution_runtime_stream(
key_id.as_str(),
)
.await;
match maybe_execute_grok_stream(&plan, report_context.as_ref()).await {
Ok(Some(grok_stream)) => {
return execute_stream_from_frame_stream(
state,
plan,
trace_id,
decision,
plan_kind,
report_kind,
grok_stream.report_context.or(report_context),
candidate_started_unix_secs,
stream_started_at,
grok_stream.frame_stream,
provider_pool_in_flight_guard.take(),
)
.await;
}
Ok(None) => {}
Err(err) => {
info!(
event_name = "grok_execution_unavailable",
log_type = "ops",
trace_id = %trace_id,
request_id = %plan_request_id_for_log,
candidate_id = ?plan.candidate_id,
provider_name = provider_name.as_str(),
endpoint_id = %endpoint_id,
key_id = %key_id,
model_name = model_name.as_str(),
candidate_index = candidate_index.as_str(),
error = %err,
"gateway Grok stream execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("grok_execution_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
}
match maybe_execute_kiro_web_search_stream(state, &plan, report_context.as_ref()).await {
Ok(Some(kiro_web_search)) => {
return execute_stream_from_frame_stream(

View File

@@ -18,7 +18,9 @@ use crate::ai_serving::api::{
normalize_provider_private_report_context, StreamingStandardTerminalObserver,
};
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
use crate::execution_runtime::transport::DirectUpstreamResponse;
use crate::execution_runtime::transport::{
format_wreq_upstream_request_error, DirectUpstreamResponse,
};
use crate::execution_runtime::DirectUpstreamStreamExecution;
use crate::GatewayError;
@@ -235,6 +237,62 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
}
DirectUpstreamResponse::BrowserWreq(response) => {
let mut bytes_stream = response.bytes_stream();
while let Some(item) = bytes_stream.next().await {
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
if !first_chunk_telemetry_emitted {
match encode_telemetry_frame(ttfb_ms, ttfb_ms, upstream_bytes) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
observe_stream_chunk(
&mut stream_terminal_observer,
&normalized_observer_context,
private_stream_normalizer.as_mut(),
&mut observer_buffered,
chunk.as_ref(),
);
match encode_data_frame(&chunk) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
}
Err(err) => {
let message = format_wreq_upstream_request_error(&err);
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error = %message,
"upstream body stream read error"
);
match encode_error_frame(status_code, message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
return;
}
}
break;
}
}
}
}
DirectUpstreamResponse::LocalTunnel(mut response) => loop {
match response.next_chunk().await {
Ok(Some(chunk)) => {
@@ -454,6 +512,35 @@ async fn buffer_non_sse_upstream_body(
}
}
}
DirectUpstreamResponse::BrowserWreq(response) => {
let mut bytes_stream = response.bytes_stream();
while let Some(item) = bytes_stream.next().await {
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
upstream_bytes += chunk.len() as u64;
body_bytes.extend_from_slice(&chunk);
}
Err(err) => {
let message = format_wreq_upstream_request_error(&err);
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
upstream_bytes,
error = %message,
"upstream body stream read error"
);
return Err(BufferedUpstreamBodyError {
message,
ttfb_ms,
upstream_bytes,
});
}
}
}
}
DirectUpstreamResponse::LocalTunnel(mut response) => loop {
match response.next_chunk().await {
Ok(Some(chunk)) => {

View File

@@ -36,14 +36,15 @@ use crate::api::response::{
use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync;
use crate::execution_runtime::grok::maybe_execute_grok_sync;
use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry;
#[cfg(test)]
use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_runtime;
use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
use crate::execution_runtime::transport::{
build_request_body, collect_response_headers, decode_response_body_bytes,
response_body_is_json, send_request, DirectSyncExecutionRuntime,
ExecutionRuntimeTransportError,
format_upstream_request_error, format_wreq_upstream_request_error, response_body_is_json,
send_request, DirectHttpResponse, DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
};
use crate::execution_runtime::{
analyze_local_candidate_failover_sync, apply_endpoint_response_header_rules,
@@ -682,23 +683,46 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
.await
.map_err(SyncExecutionFailure::from_transport)?;
let ttfb_ms = started_at.elapsed().as_millis() as u64;
let status_code = response.status().as_u16();
let headers = collect_response_headers(response.headers());
let status_code = response.status_code();
let headers = response.headers();
progress.record_response_started(status_code, ttfb_ms).await;
let mut upstream_stream = response.bytes_stream();
let mut body_bytes = Vec::new();
while let Some(chunk) = upstream_stream.next().await {
let chunk = chunk.map_err(|err| {
SyncExecutionFailure::from_transport(ExecutionRuntimeTransportError::UpstreamRequest(
crate::execution_runtime::transport::format_upstream_request_error(&err),
))
})?;
let elapsed_ms = started_at.elapsed().as_millis() as u64;
progress
.observe_chunk(&chunk, status_code, elapsed_ms)
.await;
body_bytes.extend_from_slice(&chunk);
match response {
DirectHttpResponse::Reqwest(response) => {
let mut upstream_stream = response.bytes_stream();
while let Some(chunk) = upstream_stream.next().await {
let chunk = chunk.map_err(|err| {
SyncExecutionFailure::from_transport(
ExecutionRuntimeTransportError::UpstreamRequest(
format_upstream_request_error(&err),
),
)
})?;
let elapsed_ms = started_at.elapsed().as_millis() as u64;
progress
.observe_chunk(&chunk, status_code, elapsed_ms)
.await;
body_bytes.extend_from_slice(&chunk);
}
}
DirectHttpResponse::BrowserWreq(response) => {
let mut upstream_stream = response.bytes_stream();
while let Some(chunk) = upstream_stream.next().await {
let chunk = chunk.map_err(|err| {
SyncExecutionFailure::from_transport(
ExecutionRuntimeTransportError::UpstreamRequest(
format_wreq_upstream_request_error(&err),
),
)
})?;
let elapsed_ms = started_at.elapsed().as_millis() as u64;
progress
.observe_chunk(&chunk, status_code, elapsed_ms)
.await;
body_bytes.extend_from_slice(&chunk);
}
}
}
let decoded_body_bytes =
@@ -1077,64 +1101,106 @@ async fn execute_execution_runtime_sync_impl(
.await;
#[cfg(not(test))]
let mut result = {
match maybe_execute_chatgpt_web_image_sync(state, &plan, report_context.as_ref()).await {
match maybe_execute_grok_sync(&plan, report_context.as_ref()).await {
Ok(Some(result)) => result,
Ok(None) => match execute_direct_sync_runtime_candidate(
state,
&plan,
report_context.as_ref(),
trace_id,
plan_kind,
plan_request_id_for_log.as_str(),
plan_candidate_id.as_deref(),
provider_name.as_str(),
endpoint_id.as_str(),
key_id.as_str(),
model_name.as_str(),
candidate_index.as_str(),
progress_snapshot.clone(),
)
.await
{
Ok(result) => result,
Err(err) => {
warn!(
event_name = "sync_execution_runtime_unavailable",
log_type = "ops",
trace_id = %trace_id,
request_id = %plan_request_id_for_log,
candidate_id = ?plan_candidate_id,
provider_name,
endpoint_id,
key_id,
model_name,
candidate_index = candidate_index.as_str(),
error_type = err.error_type,
error = %err.message,
"gateway in-process sync execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
Ok(None) => {
match maybe_execute_chatgpt_web_image_sync(state, &plan, report_context.as_ref())
.await
{
Ok(Some(result)) => result,
Ok(None) => match execute_direct_sync_runtime_candidate(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: err.status_code,
error_type: Some(err.error_type.to_string()),
error_message: Some(err.message),
latency_ms: err.latency_ms,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
trace_id,
plan_kind,
plan_request_id_for_log.as_str(),
plan_candidate_id.as_deref(),
provider_name.as_str(),
endpoint_id.as_str(),
key_id.as_str(),
model_name.as_str(),
candidate_index.as_str(),
progress_snapshot.clone(),
)
.await;
return Ok(None);
.await
{
Ok(result) => result,
Err(err) => {
warn!(
event_name = "sync_execution_runtime_unavailable",
log_type = "ops",
trace_id = %trace_id,
request_id = %plan_request_id_for_log,
candidate_id = ?plan_candidate_id,
provider_name,
endpoint_id,
key_id,
model_name,
candidate_index = candidate_index.as_str(),
error_type = err.error_type,
error = %err.message,
"gateway in-process sync execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: err.status_code,
error_type: Some(err.error_type.to_string()),
error_message: Some(err.message),
latency_ms: err.latency_ms,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
},
Err(err) => {
warn!(
event_name = "chatgpt_web_image_execution_unavailable",
log_type = "ops",
trace_id = %trace_id,
request_id = %plan_request_id_for_log,
candidate_id = ?plan_candidate_id,
provider_name,
endpoint_id,
key_id,
model_name,
candidate_index = candidate_index.as_str(),
error = %err,
"gateway ChatGPT-Web image execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some(
"chatgpt_web_image_execution_unavailable".to_string(),
),
error_message: Some(err.to_string()),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
}
},
}
Err(err) => {
warn!(
event_name = "chatgpt_web_image_execution_unavailable",
event_name = "grok_execution_unavailable",
log_type = "ops",
trace_id = %trace_id,
request_id = %plan_request_id_for_log,
@@ -1145,7 +1211,7 @@ async fn execute_execution_runtime_sync_impl(
model_name,
candidate_index = candidate_index.as_str(),
error = %err,
"gateway ChatGPT-Web image execution unavailable"
"gateway Grok execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
@@ -1155,7 +1221,7 @@ async fn execute_execution_runtime_sync_impl(
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("chatgpt_web_image_execution_unavailable".to_string()),
error_type: Some("grok_execution_unavailable".to_string()),
error_message: Some(err.to_string()),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
@@ -1212,30 +1278,72 @@ async fn execute_execution_runtime_sync_impl(
.trim()
.is_empty()
{
match maybe_execute_chatgpt_web_image_sync(state, &plan, report_context.as_ref()).await
{
match maybe_execute_grok_sync(&plan, report_context.as_ref()).await {
Ok(Some(result)) => result,
Ok(None) => match execute_direct_sync_runtime_candidate(
Ok(None) => match maybe_execute_chatgpt_web_image_sync(
state,
&plan,
report_context.as_ref(),
trace_id,
plan_kind,
plan_request_id_for_log.as_str(),
plan_candidate_id.as_deref(),
provider_name.as_str(),
endpoint_id.as_str(),
key_id.as_str(),
model_name.as_str(),
candidate_index.as_str(),
progress_snapshot.clone(),
)
.await
{
Ok(result) => result,
Ok(Some(result)) => result,
Ok(None) => match execute_direct_sync_runtime_candidate(
state,
&plan,
report_context.as_ref(),
trace_id,
plan_kind,
plan_request_id_for_log.as_str(),
plan_candidate_id.as_deref(),
provider_name.as_str(),
endpoint_id.as_str(),
key_id.as_str(),
model_name.as_str(),
candidate_index.as_str(),
progress_snapshot.clone(),
)
.await
{
Ok(result) => result,
Err(err) => {
warn!(
event_name = "sync_execution_runtime_unavailable",
log_type = "ops",
trace_id = %trace_id,
request_id = %plan_request_id_for_log,
candidate_id = ?plan_candidate_id,
provider_name,
endpoint_id,
key_id,
model_name,
candidate_index = candidate_index.as_str(),
error_type = err.error_type,
error = %err.message,
"gateway in-process sync execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: err.status_code,
error_type: Some(err.error_type.to_string()),
error_message: Some(err.message),
latency_ms: err.latency_ms,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
}
},
Err(err) => {
warn!(
event_name = "sync_execution_runtime_unavailable",
event_name = "chatgpt_web_image_execution_unavailable",
log_type = "ops",
trace_id = %trace_id,
request_id = %plan_request_id_for_log,
@@ -1245,9 +1353,8 @@ async fn execute_execution_runtime_sync_impl(
key_id,
model_name,
candidate_index = candidate_index.as_str(),
error_type = err.error_type,
error = %err.message,
"gateway in-process sync execution unavailable"
error = %err,
"gateway ChatGPT-Web image execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
@@ -1256,10 +1363,12 @@ async fn execute_execution_runtime_sync_impl(
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: err.status_code,
error_type: Some(err.error_type.to_string()),
error_message: Some(err.message),
latency_ms: err.latency_ms,
status_code: None,
error_type: Some(
"chatgpt_web_image_execution_unavailable".to_string(),
),
error_message: Some(err.to_string()),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
@@ -1270,7 +1379,7 @@ async fn execute_execution_runtime_sync_impl(
},
Err(err) => {
warn!(
event_name = "chatgpt_web_image_execution_unavailable",
event_name = "grok_execution_unavailable",
log_type = "ops",
trace_id = %trace_id,
request_id = %plan_request_id_for_log,
@@ -1281,7 +1390,7 @@ async fn execute_execution_runtime_sync_impl(
model_name,
candidate_index = candidate_index.as_str(),
error = %err,
"gateway ChatGPT-Web image execution unavailable"
"gateway Grok execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
@@ -1291,7 +1400,7 @@ async fn execute_execution_runtime_sync_impl(
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("chatgpt_web_image_execution_unavailable".to_string()),
error_type: Some("grok_execution_unavailable".to_string()),
error_message: Some(err.to_string()),
latency_ms: None,
started_at_unix_ms: Some(candidate_started_unix_secs),

View File

@@ -8,7 +8,8 @@ use aether_contracts::{
ExecutionPlan, ExecutionResult, ExecutionTelemetry, ProxySnapshot, ResolvedTransportProfile,
ResponseBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_HTTP1_ONLY,
TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS,
TRANSPORT_HTTP_MODE_HTTP1_ONLY,
};
use aether_data::repository::proxy_nodes::ProxyNodeTrafficMutation;
use aether_http::{apply_http_client_config, HttpClientConfig};
@@ -81,6 +82,52 @@ pub(crate) fn format_upstream_request_error(err: &reqwest::Error) -> String {
detail
}
pub(crate) fn format_wreq_upstream_request_error(err: &wreq::Error) -> String {
let mut kinds = Vec::new();
if err.is_connect() {
kinds.push("connect");
}
if err.is_timeout() {
kinds.push("timeout");
}
if err.is_redirect() {
kinds.push("redirect");
}
if err.is_body() {
kinds.push("body");
}
if err.is_decode() {
kinds.push("decode");
}
if err.is_request() {
kinds.push("request");
}
let mut detail = err.to_string();
let mut source = err.source();
while let Some(cause) = source {
let cause_text = cause.to_string();
if !cause_text.is_empty() && !detail.contains(&cause_text) {
detail.push_str(": ");
detail.push_str(&cause_text);
}
source = cause.source();
}
if let Some(uri) = err.uri() {
detail.push_str(" [uri=");
detail.push_str(&uri.to_string());
detail.push(']');
}
if !kinds.is_empty() {
detail.push_str(" [kind=");
detail.push_str(&kinds.join(","));
detail.push(']');
}
detail
}
#[derive(Debug, Error)]
pub(crate) enum ExecutionRuntimeTransportError {
#[error("stream execution is not supported for this plan")]
@@ -107,6 +154,10 @@ pub(crate) enum ExecutionRuntimeTransportError {
BodyEncode(serde_json::Error),
#[error("failed to build HTTP client: {0}")]
ClientBuild(reqwest::Error),
#[error("failed to build browser impersonation HTTP client: {0}")]
BrowserClientBuild(wreq::Error),
#[error("browser impersonation response body failed: {0}")]
BrowserBody(String),
#[error("failed to execute upstream request: {0}")]
UpstreamRequest(String),
#[error("hub relay request failed: {0}")]
@@ -136,7 +187,7 @@ struct RelayRequestMeta {
pub(crate) struct DirectSyncExecutionRuntime;
#[derive(Debug, Clone, Copy, Default)]
struct ExecutionTransportControls {
pub(crate) struct ExecutionTransportControls {
follow_redirects: Option<bool>,
http1_only: bool,
accept_invalid_certs: bool,
@@ -144,6 +195,7 @@ struct ExecutionTransportControls {
pub(crate) enum DirectUpstreamResponse {
Reqwest(reqwest::Response),
BrowserWreq(wreq::Response),
LocalTunnel(tunnel::DirectRelayResponse),
}
@@ -172,11 +224,9 @@ impl DirectSyncExecutionRuntime {
let started_at = Instant::now();
let response = send_request(plan, body_bytes).await?;
let ttfb_ms = started_at.elapsed().as_millis() as u64;
let status_code = response.status().as_u16();
let headers = collect_response_headers(response.headers());
let body_bytes = response.bytes().await.map_err(|err| {
ExecutionRuntimeTransportError::UpstreamRequest(format_upstream_request_error(&err))
})?;
let status_code = response.status_code();
let headers = response.headers();
let body_bytes = response.bytes().await?;
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)
.unwrap_or_else(|| body_bytes.to_vec());
let elapsed_ms = started_at.elapsed().as_millis() as u64;
@@ -230,8 +280,8 @@ impl DirectSyncExecutionRuntime {
let started_at = Instant::now();
let response = send_request(plan, body_bytes).await?;
let status_code = response.status().as_u16();
let headers = collect_response_headers(response.headers());
let status_code = response.status_code();
let headers = response.headers();
let stream_summary_report_context = build_stream_summary_report_context(plan);
@@ -242,7 +292,7 @@ impl DirectSyncExecutionRuntime {
headers,
provider_api_format: plan.provider_api_format.clone(),
stream_summary_report_context,
response: DirectUpstreamResponse::Reqwest(response),
response: response.into_direct_upstream_response(),
started_at,
})
}
@@ -252,6 +302,15 @@ pub(crate) async fn execute_sync_plan(
state: &AppState,
trace_id: Option<&str>,
plan: &ExecutionPlan,
) -> Result<ExecutionResult, GatewayError> {
execute_sync_plan_with_report_context(state, trace_id, plan, None).await
}
pub(crate) async fn execute_sync_plan_with_report_context(
state: &AppState,
trace_id: Option<&str>,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<ExecutionResult, GatewayError> {
#[cfg(test)]
{
@@ -275,6 +334,18 @@ pub(crate) async fn execute_sync_plan(
.map_err(|err| GatewayError::Internal(err.to_string()));
}
match super::grok::maybe_execute_grok_sync(plan, report_context).await {
Ok(Some(result)) => {
record_manual_proxy_request_outcome(state, plan, result.status_code).await;
return Ok(result);
}
Ok(None) => {}
Err(err) => {
record_manual_proxy_request_failure(state, plan).await;
return Err(GatewayError::Internal(err.to_string()));
}
}
let _ = trace_id;
match DirectSyncExecutionRuntime::new().execute_sync(plan).await {
Ok(result) => {
@@ -554,7 +625,7 @@ fn build_direct_tunnel_request_meta(
pub(crate) async fn send_request(
plan: &ExecutionPlan,
body_bytes: Vec<u8>,
) -> Result<reqwest::Response, ExecutionRuntimeTransportError> {
) -> Result<DirectHttpResponse, ExecutionRuntimeTransportError> {
if let Some(detail) = gateway_frontdoor_self_loop_guard_error(plan.url.as_str()) {
return Err(ExecutionRuntimeTransportError::UpstreamRequest(detail));
}
@@ -572,6 +643,18 @@ pub(crate) async fn send_request(
.and_then(|timeouts| timeouts.total_ms)
.map(Duration::from_millis);
if transport_profile_uses_browser_wreq(plan.transport_profile.as_ref()) {
return send_via_browser_wreq_transport(
plan,
method,
headers,
body_bytes,
total_timeout,
transport_controls,
)
.await;
}
if let Some(node_id) = resolve_tunnel_node_id(plan.proxy.as_ref()) {
return send_via_tunnel_relay(
plan,
@@ -582,7 +665,8 @@ pub(crate) async fn send_request(
total_timeout,
transport_controls,
)
.await;
.await
.map(DirectHttpResponse::Reqwest);
}
let client = build_client(
@@ -596,9 +680,95 @@ pub(crate) async fn send_request(
if let Some(timeout) = total_timeout {
request = request.timeout(timeout);
}
request.send().await.map_err(|err| {
ExecutionRuntimeTransportError::UpstreamRequest(format_upstream_request_error(&err))
})
request
.send()
.await
.map(DirectHttpResponse::Reqwest)
.map_err(|err| {
ExecutionRuntimeTransportError::UpstreamRequest(format_upstream_request_error(&err))
})
}
pub(crate) enum DirectHttpResponse {
Reqwest(reqwest::Response),
BrowserWreq(wreq::Response),
}
impl DirectHttpResponse {
pub(crate) fn status_code(&self) -> u16 {
match self {
DirectHttpResponse::Reqwest(response) => response.status().as_u16(),
DirectHttpResponse::BrowserWreq(response) => response.status().as_u16(),
}
}
pub(crate) fn headers(&self) -> BTreeMap<String, String> {
match self {
DirectHttpResponse::Reqwest(response) => collect_response_headers(response.headers()),
DirectHttpResponse::BrowserWreq(response) => {
collect_response_headers(response.headers())
}
}
}
pub(crate) async fn bytes(self) -> Result<Bytes, ExecutionRuntimeTransportError> {
match self {
DirectHttpResponse::Reqwest(response) => response.bytes().await.map_err(|err| {
ExecutionRuntimeTransportError::UpstreamRequest(format_upstream_request_error(&err))
}),
DirectHttpResponse::BrowserWreq(response) => response.bytes().await.map_err(|err| {
ExecutionRuntimeTransportError::BrowserBody(format_wreq_upstream_request_error(
&err,
))
}),
}
}
fn into_direct_upstream_response(self) -> DirectUpstreamResponse {
match self {
DirectHttpResponse::Reqwest(response) => DirectUpstreamResponse::Reqwest(response),
DirectHttpResponse::BrowserWreq(response) => {
DirectUpstreamResponse::BrowserWreq(response)
}
}
}
}
async fn send_via_browser_wreq_transport(
plan: &ExecutionPlan,
method: reqwest::Method,
headers: HeaderMap,
body_bytes: Vec<u8>,
total_timeout: Option<Duration>,
transport_controls: ExecutionTransportControls,
) -> Result<DirectHttpResponse, ExecutionRuntimeTransportError> {
let profile = plan.transport_profile.as_ref().ok_or_else(|| {
ExecutionRuntimeTransportError::UnsupportedTransportProfile(String::new())
})?;
let client = build_browser_wreq_client(
plan.timeouts.as_ref(),
plan.proxy.as_ref(),
profile,
transport_controls,
)?;
let method = wreq::Method::from_bytes(method.as_str().as_bytes())
.map_err(ExecutionRuntimeTransportError::InvalidMethod)?;
let mut request = client
.request(method, plan.url.as_str())
.headers(headers)
.body(body_bytes);
if let Some(timeout) = total_timeout {
request = request.timeout(timeout);
}
request
.send()
.await
.map(DirectHttpResponse::BrowserWreq)
.map_err(|err| {
ExecutionRuntimeTransportError::UpstreamRequest(format_wreq_upstream_request_error(
&err,
))
})
}
async fn send_via_tunnel_relay(
@@ -905,6 +1075,96 @@ fn build_client(
.map_err(ExecutionRuntimeTransportError::ClientBuild)
}
pub(crate) fn build_browser_wreq_client(
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
proxy: Option<&ProxySnapshot>,
transport_profile: &ResolvedTransportProfile,
transport_controls: ExecutionTransportControls,
) -> Result<wreq::Client, ExecutionRuntimeTransportError> {
let emulation = browser_wreq_emulation_from_profile(transport_profile)?;
let mut builder = wreq::Client::builder().emulation(emulation);
if transport_controls.follow_redirects == Some(true) {
builder = builder.redirect(wreq::redirect::Policy::limited(10));
}
if transport_controls.http1_only || transport_profile_http1_only(Some(transport_profile)) {
builder = builder.http1_only();
}
if transport_controls.accept_invalid_certs {
builder = builder.cert_verification(false).verify_hostname(false);
}
if let Some(connect_ms) = timeouts.and_then(|timeouts| timeouts.connect_ms) {
builder = builder.connect_timeout(Duration::from_millis(connect_ms));
}
if let Some(total_ms) = timeouts.and_then(|timeouts| timeouts.total_ms) {
builder = builder.timeout(Duration::from_millis(total_ms));
}
if let Some(read_ms) = timeouts.and_then(|timeouts| timeouts.read_ms) {
builder = builder.read_timeout(Duration::from_millis(read_ms));
}
if let Some(proxy_url) = resolve_proxy_url(proxy)? {
let proxy = wreq::Proxy::all(proxy_url.as_str())
.map_err(ExecutionRuntimeTransportError::BrowserClientBuild)?;
builder = builder.proxy(proxy);
}
builder
.build()
.map_err(ExecutionRuntimeTransportError::BrowserClientBuild)
}
fn browser_wreq_emulation_from_profile(
profile: &ResolvedTransportProfile,
) -> Result<wreq_util::Emulation, ExecutionRuntimeTransportError> {
match normalize_browser_profile_name(browser_transport_profile_name(profile)).as_str() {
"chrome100" => Ok(wreq_util::Emulation::Chrome100),
"chrome101" => Ok(wreq_util::Emulation::Chrome101),
"chrome104" => Ok(wreq_util::Emulation::Chrome104),
"chrome105" => Ok(wreq_util::Emulation::Chrome105),
"chrome106" => Ok(wreq_util::Emulation::Chrome106),
"chrome107" => Ok(wreq_util::Emulation::Chrome107),
"chrome108" => Ok(wreq_util::Emulation::Chrome108),
"chrome109" => Ok(wreq_util::Emulation::Chrome109),
"chrome110" => Ok(wreq_util::Emulation::Chrome110),
"chrome114" => Ok(wreq_util::Emulation::Chrome114),
"chrome116" => Ok(wreq_util::Emulation::Chrome116),
"chrome117" => Ok(wreq_util::Emulation::Chrome117),
"chrome118" => Ok(wreq_util::Emulation::Chrome118),
"chrome119" => Ok(wreq_util::Emulation::Chrome119),
"chrome120" => Ok(wreq_util::Emulation::Chrome120),
"chrome123" => Ok(wreq_util::Emulation::Chrome123),
"chrome124" => Ok(wreq_util::Emulation::Chrome124),
"chrome126" => Ok(wreq_util::Emulation::Chrome126),
"chrome127" => Ok(wreq_util::Emulation::Chrome127),
"chrome128" => Ok(wreq_util::Emulation::Chrome128),
"chrome129" => Ok(wreq_util::Emulation::Chrome129),
"chrome130" => Ok(wreq_util::Emulation::Chrome130),
"chrome131" => Ok(wreq_util::Emulation::Chrome131),
"chrome132" => Ok(wreq_util::Emulation::Chrome132),
"chrome133" => Ok(wreq_util::Emulation::Chrome133),
"chrome134" => Ok(wreq_util::Emulation::Chrome134),
"chrome135" => Ok(wreq_util::Emulation::Chrome135),
"chrome136" => Ok(wreq_util::Emulation::Chrome136),
"chrome137" => Ok(wreq_util::Emulation::Chrome137),
"chrome138" => Ok(wreq_util::Emulation::Chrome138),
"chrome139" => Ok(wreq_util::Emulation::Chrome139),
"chrome140" => Ok(wreq_util::Emulation::Chrome140),
"chrome141" => Ok(wreq_util::Emulation::Chrome141),
"chrome142" => Ok(wreq_util::Emulation::Chrome142),
"chrome143" => Ok(wreq_util::Emulation::Chrome143),
"chrome144" => Ok(wreq_util::Emulation::Chrome144),
"chrome145" => Ok(wreq_util::Emulation::Chrome145),
other => Err(ExecutionRuntimeTransportError::UnsupportedTransportProfile(
format!("browser_wreq:{other}"),
)),
}
}
fn normalize_browser_profile_name(value: String) -> String {
value
.trim()
.to_ascii_lowercase()
.replace(['_', '-', ' '], "")
}
fn validate_reqwest_transport_profile(
transport_profile: Option<&ResolvedTransportProfile>,
) -> Result<(), ExecutionRuntimeTransportError> {
@@ -923,6 +1183,56 @@ fn validate_reqwest_transport_profile(
))
}
fn transport_profile_uses_browser_wreq(
transport_profile: Option<&ResolvedTransportProfile>,
) -> bool {
transport_profile
.map(|profile| {
profile
.backend
.trim()
.eq_ignore_ascii_case(TRANSPORT_BACKEND_BROWSER_WREQ)
})
.unwrap_or(false)
}
fn browser_transport_profile_name(profile: &ResolvedTransportProfile) -> String {
profile
.extra
.as_ref()
.and_then(|value| {
value
.get("browser_profile")
.or_else(|| value.get("impersonate"))
.and_then(Value::as_str)
})
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| {
profile
.profile_id
.trim()
.is_empty()
.then_some("chrome136".to_string())
.or_else(|| Some(profile.profile_id.trim().to_string()))
})
.unwrap_or_else(|| "chrome136".to_string())
}
fn insert_browser_control_header(
headers: &mut HeaderMap,
name: &'static str,
value: &str,
) -> Result<(), ExecutionRuntimeTransportError> {
headers.insert(
HeaderName::from_static(name),
HeaderValue::from_str(value)
.map_err(|_| ExecutionRuntimeTransportError::InvalidHeaderValue(name.to_string()))?,
);
Ok(())
}
fn transport_profile_http1_only(transport_profile: Option<&ResolvedTransportProfile>) -> bool {
transport_profile
.map(|profile| {
@@ -991,7 +1301,7 @@ fn resolve_proxy_url(
Ok(None)
}
fn build_request_headers(
pub(crate) fn build_request_headers(
headers: &BTreeMap<String, String>,
content_encoding: Option<&str>,
allow_passthrough_content_encoding: bool,
@@ -1180,22 +1490,22 @@ mod tests {
use aether_contracts::{
ExecutionPlan, ExecutionTimeouts, ProxySnapshot, RequestBody, ResolvedTransportProfile,
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
TRANSPORT_BACKEND_REQWEST_RUSTLS,
TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS,
};
use aether_data::repository::proxy_nodes::{
InMemoryProxyNodeRepository, ProxyNodeReadRepository, StoredProxyNode,
};
use axum::body::Bytes;
use axum::body::{Body, Bytes};
use axum::extract::ws::Message;
use axum::extract::Path;
use axum::http::HeaderMap as AxumHeaderMap;
use axum::routing::post;
use axum::routing::{any, post};
use axum::{Json, Router};
use serde_json::json;
use tokio::sync::watch;
use super::{
build_client, build_request_headers, execute_sync_plan,
build_browser_wreq_client, build_client, build_request_headers, execute_sync_plan,
record_manual_proxy_request_failure, record_manual_proxy_request_outcome,
record_manual_proxy_request_success, record_manual_proxy_stream_error,
resolve_execution_transport_controls, DirectSyncExecutionRuntime,
@@ -1431,6 +1741,228 @@ mod tests {
);
}
#[tokio::test]
async fn direct_sync_execution_runtime_routes_browser_wreq_transport_in_process() {
async fn browser_upstream(headers: AxumHeaderMap, body: Bytes) -> axum::response::Response {
assert_eq!(
headers
.get("content-type")
.and_then(|value| value.to_str().ok()),
Some("application/json")
);
assert!(
headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.is_none(),
"internal execution control headers must not leak upstream"
);
assert_eq!(body.as_ref(), br#"{"modelName":"auto"}"#);
axum::response::Response::builder()
.status(http::StatusCode::ACCEPTED)
.header("content-type", "application/json")
.body(Body::from(
json!({
"ok": true,
"via": "browser_wreq"
})
.to_string(),
))
.expect("response should build")
}
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let app = Router::new().route("/request", any(browser_upstream));
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("test server should run");
});
let plan = ExecutionPlan {
request_id: "req-browser-wreq".into(),
candidate_id: None,
provider_name: Some("grok".into()),
provider_id: "provider-1".into(),
endpoint_id: "endpoint-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: format!("http://{addr}/request"),
headers: BTreeMap::from([
("content-type".into(), "application/json".into()),
(
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER.into(),
"true".into(),
),
]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({"modelName":"auto"})),
stream: false,
client_api_format: "openai:responses".into(),
provider_api_format: "grok:rate_limits".into(),
model_name: Some("grok-quota".into()),
proxy: None,
transport_profile: Some(ResolvedTransportProfile {
profile_id: "chrome136".into(),
backend: TRANSPORT_BACKEND_BROWSER_WREQ.into(),
http_mode: "auto".into(),
pool_scope: "key".into(),
header_fingerprint: None,
extra: Some(json!({
"browser_profile": "chrome136"
})),
}),
timeouts: Some(ExecutionTimeouts {
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let result = DirectSyncExecutionRuntime::new()
.execute_sync(&plan)
.await
.expect("browser wreq transport plan should execute in-process");
server.abort();
assert_eq!(result.status_code, http::StatusCode::ACCEPTED.as_u16());
assert_eq!(
result
.body
.and_then(|body| body.json_body)
.and_then(|body| body.get("via").cloned()),
Some(json!("browser_wreq"))
);
}
#[test]
fn browser_wreq_transport_rejects_unknown_profile() {
let profile = ResolvedTransportProfile {
profile_id: "firefox999".into(),
backend: TRANSPORT_BACKEND_BROWSER_WREQ.into(),
http_mode: "auto".into(),
pool_scope: "key".into(),
header_fingerprint: None,
extra: None,
};
let error = match build_browser_wreq_client(
None,
None,
&profile,
ExecutionTransportControls::default(),
) {
Ok(_) => panic!("unknown browser profile should fail loudly"),
Err(error) => error,
};
assert!(matches!(
error,
ExecutionRuntimeTransportError::UnsupportedTransportProfile(backend)
if backend == "browser_wreq:firefox999"
));
}
#[tokio::test]
async fn execute_sync_plan_routes_grok_marker_through_grok_runtime() {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let app = Router::new().route(
"/rest/app-chat/conversations/new",
post(|body: Bytes| async move {
let body_json: serde_json::Value =
serde_json::from_slice(&body).expect("request body should be json");
if body_json.get("message").and_then(serde_json::Value::as_str)
!= Some("[user]: hello")
{
return (
axum::http::StatusCode::BAD_REQUEST,
Json(json!({
"error": {
"message": "expected grok app-chat message",
"body": body_json,
}
})),
);
}
(
axum::http::StatusCode::OK,
Json(json!({
"result": {
"response": {
"token": "pong",
"messageTag": "final"
}
}
})),
)
}),
);
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("test server should run");
});
let plan = ExecutionPlan {
request_id: "req-grok-runtime".into(),
candidate_id: Some("cand-grok".into()),
provider_name: Some("grok".into()),
provider_id: "provider-grok".into(),
endpoint_id: "endpoint-grok".into(),
key_id: "key-grok".into(),
method: "POST".into(),
url: format!("http://{addr}/rest/app-chat/conversations/new"),
headers: BTreeMap::from([
("content-type".into(), "application/json".into()),
(
aether_provider_transport::GROK_INTERNAL_HEADER.into(),
"1".into(),
),
]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "grok-4.20-0309-non-reasoning",
"messages": [{"role": "user", "content": "hello"}],
})),
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("grok-4.20-0309-non-reasoning".into()),
proxy: None,
transport_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let report_context = json!({"mapped_model": "grok-4.20-fast"});
let result = super::super::grok::maybe_execute_grok_sync(&plan, Some(&report_context))
.await
.expect("grok runtime plan should execute")
.expect("grok runtime should handle marked plan");
server.abort();
assert_eq!(result.status_code, http::StatusCode::OK.as_u16());
assert_eq!(
result
.body
.and_then(|body| body.json_body)
.and_then(|body| body["choices"][0]["message"]["content"]
.as_str()
.map(str::to_string)),
Some("pong".to_string())
);
}
#[tokio::test]
async fn execute_sync_plan_records_manual_proxy_success() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![

View File

@@ -1,6 +1,9 @@
use crate::async_task::CancelVideoTaskError;
use crate::control::GatewayControlDecision;
use crate::control::GatewayPublicRequestContext;
use crate::image_capabilities::{
openai_image_gateway_max_generation_count, openai_image_gateway_max_generation_count_for_model,
};
use crate::{AppState, GatewayError};
use aether_data_contracts::repository::video_tasks::{
StoredVideoTask, VideoTaskQueryFilter, VideoTaskStatus,
@@ -18,9 +21,6 @@ const AI_PUBLIC_METHOD_NOT_ALLOWED_DETAIL: &str = "Method not allowed";
const AI_PUBLIC_UNAUTHORIZED_DETAIL: &str = "Unauthorized";
const OPENAI_IMAGE_PROMPT_DETAIL: &str = "图片生成/编辑请求缺少 prompt";
const OPENAI_IMAGE_EDIT_INPUT_DETAIL: &str = "图片编辑请求至少需要 1 张输入图片";
const OPENAI_IMAGE_VARIATION_INPUT_DETAIL: &str = "图片变体请求需要 image 文件";
const OPENAI_IMAGE_N_DETAIL: &str = "当前 Codex 图片反代仅支持 n=1";
const OPENAI_IMAGE_STREAM_VARIATION_DETAIL: &str = "图片变体接口当前仅支持同步响应";
const OPENAI_IMAGE_PARTIAL_IMAGES_DETAIL: &str =
"partial_images 仅支持 0-3且必须配合 stream=true";
const OPENAI_IMAGE_STYLE_DETAIL: &str = "当前 Codex 图片反代暂不支持 style 参数";
@@ -57,7 +57,6 @@ const OPENAI_RERANK_STREAM_UNSUPPORTED_DETAIL: &str = "Rerank requests do not su
enum OpenAiImageOperation {
Generate,
Edit,
Variation,
}
impl OpenAiImageOperation {
@@ -65,7 +64,6 @@ impl OpenAiImageOperation {
match path {
"/v1/images/generations" => Some(Self::Generate),
"/v1/images/edits" => Some(Self::Edit),
"/v1/images/variations" => Some(Self::Variation),
_ => None,
}
}
@@ -202,7 +200,7 @@ fn maybe_build_local_openai_request_validation_response(
if decision.route_kind.as_deref() != Some("image")
|| !matches!(
request_context.request_path.as_str(),
"/v1/images/generations" | "/v1/images/edits" | "/v1/images/variations"
"/v1/images/generations" | "/v1/images/edits"
)
{
return None;
@@ -240,19 +238,13 @@ fn maybe_build_local_openai_request_validation_response(
OPENAI_IMAGE_EDIT_INPUT_DETAIL,
));
}
OpenAiImageOperation::Variation if validation.image_count == 0 => {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_VARIATION_INPUT_DETAIL,
));
}
_ => {}
}
if validation.n.is_some_and(|value| value != 1) {
if let Some(detail) = validate_openai_image_n(&validation) {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_N_DETAIL,
detail,
));
}
@@ -272,15 +264,6 @@ fn maybe_build_local_openai_request_validation_response(
));
}
if validation.stream {
if operation == OpenAiImageOperation::Variation {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_IMAGE_STREAM_VARIATION_DETAIL,
));
}
}
if validation
.response_format
.as_deref()
@@ -360,6 +343,23 @@ fn maybe_build_local_openai_request_validation_response(
None
}
fn openai_image_n_detail(max_generation_count: u64) -> String {
if max_generation_count >= openai_image_gateway_max_generation_count() {
format!("当前图片反代仅支持 n=1..{max_generation_count}")
} else {
format!("当前图片模型仅支持 n=1..{max_generation_count}")
}
}
fn validate_openai_image_n(validation: &OpenAiImageValidationInput) -> Option<String> {
let max_generation_count =
openai_image_gateway_max_generation_count_for_model(validation.model.as_deref());
validation
.n
.is_some_and(|value| value == 0 || value > max_generation_count)
.then(|| openai_image_n_detail(max_generation_count))
}
fn validate_openai_embedding_request(
content_type: Option<&str>,
request_body: &Bytes,
@@ -535,7 +535,6 @@ fn parse_openai_image_validation_input(
OpenAiImageOperation::Generate | OpenAiImageOperation::Edit => {
OPENAI_IMAGE_PROMPT_DETAIL
}
OpenAiImageOperation::Variation => OPENAI_IMAGE_VARIATION_INPUT_DETAIL,
});
}
@@ -1260,7 +1259,8 @@ fn estimate_text_tokens(text: &str) -> u64 {
#[cfg(test)]
mod tests {
use super::{
estimate_claude_count_tokens, parse_openai_image_validation_input, OpenAiImageOperation,
estimate_claude_count_tokens, parse_openai_image_validation_input, validate_openai_image_n,
OpenAiImageOperation,
};
use axum::body::Bytes;
use serde_json::json;
@@ -1344,4 +1344,31 @@ mod tests {
assert_eq!(validation.prompt.as_deref(), Some("edit this image"));
assert_eq!(validation.image_count, 1);
}
#[test]
fn image_validation_restricts_multi_image_count_to_grok_models() {
let openai_body = Bytes::from_static(br#"{"model":"gpt-image-2","prompt":"draw","n":2}"#);
let openai_validation = parse_openai_image_validation_input(
OpenAiImageOperation::Generate,
Some("application/json"),
&openai_body,
)
.expect("valid image payload should parse");
assert_eq!(
validate_openai_image_n(&openai_validation).as_deref(),
Some("当前图片模型仅支持 n=1..1")
);
let grok_body =
Bytes::from_static(br#"{"model":"grok-imagine-image-lite","prompt":"draw","n":4}"#);
let grok_validation = parse_openai_image_validation_input(
OpenAiImageOperation::Generate,
Some("application/json"),
&grok_body,
)
.expect("valid grok image payload should parse");
assert!(validate_openai_image_n(&grok_validation).is_none());
}
}

View File

@@ -0,0 +1,60 @@
use crate::ai_serving::OpenAiImageNormalizeOptions;
const DEFAULT_OPENAI_IMAGE_MAX_GENERATION_COUNT: u64 = 1;
const GROK_OPENAI_IMAGE_MAX_GENERATION_COUNT: u64 = 4;
pub(crate) fn openai_image_gateway_max_generation_count() -> u64 {
GROK_OPENAI_IMAGE_MAX_GENERATION_COUNT
}
pub(crate) fn openai_image_gateway_max_generation_count_for_model(model: Option<&str>) -> u64 {
if model.is_some_and(is_grok_openai_image_model) {
GROK_OPENAI_IMAGE_MAX_GENERATION_COUNT
} else {
DEFAULT_OPENAI_IMAGE_MAX_GENERATION_COUNT
}
}
pub(crate) fn openai_image_provider_max_generation_count(provider_type: &str) -> u64 {
if provider_type.trim().eq_ignore_ascii_case("grok") {
GROK_OPENAI_IMAGE_MAX_GENERATION_COUNT
} else {
DEFAULT_OPENAI_IMAGE_MAX_GENERATION_COUNT
}
}
pub(crate) fn openai_image_normalize_options_for_provider(
provider_type: &str,
) -> OpenAiImageNormalizeOptions {
OpenAiImageNormalizeOptions::with_max_generation_count(
openai_image_provider_max_generation_count(provider_type),
)
}
fn is_grok_openai_image_model(model: &str) -> bool {
model
.trim()
.to_ascii_lowercase()
.contains("grok-imagine-image")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn grok_owns_gateway_wide_image_generation_count_ceiling() {
assert_eq!(openai_image_gateway_max_generation_count(), 4);
assert_eq!(
openai_image_gateway_max_generation_count_for_model(Some("grok-imagine-image-lite")),
4
);
assert_eq!(
openai_image_gateway_max_generation_count_for_model(Some("gpt-image-2")),
1
);
assert_eq!(openai_image_gateway_max_generation_count_for_model(None), 1);
assert_eq!(openai_image_provider_max_generation_count("grok"), 4);
assert_eq!(openai_image_provider_max_generation_count("openai"), 1);
}
}

View File

@@ -45,6 +45,7 @@ mod frontdoor_loop_guard;
mod handlers;
mod headers;
mod hooks;
mod image_capabilities;
mod log_ids;
mod maintenance;
pub(crate) mod middleware;

View File

@@ -3,12 +3,14 @@ use std::sync::{Mutex, OnceLock};
use std::time::{Duration, Instant};
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_provider_pool::grok_quota_window_key_for_model;
use aether_usage_runtime::{
extract_gemini_file_mapping_entries, gemini_file_mapping_cache_key, normalize_gemini_file_name,
report_request_id, GatewayStreamReportRequest, GatewaySyncReportRequest,
GEMINI_FILE_MAPPING_TTL_SECONDS,
};
use serde_json::Value;
use regex::Regex;
use serde_json::{json, Value};
use tracing::warn;
use uuid::Uuid;
@@ -23,6 +25,8 @@ const CODEX_QUOTA_CACHE_MAX_ENTRIES: usize = 4096;
type HeaderFingerprintCache = Mutex<HashMap<String, (String, Instant)>>;
static CODEX_QUOTA_HEADER_FINGERPRINT_CACHE: OnceLock<HeaderFingerprintCache> = OnceLock::new();
static GROK_CHINESE_WAIT_DURATION_RE: OnceLock<Regex> = OnceLock::new();
static GROK_ENGLISH_WAIT_DURATION_RE: OnceLock<Regex> = OnceLock::new();
#[derive(Debug, Clone, Copy)]
pub(crate) enum LocalReportEffect<'a> {
@@ -167,6 +171,264 @@ fn merge_metadata_object(
Some(Value::Object(merged))
}
fn grok_report_context_model(report_context: Option<&Value>) -> Option<String> {
report_context
.and_then(|context| context.get("mapped_model"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn grok_chinese_wait_duration_regex() -> &'static Regex {
GROK_CHINESE_WAIT_DURATION_RE.get_or_init(|| {
Regex::new(r"(?i)(?:(\d+)\s*天)?\s*(?:(\d+)\s*(?:小时|小時))?\s*(?:(\d+)\s*分钟)?\s*(?:(\d+)\s*秒)?")
.expect("grok Chinese wait duration regex should compile")
})
}
fn grok_english_wait_duration_regex() -> &'static Regex {
GROK_ENGLISH_WAIT_DURATION_RE.get_or_init(|| {
Regex::new(
r"(?i)(?:(\d+)\s*(?:d|day|days))?\s*(?:(\d+)\s*(?:h|hour|hours))?\s*(?:(\d+)\s*(?:m|min|mins|minute|minutes))?\s*(?:(\d+)\s*(?:s|sec|secs|second|seconds))?",
)
.expect("grok English wait duration regex should compile")
})
}
fn grok_duration_capture_seconds(captures: regex::Captures<'_>) -> Option<u64> {
let values = [1usize, 2, 3, 4]
.into_iter()
.map(|index| {
captures
.get(index)
.and_then(|item| item.as_str().parse::<u64>().ok())
.unwrap_or(0)
})
.collect::<Vec<_>>();
let seconds = values[0]
.saturating_mul(86_400)
.saturating_add(values[1].saturating_mul(3_600))
.saturating_add(values[2].saturating_mul(60))
.saturating_add(values[3]);
(seconds > 0).then_some(seconds)
}
fn grok_wait_duration_seconds_from_text(text: &str) -> Option<u64> {
let text = text.trim();
if text.is_empty() {
return None;
}
for captures in grok_chinese_wait_duration_regex().captures_iter(text) {
if let Some(seconds) = grok_duration_capture_seconds(captures) {
return Some(seconds);
}
}
for captures in grok_english_wait_duration_regex().captures_iter(text) {
if let Some(seconds) = grok_duration_capture_seconds(captures) {
return Some(seconds);
}
}
None
}
fn grok_response_error_text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.trim().to_string()).filter(|text| !text.is_empty()),
Value::Object(object) => {
if let Some(error) = object.get("error") {
if let Some(text) = grok_response_error_text(error) {
return Some(text);
}
}
for key in ["message", "detail", "reason", "error"] {
if let Some(text) = object
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|text| !text.is_empty())
{
return Some(text.to_string());
}
}
None
}
_ => None,
}
}
fn grok_upstream_response_body(report_context: Option<&Value>) -> Option<&Value> {
report_context
.and_then(|context| context.get("upstream_response"))
.and_then(|response| response.get("body"))
}
fn grok_quota_reset_after_seconds(
body_json: Option<&Value>,
report_context: Option<&Value>,
) -> Option<u64> {
body_json
.and_then(grok_response_error_text)
.and_then(|text| grok_wait_duration_seconds_from_text(text.as_str()))
.or_else(|| {
grok_upstream_response_body(report_context)
.and_then(grok_response_error_text)
.and_then(|text| grok_wait_duration_seconds_from_text(text.as_str()))
})
}
fn grok_apply_quota_feedback(
bucket: &mut serde_json::Map<String, Value>,
model: &str,
status_code: u16,
reset_after_seconds: Option<u64>,
now_unix_secs: u64,
) -> bool {
let Some(quota_key) = grok_quota_window_key_for_model(Some(model)) else {
return false;
};
let quota_by_model = if bucket.contains_key("quota_by_model") {
bucket.get_mut("quota_by_model")
} else {
bucket.get_mut("models")
};
let Some(window) = quota_by_model
.and_then(Value::as_object_mut)
.and_then(|models| models.get_mut(quota_key))
.and_then(Value::as_object_mut)
else {
return false;
};
let total = window
.get("total")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.filter(|value| *value > 0.0);
let current_remaining = window
.get("remaining")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.or(total)
.unwrap_or(0.0);
let next_remaining = match status_code {
429 => 0.0,
code if code >= 400 && reset_after_seconds.is_some() => 0.0,
code if code < 300 => (current_remaining - 1.0).max(0.0),
_ => return false,
};
window.insert("remaining".to_string(), json!(next_remaining));
if let Some(total) = total {
window.insert("total".to_string(), json!(total));
window.insert(
"remaining_fraction".to_string(),
json!((next_remaining / total).clamp(0.0, 1.0)),
);
window.insert(
"used_percent".to_string(),
json!(((total - next_remaining).max(0.0) / total * 100.0).clamp(0.0, 100.0)),
);
} else if status_code == 429 {
window.insert("remaining_fraction".to_string(), json!(0.0));
window.insert("used_percent".to_string(), json!(100.0));
}
if let Some(reset_after_seconds) = reset_after_seconds.filter(|seconds| *seconds > 0) {
let reset_at = now_unix_secs.saturating_add(reset_after_seconds);
window.insert("reset_at".to_string(), json!(reset_at));
window.insert("next_reset_at".to_string(), json!(reset_at));
window.insert(
"reset_after_seconds".to_string(),
json!(reset_after_seconds),
);
window.insert("reset_at_source".to_string(), json!("grok_upstream_error"));
}
window.insert("is_exhausted".to_string(), json!(next_remaining <= 0.0));
true
}
fn grok_mark_quota_bucket_updated(bucket: &mut serde_json::Map<String, Value>, now_unix_secs: u64) {
bucket.insert("updated_at".to_string(), json!(now_unix_secs));
}
async fn sync_grok_quota_from_report_context(
state: &AppState,
report_context: Option<&Value>,
status_code: u16,
body_json: Option<&Value>,
) -> Result<bool, GatewayError> {
let key_id = match report_context_key_id(report_context) {
Some(value) => value,
None => return Ok(false),
};
let Some(model) = grok_report_context_model(report_context) else {
return Ok(false);
};
let Some(key) = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next()
else {
return Ok(false);
};
let Some(provider) = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
.await?
.into_iter()
.next()
else {
return Ok(false);
};
if !provider.provider_type.trim().eq_ignore_ascii_case("grok") {
return Ok(false);
}
let Some(grok_bucket) = key
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("grok"))
.and_then(Value::as_object)
.cloned()
else {
return Ok(false);
};
let mut updated_grok_bucket = grok_bucket;
let now_unix_secs = current_unix_secs();
if !grok_apply_quota_feedback(
&mut updated_grok_bucket,
model.as_str(),
status_code,
grok_quota_reset_after_seconds(body_json, report_context),
now_unix_secs,
) {
return Ok(false);
}
grok_mark_quota_bucket_updated(&mut updated_grok_bucket, now_unix_secs);
let updated_upstream_metadata = merge_metadata_object(
key.upstream_metadata.as_ref(),
"grok",
Value::Object(updated_grok_bucket),
);
let updated_status_snapshot = sync_provider_key_quota_status_snapshot(
key.status_snapshot.as_ref(),
provider.provider_type.as_str(),
updated_upstream_metadata.as_ref(),
"report_effect",
);
let mut updated_key = key;
updated_key.upstream_metadata = updated_upstream_metadata;
updated_key.status_snapshot = updated_status_snapshot;
updated_key.updated_at_unix_secs = Some(now_unix_secs);
Ok(state
.update_provider_catalog_key(&updated_key)
.await?
.is_some())
}
async fn apply_local_sync_report_effect(state: &AppState, payload: &GatewaySyncReportRequest) {
apply_local_gemini_file_mapping_report_effect(state, payload).await;
if let Err(err) = sync_codex_quota_from_response_headers(
@@ -185,6 +447,23 @@ async fn apply_local_sync_report_effect(state: &AppState, payload: &GatewaySyncR
"gateway failed to persist codex realtime quota from sync response headers"
);
}
if let Err(err) = sync_grok_quota_from_report_context(
state,
payload.report_context.as_ref(),
payload.status_code,
payload.body_json.as_ref(),
)
.await
{
warn!(
event_name = "grok_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist grok realtime quota from sync response"
);
}
}
async fn apply_local_stream_report_effect(state: &AppState, payload: &GatewayStreamReportRequest) {
@@ -204,6 +483,23 @@ async fn apply_local_stream_report_effect(state: &AppState, payload: &GatewayStr
"gateway failed to persist codex realtime quota from stream response headers"
);
}
if let Err(err) = sync_grok_quota_from_report_context(
state,
payload.report_context.as_ref(),
payload.status_code,
grok_upstream_response_body(payload.report_context.as_ref()),
)
.await
{
warn!(
event_name = "grok_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist grok realtime quota from stream response"
);
}
}
async fn apply_local_gemini_file_mapping_report_effect(
@@ -443,3 +739,196 @@ pub(crate) fn clear_local_report_effect_caches_for_tests() {
.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn grok_quota_feedback_decrements_the_matching_window() {
let mut bucket = json!({
"quota_by_model": {
"quota_fast": {
"display_name": "fast",
"remaining": 30.0,
"total": 30.0,
"remaining_fraction": 1.0,
"used_percent": 0.0,
"is_exhausted": false
}
}
})
.as_object()
.cloned()
.expect("bucket should be object");
assert!(grok_apply_quota_feedback(
&mut bucket,
"grok-4.20-fast",
200,
None,
1_700_000_000
));
let fast = bucket
.get("quota_by_model")
.and_then(Value::as_object)
.and_then(|models| models.get("quota_fast"))
.and_then(Value::as_object)
.expect("fast window should exist");
assert_eq!(fast.get("remaining"), Some(&json!(29.0)));
assert_eq!(fast.get("remaining_fraction"), Some(&json!(29.0 / 30.0)));
assert_eq!(fast.get("used_percent"), Some(&json!(100.0 / 30.0)));
assert_eq!(fast.get("is_exhausted"), Some(&json!(false)));
}
#[test]
fn grok_report_context_model_requires_mapped_model() {
assert_eq!(
grok_report_context_model(Some(&json!({
"model": "grok-4.20-0309-reasoning"
}))),
None
);
assert_eq!(
grok_report_context_model(Some(&json!({
"mapped_model": "grok-4.20-fast"
}))),
Some("grok-4.20-fast".to_string())
);
}
#[test]
fn grok_quota_feedback_zeros_rate_limited_window() {
let mut bucket = json!({
"quota_by_model": {
"quota_fast": {
"display_name": "fast",
"remaining": 1.0,
"total": 30.0,
"remaining_fraction": 1.0 / 30.0,
"used_percent": 29.0 / 30.0 * 100.0,
"is_exhausted": false
}
}
})
.as_object()
.cloned()
.expect("bucket should be object");
assert!(grok_apply_quota_feedback(
&mut bucket,
"grok-4.20-fast",
429,
None,
1_700_000_000
));
let fast = bucket
.get("quota_by_model")
.and_then(Value::as_object)
.and_then(|models| models.get("quota_fast"))
.and_then(Value::as_object)
.expect("fast window should exist");
assert_eq!(fast.get("remaining"), Some(&json!(0.0)));
assert_eq!(fast.get("remaining_fraction"), Some(&json!(0.0)));
assert_eq!(fast.get("used_percent"), Some(&json!(100.0)));
assert_eq!(fast.get("is_exhausted"), Some(&json!(true)));
}
#[test]
fn grok_quota_feedback_records_reset_after_when_upstream_mentions_wait_time() {
let mut bucket = json!({
"quota_by_model": {
"quota_fast": {
"display_name": "fast",
"remaining": 1.0,
"total": 30.0,
"remaining_fraction": 1.0 / 30.0,
"used_percent": 29.0 / 30.0 * 100.0,
"is_exhausted": false,
"reset_at": 10,
"next_reset_at": 10
}
}
})
.as_object()
.cloned()
.expect("bucket should be object");
let parsed = grok_quota_reset_after_seconds(
Some(&json!({
"error": {
"message": "升级到 SuperGrok 获得更高使用上限,或等待 6小时 13分钟。"
}
})),
None,
);
assert_eq!(parsed, Some(22_380));
assert!(grok_apply_quota_feedback(
&mut bucket,
"grok-4.20-fast",
503,
parsed,
1_700_000_000
));
let fast = bucket
.get("quota_by_model")
.and_then(Value::as_object)
.and_then(|models| models.get("quota_fast"))
.and_then(Value::as_object)
.expect("fast window should exist");
assert_eq!(fast.get("remaining"), Some(&json!(0.0)));
assert_eq!(fast.get("is_exhausted"), Some(&json!(true)));
assert_eq!(fast.get("reset_after_seconds"), Some(&json!(22_380u64)));
assert_eq!(fast.get("reset_at"), Some(&json!(1_700_022_380u64)));
assert_eq!(fast.get("next_reset_at"), Some(&json!(1_700_022_380u64)));
assert_eq!(
fast.get("reset_at_source"),
Some(&json!("grok_upstream_error"))
);
}
#[test]
fn grok_realtime_quota_bucket_updates_observed_timestamp() {
let mut bucket = json!({
"updated_at": 1_600_000_000u64,
"quota_by_model": {
"quota_fast": {
"display_name": "fast",
"remaining": 1.0,
"total": 30.0,
"remaining_fraction": 1.0 / 30.0,
"used_percent": 29.0 / 30.0 * 100.0,
"is_exhausted": false
}
}
})
.as_object()
.cloned()
.expect("bucket should be object");
grok_mark_quota_bucket_updated(&mut bucket, 1_700_000_000);
assert_eq!(bucket.get("updated_at"), Some(&json!(1_700_000_000u64)));
}
#[test]
fn grok_wait_duration_parser_handles_english_and_chinese_messages() {
assert_eq!(
grok_wait_duration_seconds_from_text("wait 6h 13m"),
Some(22_380)
);
assert_eq!(
grok_wait_duration_seconds_from_text("等待 6小时13分钟"),
Some(22_380)
);
assert_eq!(
grok_wait_duration_seconds_from_text("no duration here"),
None
);
}
}

View File

@@ -4,6 +4,7 @@ use super::{
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
VideoTaskStatus, VideoTaskWriteRepository, DEVELOPMENT_ENCRYPTION_KEY,
};
use crate::image_capabilities::openai_image_gateway_max_generation_count;
use crate::tests::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, AppState, Arc, Body, Json, Mutex, Request, Router, StatusCode, EXECUTION_PATH_HEADER,
@@ -582,7 +583,7 @@ async fn gateway_does_not_locally_reject_image_model_name_on_chat_completions()
}
#[tokio::test]
async fn gateway_rejects_image_request_with_n_greater_than_one_without_hitting_fallback_probe() {
async fn gateway_rejects_image_request_with_n_greater_than_four_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
@@ -615,9 +616,9 @@ async fn gateway_rejects_image_request_with_n_greater_than_one_without_hitting_f
.header(http::header::CONTENT_TYPE, "application/json")
.body(
serde_json::to_vec(&json!({
"model": "gpt-image-2",
"model": "grok-imagine-image-lite",
"prompt": "draw",
"n": 2,
"n": 5,
"response_format": "b64_json"
}))
.expect("request body should encode"),
@@ -635,7 +636,13 @@ async fn gateway_rejects_image_request_with_n_greater_than_one_without_hitting_f
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "当前 Codex 图片反代仅支持 n=1");
assert_eq!(
payload["detail"],
format!(
"当前图片反代仅支持 n=1..{}",
openai_image_gateway_max_generation_count()
)
);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -643,7 +650,7 @@ async fn gateway_rejects_image_request_with_n_greater_than_one_without_hitting_f
}
#[tokio::test]
async fn gateway_rejects_variation_request_without_image_without_hitting_fallback_probe() {
async fn gateway_does_not_mount_image_variation_route_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
@@ -685,16 +692,7 @@ async fn gateway_rejects_variation_request_without_image_without_hitting_fallbac
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AI_PUBLIC)
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "图片变体请求需要 image 文件");
assert_eq!(response.status(), StatusCode::NOT_FOUND);
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();

View File

@@ -180,9 +180,10 @@ pub use crate::formats::{
request::{
build_chatgpt_web_image_request_body, build_openai_image_provider_request_body,
default_model_for_openai_image_operation, is_openai_image_stream_request,
normalize_openai_image_request, openai_image_operation_from_path,
resolve_requested_openai_image_model_for_request, ChatGptWebImageRequestError,
NormalizedOpenAiImageRequest, OpenAiImageOperation, OpenAiImageResponseFormat,
normalize_openai_image_request, normalize_openai_image_request_with_options,
openai_image_operation_from_path, resolve_requested_openai_image_model_for_request,
ChatGptWebImageRequestError, NormalizedOpenAiImageRequest, OpenAiImageNormalizeOptions,
OpenAiImageOperation, OpenAiImageResponseFormat,
},
spec::{
resolve_stream_spec as resolve_local_image_stream_spec,

View File

@@ -669,6 +669,12 @@ impl ClaudeClientEmitter {
Ok(out)
}
CanonicalStreamEvent::ContentPart(part) => self.emit_content_part(part),
CanonicalStreamEvent::ImageGenerationCall { item, .. } => {
let Some(part) = content_part_from_openai_image_generation_item(&item) else {
return Ok(Vec::new());
};
self.emit_content_part(part)
}
CanonicalStreamEvent::ToolCallStart {
index,
call_id,

View File

@@ -485,6 +485,16 @@ impl GeminiClientEmitter {
None,
None,
),
CanonicalStreamEvent::ImageGenerationCall { item, .. } => {
let Some(part) = content_part_from_openai_image_generation_item(&item) else {
return Ok(Vec::new());
};
self.emit_candidate(
vec![gemini_part_from_canonical_content_part(part)],
None,
None,
)
}
CanonicalStreamEvent::ToolCallStart {
index,
call_id,

View File

@@ -1,4 +1,4 @@
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use serde_json::{json, Map, Value};
@@ -50,6 +50,7 @@ pub struct OpenAIResponsesProviderState {
tool_calls: BTreeMap<usize, OpenAIResponsesProviderToolState>,
tool_results: BTreeMap<usize, OpenAIResponsesProviderToolResultState>,
tool_index_by_key: BTreeMap<String, usize>,
image_item_keys: BTreeSet<String>,
last_tool_index: Option<usize>,
}
@@ -718,6 +719,55 @@ impl OpenAIResponsesProviderState {
}
}
fn emit_image_generation_item(
&mut self,
report_context: &Value,
out: &mut Vec<CanonicalStreamFrame>,
item: &Map<String, Value>,
output_index: Option<usize>,
final_item: bool,
) {
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
return;
}
if !final_item
&& !item
.get("status")
.and_then(Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("completed"))
{
return;
}
let has_image_payload = item
.get("result")
.or_else(|| item.get("url"))
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty());
if !has_image_payload {
return;
}
let index = output_index.unwrap_or(self.image_item_keys.len());
let key = item
.get("id")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.unwrap_or_else(|| format!("image_generation_call:{index}"));
if !self.image_item_keys.insert(key) {
return;
}
self.ensure_started(report_context, out);
let (id, model) = self.identity(report_context);
out.push(CanonicalStreamFrame {
id,
model,
event: CanonicalStreamEvent::ImageGenerationCall {
index,
item: Value::Object(item.clone()),
},
});
}
pub fn push_line(
&mut self,
report_context: &Value,
@@ -894,6 +944,15 @@ impl OpenAIResponsesProviderState {
"reasoning" => {
self.ensure_started(report_context, &mut out);
}
"image_generation_call" => {
self.emit_image_generation_item(
report_context,
&mut out,
item,
output_index,
false,
);
}
_ => {
out.push(self.unknown_frame(report_context, Value::Object(item.clone())));
}
@@ -1107,6 +1166,15 @@ impl OpenAIResponsesProviderState {
"reasoning" => {
self.emit_reasoning_item(report_context, &mut out, item);
}
"image_generation_call" => {
self.emit_image_generation_item(
report_context,
&mut out,
item,
output_index,
true,
);
}
_ => {
out.push(self.unknown_frame(report_context, Value::Object(item.clone())));
}
@@ -1152,6 +1220,15 @@ impl OpenAIResponsesProviderState {
"reasoning" => {
self.emit_reasoning_item(report_context, &mut out, item);
}
"image_generation_call" => {
self.emit_image_generation_item(
report_context,
&mut out,
item,
Some(output_index),
true,
);
}
_ => {
out.push(
self.unknown_frame(report_context, Value::Object(item.clone())),
@@ -1255,6 +1332,7 @@ pub struct OpenAIResponsesClientEmitter {
reasoning_summary_parts: Vec<String>,
tool_calls: BTreeMap<usize, OpenAIResponsesClientToolState>,
tool_results: BTreeMap<usize, OpenAIResponsesClientToolResultState>,
image_generation_items: BTreeMap<usize, Value>,
}
impl OpenAIChatClientEmitter {
@@ -1361,6 +1439,26 @@ impl OpenAIChatClientEmitter {
)?);
Ok(out)
}
CanonicalStreamEvent::ImageGenerationCall { item, .. } => {
let Some(part) = content_part_from_openai_image_generation_item(&item) else {
return Ok(Vec::new());
};
let placeholder = openai_stream_placeholder_for_content_part(&part);
let mut out = self.ensure_started()?;
out.extend(encode_json_sse(
None,
&build_openai_chat_chunk(
self.response_id
.as_deref()
.unwrap_or("chatcmpl-local-stream"),
self.model.as_deref().unwrap_or("unknown"),
placeholder,
None,
None,
),
)?);
Ok(out)
}
CanonicalStreamEvent::ToolCallStart {
index,
call_id,
@@ -1653,6 +1751,11 @@ impl OpenAIResponsesClientEmitter {
output_index
}
fn ensure_image_generation_output_index(&mut self, index: usize) -> usize {
self.next_output_index = self.next_output_index.max(index.saturating_add(1));
index
}
fn ensure_reasoning_item_started(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = self.ensure_started()?;
let output_index = self.ensure_reasoning_output_index();
@@ -1956,6 +2059,42 @@ impl OpenAIResponsesClientEmitter {
Ok(out)
}
fn emit_image_generation_call_item(
&mut self,
index: usize,
item: Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = self.ensure_started()?;
let output_index = self.ensure_image_generation_output_index(index);
let mut item = item.as_object().cloned().unwrap_or_default();
item.insert(
"type".to_string(),
Value::String("image_generation_call".to_string()),
);
if !item.contains_key("id") {
item.insert(
"id".to_string(),
Value::String(format!("{}_ig_{}", self.response_id(), output_index)),
);
}
if !item.contains_key("status") {
item.insert("status".to_string(), Value::String("completed".to_string()));
}
let item = Value::Object(item);
self.image_generation_items
.insert(output_index, item.clone());
out.extend(self.encode_response_event(
"response.output_item.done",
json!({
"type": "response.output_item.done",
"response_id": self.response_id(),
"output_index": output_index,
"item": item,
}),
)?);
Ok(out)
}
fn completed_response(&self, usage: CanonicalUsage) -> Value {
let mut ordered_output = Vec::new();
let summary = if self.reasoning_summary_parts.is_empty() {
@@ -2058,6 +2197,9 @@ impl OpenAIResponsesClientEmitter {
ordered_output.push((output_index, Value::Object(item)));
}
}
for (output_index, item) in &self.image_generation_items {
ordered_output.push((*output_index, item.clone()));
}
ordered_output.sort_by_key(|(output_index, _)| *output_index);
let mut usage_payload = Map::new();
@@ -2183,6 +2325,9 @@ impl OpenAIResponsesClientEmitter {
)?);
Ok(out)
}
CanonicalStreamEvent::ImageGenerationCall { index, item } => {
self.emit_image_generation_call_item(index, item)
}
CanonicalStreamEvent::ToolCallStart {
index,
call_id,
@@ -2876,6 +3021,134 @@ mod tests {
)));
}
#[test]
fn openai_responses_provider_state_preserves_image_generation_calls() {
let mut state = OpenAIResponsesProviderState::default();
let report_context = json!({});
let frames = state
.push_line(
&report_context,
data_line(json!({
"type": "response.completed",
"response": {
"id": "resp_img_123",
"model": "gpt-image-2",
"output": [{
"id": "ig_123",
"type": "image_generation_call",
"status": "completed",
"output_format": "png",
"result": "aGVsbG8="
}],
"usage": {"input_tokens": 1, "output_tokens": 2, "total_tokens": 3}
}
})),
)
.expect("completed event should parse");
assert!(frames.iter().any(|frame| matches!(
frame.event,
CanonicalStreamEvent::ImageGenerationCall {
index: 0,
ref item,
} if item["type"] == json!("image_generation_call")
&& item["result"] == json!("aGVsbG8=")
)));
}
#[test]
fn openai_responses_provider_state_waits_for_final_image_generation_item() {
let mut state = OpenAIResponsesProviderState::default();
let report_context = json!({});
let added_frames = state
.push_line(
&report_context,
data_line(json!({
"type": "response.output_item.added",
"output_index": 0,
"item": {
"id": "ig_123",
"type": "image_generation_call",
"status": "generating",
"output_format": "png",
"result": "early"
}
})),
)
.expect("added event should parse");
assert!(!added_frames.iter().any(|frame| matches!(
frame.event,
CanonicalStreamEvent::ImageGenerationCall { .. }
)));
let done_frames = state
.push_line(
&report_context,
data_line(json!({
"type": "response.output_item.done",
"output_index": 0,
"item": {
"id": "ig_123",
"type": "image_generation_call",
"status": "completed",
"output_format": "png",
"result": "final"
}
})),
)
.expect("done event should parse");
assert!(done_frames.iter().any(|frame| matches!(
frame.event,
CanonicalStreamEvent::ImageGenerationCall {
index: 0,
ref item,
} if item["status"] == json!("completed") && item["result"] == json!("final")
)));
}
#[test]
fn openai_responses_client_emitter_emits_image_generation_call_events() {
let mut emitter = OpenAIResponsesClientEmitter::default();
let mut bytes = emitter
.emit(CanonicalStreamFrame {
id: "resp_img_123".to_string(),
model: "gpt-image-2".to_string(),
event: CanonicalStreamEvent::ImageGenerationCall {
index: 0,
item: json!({
"id": "ig_123",
"type": "image_generation_call",
"status": "completed",
"output_format": "png",
"result": "aGVsbG8="
}),
},
})
.expect("image event should encode");
bytes.extend(
emitter
.emit(CanonicalStreamFrame {
id: "resp_img_123".to_string(),
model: "gpt-image-2".to_string(),
event: CanonicalStreamEvent::Finish {
finish_reason: Some("stop".to_string()),
usage: None,
},
})
.expect("finish should encode"),
);
let sse = String::from_utf8(bytes).expect("sse should be utf8");
assert!(sse.contains("event: response.output_item.done\n"));
assert!(sse.contains("\"type\":\"image_generation_call\""));
assert!(sse.contains("\"result\":\"aGVsbG8=\""));
assert!(sse.contains("\"output\":["));
assert!(sse.contains("\"id\":\"ig_123\""));
}
#[test]
fn openai_responses_client_emitter_emits_function_call_output_events() {
let mut emitter = OpenAIResponsesClientEmitter::default();

View File

@@ -3,15 +3,12 @@ use std::collections::BTreeMap;
use base64::Engine as _;
use serde_json::{json, Map, Number, Value};
use crate::formats::openai::responses::codex::{
CODEX_OPENAI_IMAGE_DEFAULT_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL,
};
use crate::formats::openai::responses::codex::CODEX_OPENAI_IMAGE_DEFAULT_MODEL;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OpenAiImageOperation {
Generate,
Edit,
Variation,
}
impl OpenAiImageOperation {
@@ -19,7 +16,6 @@ impl OpenAiImageOperation {
match self {
Self::Generate => "generate",
Self::Edit => "edit",
Self::Variation => "variation",
}
}
}
@@ -47,9 +43,31 @@ pub struct NormalizedOpenAiImageRequest {
prompt: Option<String>,
images: Vec<Value>,
tool: Map<String, Value>,
image_count: Option<u64>,
user: Option<String>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct OpenAiImageNormalizeOptions {
max_generation_count: u64,
}
impl Default for OpenAiImageNormalizeOptions {
fn default() -> Self {
Self {
max_generation_count: 1,
}
}
}
impl OpenAiImageNormalizeOptions {
pub fn with_max_generation_count(max_generation_count: u64) -> Self {
Self {
max_generation_count: max_generation_count.max(1),
}
}
}
pub const CHATGPT_WEB_IMAGE_MAX_AREA: u64 = 1_500_000;
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -123,7 +141,6 @@ pub fn build_chatgpt_web_image_request_body(
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(match request.operation {
OpenAiImageOperation::Variation => "Create a faithful variation of the provided image.",
OpenAiImageOperation::Generate | OpenAiImageOperation::Edit => {
"Generate a high quality image."
}
@@ -244,7 +261,6 @@ pub fn openai_image_operation_from_path(path: &str) -> Option<OpenAiImageOperati
match path {
"/v1/images/generations" => Some(OpenAiImageOperation::Generate),
"/v1/images/edits" => Some(OpenAiImageOperation::Edit),
"/v1/images/variations" => Some(OpenAiImageOperation::Variation),
_ => None,
}
}
@@ -435,7 +451,6 @@ pub fn resolve_requested_openai_image_model_for_request(
pub fn default_model_for_openai_image_operation(operation: OpenAiImageOperation) -> &'static str {
match operation {
OpenAiImageOperation::Variation => CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL,
OpenAiImageOperation::Generate | OpenAiImageOperation::Edit => {
CODEX_OPENAI_IMAGE_DEFAULT_MODEL
}
@@ -446,12 +461,26 @@ pub fn normalize_openai_image_request(
parts: &http::request::Parts,
body_json: &Value,
body_base64: Option<&str>,
) -> Option<NormalizedOpenAiImageRequest> {
normalize_openai_image_request_with_options(
parts,
body_json,
body_base64,
OpenAiImageNormalizeOptions::default(),
)
}
pub fn normalize_openai_image_request_with_options(
parts: &http::request::Parts,
body_json: &Value,
body_base64: Option<&str>,
options: OpenAiImageNormalizeOptions,
) -> Option<NormalizedOpenAiImageRequest> {
let operation = openai_image_operation_from_path(parts.uri.path())?;
if let Some(body_base64) = body_base64 {
normalize_openai_image_multipart_request(parts, body_base64, operation)
normalize_openai_image_multipart_request(parts, body_base64, operation, options)
} else {
normalize_openai_image_json_request(body_json, operation)
normalize_openai_image_json_request(body_json, operation, options)
}
}
@@ -486,12 +515,16 @@ pub fn build_openai_image_provider_request_body(request: &NormalizedOpenAiImageR
if let Some(user) = request.user.as_ref() {
body.insert("user".to_string(), Value::String(user.clone()));
}
if let Some(image_count) = request.image_count.filter(|value| *value > 1) {
body.insert("n".to_string(), Value::Number(Number::from(image_count)));
}
Value::Object(body)
}
fn normalize_openai_image_json_request(
body_json: &Value,
operation: OpenAiImageOperation,
options: OpenAiImageNormalizeOptions,
) -> Option<NormalizedOpenAiImageRequest> {
let object = body_json.as_object()?;
if object
@@ -502,10 +535,9 @@ fn normalize_openai_image_json_request(
{
return None;
}
if object
.get("n")
.and_then(image_request_count)
.is_some_and(|value| value != 1)
let image_count = object.get("n").and_then(image_request_count);
if image_count
.is_some_and(|value| value == 0 || value > max_count_for_operation(operation, options))
{
return None;
}
@@ -534,11 +566,7 @@ fn normalize_openai_image_json_request(
}
}
let mask = object.get("mask").and_then(normalize_mask_value);
if matches!(
operation,
OpenAiImageOperation::Edit | OpenAiImageOperation::Variation
) && images.is_empty()
{
if matches!(operation, OpenAiImageOperation::Edit) && images.is_empty() {
return None;
}
@@ -550,6 +578,7 @@ fn normalize_openai_image_json_request(
prompt,
images,
tool,
image_count,
user,
summary_json: build_image_request_summary_json(
operation,
@@ -564,6 +593,7 @@ fn normalize_openai_image_multipart_request(
parts: &http::request::Parts,
body_base64: &str,
operation: OpenAiImageOperation,
options: OpenAiImageNormalizeOptions,
) -> Option<NormalizedOpenAiImageRequest> {
let multipart_fields = parse_multipart_fields_from_base64(parts, body_base64)?;
let requested_model = normalize_requested_image_model(
@@ -572,9 +602,10 @@ fn normalize_openai_image_multipart_request(
if find_multipart_text_field(&multipart_fields, "style").is_some() {
return None;
}
if find_multipart_text_field(&multipart_fields, "n")
.and_then(|value| value.trim().parse::<u64>().ok())
.is_some_and(|value| value != 1)
let image_count = find_multipart_text_field(&multipart_fields, "n")
.and_then(|value| value.trim().parse::<u64>().ok());
if image_count
.is_some_and(|value| value == 0 || value > max_count_for_operation(operation, options))
{
return None;
}
@@ -635,11 +666,7 @@ fn normalize_openai_image_multipart_request(
}
}
if matches!(
operation,
OpenAiImageOperation::Edit | OpenAiImageOperation::Variation
) && images.is_empty()
{
if matches!(operation, OpenAiImageOperation::Edit) && images.is_empty() {
return None;
}
@@ -651,6 +678,7 @@ fn normalize_openai_image_multipart_request(
prompt,
images,
tool,
image_count,
user,
summary_json: build_image_request_summary_json(
operation,
@@ -677,10 +705,8 @@ fn normalize_prompt(
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
match operation {
OpenAiImageOperation::Generate | OpenAiImageOperation::Edit => prompt.map(Some),
OpenAiImageOperation::Variation => Some(prompt),
}
let _ = operation;
Some(prompt)
}
fn normalize_image_response_format(
@@ -804,7 +830,7 @@ fn build_tool_options(
Value::String(
match operation {
OpenAiImageOperation::Generate => "generate",
OpenAiImageOperation::Edit | OpenAiImageOperation::Variation => "edit",
OpenAiImageOperation::Edit => "edit",
}
.to_string(),
),
@@ -892,6 +918,16 @@ fn image_request_count(value: &Value) -> Option<u64> {
})
}
fn max_count_for_operation(
operation: OpenAiImageOperation,
options: OpenAiImageNormalizeOptions,
) -> u64 {
match operation {
OpenAiImageOperation::Generate => options.max_generation_count.max(1),
OpenAiImageOperation::Edit => 1,
}
}
fn normalize_image_value(value: &Value) -> Vec<Value> {
match value {
Value::Array(values) => values.iter().flat_map(normalize_image_value).collect(),
@@ -1101,7 +1137,9 @@ mod tests {
use super::{
build_chatgpt_web_image_request_body, build_openai_image_provider_request_body,
is_openai_image_stream_request, normalize_openai_image_request, OpenAiImageOperation,
is_openai_image_stream_request, normalize_openai_image_request,
normalize_openai_image_request_with_options, openai_image_operation_from_path,
OpenAiImageNormalizeOptions, OpenAiImageOperation,
};
use crate::formats::openai::image::spec::{resolve_stream_spec, resolve_sync_spec};
use crate::formats::openai::responses::codex::{
@@ -1168,7 +1206,7 @@ mod tests {
}
#[test]
fn normalize_variation_multipart_request_leaves_defaults_empty_until_codex_adapter() {
fn openai_image_variation_path_is_not_supported() {
let boundary = "boundary-variation-123";
let body = format!(
concat!(
@@ -1176,9 +1214,6 @@ mod tests {
"Content-Disposition: form-data; name=\"image\"; filename=\"image.png\"\r\n",
"Content-Type: image/png\r\n\r\n",
"hello\r\n",
"--{boundary}\r\n",
"Content-Disposition: form-data; name=\"response_format\"\r\n\r\n",
"url\r\n",
"--{boundary}--\r\n"
),
boundary = boundary,
@@ -1189,40 +1224,8 @@ mod tests {
Some(&format!("multipart/form-data; boundary={boundary}")),
);
let request = normalize_openai_image_request(&parts, &json!({}), Some(&body_base64))
.expect("variation request should normalize");
assert_eq!(request.operation, OpenAiImageOperation::Variation);
assert!(request.requested_model.is_none());
assert_eq!(request.summary_json["response_format"], json!("url"));
assert_eq!(
request.tool.get("action").and_then(|value| value.as_str()),
Some("edit")
);
assert!(request.tool.get("output_format").is_none());
assert_eq!(request.images.len(), 1);
let mut provider_request_body = build_openai_image_provider_request_body(&request);
apply_codex_openai_responses_special_body_edits(
&mut provider_request_body,
"codex",
"openai:image",
None,
None,
);
assert_eq!(
provider_request_body["input"][0]["content"][0]["text"],
json!("Create a faithful variation of the provided image.")
);
assert_eq!(
provider_request_body["model"],
json!(CODEX_OPENAI_IMAGE_INTERNAL_MODEL)
);
assert_eq!(
provider_request_body["tools"][0]["output_format"],
json!("png")
);
assert!(openai_image_operation_from_path("/v1/images/variations").is_none());
assert!(normalize_openai_image_request(&parts, &json!({}), Some(&body_base64)).is_none());
}
#[test]
@@ -1318,6 +1321,60 @@ mod tests {
);
}
#[test]
fn normalize_generate_json_request_keeps_allowed_multi_image_count() {
let parts = request_parts("/v1/images/generations", Some("application/json"));
let request = normalize_openai_image_request_with_options(
&parts,
&json!({
"model": "grok-imagine-image",
"prompt": "generate image",
"n": 4
}),
None,
OpenAiImageNormalizeOptions::with_max_generation_count(4),
)
.expect("grok generation request should allow n up to four");
let provider_request_body = build_openai_image_provider_request_body(&request);
assert_eq!(provider_request_body["n"], json!(4));
}
#[test]
fn normalize_generate_json_request_rejects_multi_image_count_by_default() {
let parts = request_parts("/v1/images/generations", Some("application/json"));
assert!(normalize_openai_image_request(
&parts,
&json!({
"model": "gpt-image-2",
"prompt": "generate image",
"n": 2
}),
None,
)
.is_none());
}
#[test]
fn normalize_edit_request_rejects_multi_image_count_even_with_generation_override() {
let parts = request_parts("/v1/images/edits", Some("application/json"));
assert!(normalize_openai_image_request_with_options(
&parts,
&json!({
"model": "grok-imagine-image-edit",
"prompt": "edit image",
"n": 2,
"image": {
"b64_json": "aGVsbG8=",
"mime_type": "image/png"
}
}),
None,
OpenAiImageNormalizeOptions::with_max_generation_count(4),
)
.is_none());
}
#[test]
fn build_generate_request_defaults_codex_image_tool_and_tool_choice() {
let parts = request_parts("/v1/images/generations", Some("application/json"));

View File

@@ -44,6 +44,18 @@ fn is_openai_image_request(provider_api_format: &str) -> bool {
.eq_ignore_ascii_case("openai:image")
}
fn codex_openai_responses_body_uses_image_generation_tool(
body_object: &serde_json::Map<String, Value>,
) -> bool {
body_object
.get("tools")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
.any(|tool| tool.get("type").and_then(Value::as_str) == Some("image_generation"))
}
fn apply_codex_openai_image_tool_overrides(body_object: &mut serde_json::Map<String, Value>) {
let mut tool = body_object
.get("tools")
@@ -397,7 +409,9 @@ pub fn apply_codex_openai_responses_special_body_edits(
{
body_object.insert("instructions".to_string(), json!(""));
}
if is_openai_image_request(provider_api_format) {
if is_openai_image_request(provider_api_format)
|| codex_openai_responses_body_uses_image_generation_tool(body_object)
{
body_object.insert(
"model".to_string(),
json!(CODEX_OPENAI_IMAGE_INTERNAL_MODEL),
@@ -712,6 +726,39 @@ mod tests {
);
}
#[test]
fn codex_responses_image_tool_edits_force_internal_model_and_tool_defaults() {
let mut provider_request_body = json!({
"model": "gpt-image-2",
"input": "generate image",
"tools": [{
"type": "image_generation"
}]
});
apply_codex_openai_responses_special_body_edits(
&mut provider_request_body,
"codex",
"openai:responses",
None,
None,
);
assert_eq!(
provider_request_body["model"],
json!(CODEX_OPENAI_IMAGE_INTERNAL_MODEL)
);
assert_eq!(provider_request_body["stream"], json!(true));
assert_eq!(
provider_request_body["tools"][0]["type"],
json!("image_generation")
);
assert_eq!(
provider_request_body["tool_choice"]["type"],
json!("image_generation")
);
}
#[test]
fn codex_image_body_edits_preserve_edit_action_without_generate_defaults() {
let mut provider_request_body = json!({

View File

@@ -89,13 +89,37 @@ pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value, _compact: b
for block in &canonical.content {
match block {
CanonicalContentBlock::Text { .. }
| CanonicalContentBlock::Image { .. }
| CanonicalContentBlock::File { .. }
| CanonicalContentBlock::Audio { .. } => {
if let Some(part) = canonical_content_block_to_openai_responses_part(block) {
message_content.push(part);
}
}
CanonicalContentBlock::Image {
data,
url,
media_type,
extensions,
..
} => {
if image_block_is_generation_call(extensions) {
flush_openai_responses_message_item(
&mut output,
&mut message_content,
&response_id,
&mut message_index,
);
output.push(openai_responses_image_generation_call_item(
&response_id,
output.len(),
data,
url,
media_type,
));
} else if let Some(part) = canonical_content_block_to_openai_responses_part(block) {
message_content.push(part);
}
}
CanonicalContentBlock::Thinking {
text,
encrypted_content,
@@ -248,3 +272,56 @@ pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value, _compact: b
));
Value::Object(response)
}
fn image_block_is_generation_call(extensions: &BTreeMap<String, Value>) -> bool {
extensions
.get(OPENAI_RESPONSES_EXTENSION_NAMESPACE)
.or_else(|| extensions.get(OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE))
.and_then(|value| value.get("item_type"))
.and_then(Value::as_str)
.is_some_and(|value| value == "image_generation_call")
}
fn openai_responses_image_generation_call_item(
response_id: &str,
index: usize,
data: &Option<String>,
url: &Option<String>,
media_type: &Option<String>,
) -> Value {
let mut item = Map::new();
item.insert(
"id".to_string(),
Value::String(format!("{response_id}_ig_{index}")),
);
item.insert(
"type".to_string(),
Value::String("image_generation_call".to_string()),
);
item.insert("status".to_string(), Value::String("completed".to_string()));
item.insert("action".to_string(), Value::String("generate".to_string()));
item.insert(
"output_format".to_string(),
Value::String(openai_responses_output_format_from_mime_type(
media_type.as_deref().unwrap_or("image/png"),
)),
);
if let Some(data) = data.as_ref().filter(|value| !value.trim().is_empty()) {
item.insert("result".to_string(), Value::String(data.clone()));
} else if let Some(url) = url.as_ref().filter(|value| !value.trim().is_empty()) {
item.insert("url".to_string(), Value::String(url.clone()));
} else {
item.insert("result".to_string(), Value::String(String::new()));
}
Value::Object(item)
}
fn openai_responses_output_format_from_mime_type(mime_type: &str) -> String {
match mime_type.trim().to_ascii_lowercase().as_str() {
"image/jpeg" | "image/jpg" => "jpeg",
"image/webp" => "webp",
"image/gif" => "gif",
_ => "png",
}
.to_string()
}

View File

@@ -1,6 +1,5 @@
use serde_json::{json, Map, Number, Value};
use crate::formats::openai::responses::codex::CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT;
use crate::formats::shared::model_directives::extract_gemini_model_from_path;
#[derive(Clone, Debug, PartialEq)]
@@ -32,11 +31,6 @@ pub fn build_gemini_image_request_body_from_openai_image_request(
}
let prompt = normalized_request_prompt(normalized_request)
.or_else(|| {
(normalized_request.operation
== crate::formats::openai::image::request::OpenAiImageOperation::Variation)
.then(|| CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT.to_string())
})
.unwrap_or_else(|| "Generate a high quality image.".to_string());
let mut parts = Vec::new();
if !prompt.trim().is_empty() {
@@ -393,20 +387,7 @@ pub fn build_openai_image_response_from_response_stream_sync_body(
let output = provider_body_json.get("output").and_then(Value::as_array)?;
let images = output
.iter()
.filter_map(|item| {
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
return None;
}
let b64_json = item
.get("result")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
Some(json!({
"b64_json": b64_json,
"revised_prompt": item.get("revised_prompt").cloned().unwrap_or(Value::Null),
}))
})
.filter_map(openai_response_image_generation_item_to_image_data)
.collect::<Vec<_>>();
if images.is_empty() {
return None;
@@ -439,6 +420,48 @@ pub fn build_openai_image_response_from_response_stream_sync_body(
Some(Value::Object(response))
}
fn openai_response_image_generation_item_to_image_data(item: &Value) -> Option<Value> {
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
return None;
}
let result = item
.get("result")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let url = item
.get("url")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let mut image = Map::new();
match result {
Some(value) if value.starts_with("data:") => {
let (_, b64_json) = parse_data_url(value)?;
image.insert("b64_json".to_string(), Value::String(b64_json));
}
Some(value) if value.starts_with("http://") || value.starts_with("https://") => {
image.insert("url".to_string(), Value::String(value.to_string()));
}
Some(value) => {
image.insert("b64_json".to_string(), Value::String(value.to_string()));
}
None => {
let url = url?;
if let Some((_, b64_json)) = parse_data_url(url) {
image.insert("b64_json".to_string(), Value::String(b64_json));
} else {
image.insert("url".to_string(), Value::String(url.to_string()));
}
}
}
image.insert(
"revised_prompt".to_string(),
item.get("revised_prompt").cloned().unwrap_or(Value::Null),
);
Some(Value::Object(image))
}
pub fn build_openai_image_provider_body_from_response_stream_sync_body(
provider_body_json: &Value,
report_context: Option<&Value>,
@@ -882,7 +905,9 @@ mod tests {
build_gemini_image_request_body_from_openai_image_request,
build_gemini_image_response_from_openai_image_response,
build_openai_image_request_body_from_gemini_image_request,
build_openai_image_response_from_gemini_response, gemini_request_is_image_generation,
build_openai_image_response_from_gemini_response,
build_openai_image_response_from_response_stream_sync_body,
gemini_request_is_image_generation,
};
use crate::formats::openai::image::request::normalize_openai_image_request;
@@ -1013,6 +1038,29 @@ mod tests {
assert_eq!(converted["usage"]["total_tokens"], 3);
}
#[test]
fn converts_responses_image_generation_url_to_openai_image_url() {
let converted = build_openai_image_response_from_response_stream_sync_body(
&json!({
"created_at": 1776839946,
"model": "gpt-image-2",
"output": [{
"type": "image_generation_call",
"status": "completed",
"url": "https://assets.example/generated.png"
}]
}),
None,
)
.expect("response image output should convert");
assert_eq!(
converted["data"][0]["url"],
"https://assets.example/generated.png"
);
assert!(converted["data"][0].get("b64_json").is_none());
}
#[test]
fn converts_openai_image_response_to_gemini_image_response() {
let converted = build_gemini_image_response_from_openai_image_response(

View File

@@ -193,10 +193,7 @@ pub fn resolve_execution_runtime_sync_plan_kind(
if route_family == Some("openai")
&& route_kind == Some("image")
&& *method == Method::POST
&& matches!(
path,
"/v1/images/generations" | "/v1/images/edits" | "/v1/images/variations"
)
&& matches!(path, "/v1/images/generations" | "/v1/images/edits")
{
return Some(OPENAI_IMAGE_SYNC_PLAN_KIND);
}
@@ -761,7 +758,7 @@ mod tests {
&Method::POST,
"/v1/images/variations",
),
Some(OPENAI_IMAGE_SYNC_PLAN_KIND)
None
);
assert!(supports_sync_execution_decision_kind(
OPENAI_IMAGE_SYNC_PLAN_KIND

View File

@@ -160,6 +160,57 @@ pub fn canonical_usage_from_claude_usage(value: Option<&Value>) -> Option<Canoni
})
}
pub fn content_part_from_openai_image_generation_item(
item: &Value,
) -> Option<CanonicalContentPart> {
let item = item.as_object()?;
let result = item
.get("result")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let url = item
.get("url")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let image = if let Some(result) = result {
if result.starts_with("data:image/")
|| result.starts_with("http://")
|| result.starts_with("https://")
{
result.to_string()
} else {
let mime_type = item
.get("mime_type")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| {
item.get("output_format")
.and_then(Value::as_str)
.map(openai_image_output_format_to_mime_type)
})
.unwrap_or_else(|| "image/png".to_string());
format!("data:{mime_type};base64,{result}")
}
} else {
url?.to_string()
};
Some(CanonicalContentPart::ImageUrl(image))
}
fn openai_image_output_format_to_mime_type(output_format: &str) -> String {
match output_format.trim().to_ascii_lowercase().as_str() {
"jpeg" | "jpg" => "image/jpeg",
"webp" => "image/webp",
"gif" => "image/gif",
_ => "image/png",
}
.to_string()
}
pub fn canonical_usage_from_gemini_usage(value: Option<&Value>) -> Option<CanonicalUsage> {
let usage = value?.as_object()?;
let input_tokens = usage

View File

@@ -22,8 +22,8 @@ use crate::formats::gemini::generate_content::stream::GeminiProviderState;
use crate::formats::shared::model_directives::model_directive_display_model_from_report_context;
use crate::formats::shared::response::remove_empty_pages_from_tool_arguments;
use crate::formats::shared::stream_core::common::{
map_openai_finish_reason_to_gemini, parse_json_arguments_value, CanonicalContentPart,
CanonicalStreamEvent, CanonicalUsage,
content_part_from_openai_image_generation_item, map_openai_finish_reason_to_gemini,
parse_json_arguments_value, CanonicalContentPart, CanonicalStreamEvent, CanonicalUsage,
};
#[derive(Clone, Debug, PartialEq)]
@@ -1577,6 +1577,7 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option<Va
let mut message_states: BTreeMap<usize, OpenAIResponsesSyncMessageState> = BTreeMap::new();
let mut reasoning_states: BTreeMap<usize, OpenAIResponsesSyncReasoningState> = BTreeMap::new();
let mut tool_states: BTreeMap<usize, OpenAIResponsesSyncToolState> = BTreeMap::new();
let mut image_items: BTreeMap<usize, Value> = BTreeMap::new();
let mut item_output_indexes = BTreeMap::<String, usize>::new();
for event in events {
@@ -1734,6 +1735,9 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option<Va
item,
);
}
"image_generation_call" => {
image_items.insert(output_index, Value::Object(item.clone()));
}
_ => {}
}
}
@@ -1833,6 +1837,9 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option<Va
item,
);
}
"image_generation_call" => {
image_items.insert(output_index, Value::Object(item.clone()));
}
_ => {}
}
}
@@ -1874,6 +1881,7 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option<Va
.keys()
.chain(reasoning_states.keys())
.chain(tool_states.keys())
.chain(image_items.keys())
.copied()
.collect::<Vec<_>>();
output_indexes.sort_unstable();
@@ -1897,6 +1905,9 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option<Va
if let Some(state) = tool_states.remove(&output_index) {
output.push(materialize_openai_responses_tool_item(output_index, state));
}
if let Some(item) = image_items.remove(&output_index) {
output.push(item);
}
}
response.insert("output".to_string(), Value::Array(output));
}
@@ -2541,6 +2552,11 @@ pub fn aggregate_gemini_stream_sync_response(body: &[u8]) -> Option<Value> {
CanonicalStreamEvent::ContentPart(part) => {
parts.push(gemini_sync_part_from_canonical_content_part(part));
}
CanonicalStreamEvent::ImageGenerationCall { item, .. } => {
if let Some(part) = content_part_from_openai_image_generation_item(&item) {
parts.push(gemini_sync_part_from_canonical_content_part(part));
}
}
CanonicalStreamEvent::ToolCallStart {
index,
call_id,
@@ -3464,6 +3480,25 @@ mod tests {
assert_eq!(result["output"][0]["content"][0]["text"], "Authoritative");
}
#[test]
fn reconstructs_openai_responses_image_generation_call_from_output_item_done() {
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_image_123\",\"object\":\"response\",\"model\":\"gpt-5.4-mini\",\"status\":\"in_progress\",\"output\":[]}}\n\n",
"event: response.output_item.done\n",
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"status\":\"completed\",\"output_format\":\"png\",\"result\":\"aGVsbG8=\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_image_123\",\"object\":\"response\",\"model\":\"gpt-5.4-mini\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n",
);
let result = aggregate_openai_responses_stream_sync_response(body.as_bytes())
.expect("openai-responses stream should aggregate into a sync body");
assert_eq!(result["output"][0]["type"], "image_generation_call");
assert_eq!(result["output"][0]["result"], "aGVsbG8=");
assert_eq!(result["output"][0]["output_format"], "png");
}
#[test]
fn reconstructs_openai_responses_multi_part_message_content_order() {
let body = concat!(

View File

@@ -1762,6 +1762,11 @@ pub(crate) fn openai_responses_output_to_canonical_blocks(
),
});
}
"image_generation_call" => {
blocks.push(openai_responses_image_generation_call_to_block(
item_object,
)?);
}
"output_text" | "text" | "output_image" | "image_url" | "file" | "input_file"
| "input_audio" => blocks.push(openai_responses_part_to_canonical_block(item)?),
_ => blocks.push(CanonicalContentBlock::Unknown {
@@ -1774,6 +1779,79 @@ pub(crate) fn openai_responses_output_to_canonical_blocks(
Some(blocks)
}
fn openai_responses_image_generation_call_to_block(
item_object: &Map<String, Value>,
) -> Option<CanonicalContentBlock> {
let result = item_object
.get("result")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let url = item_object
.get("url")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let raw_image = result.or(url)?;
let fallback_media_type = item_object
.get("mime_type")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.or_else(|| {
item_object
.get("output_format")
.and_then(Value::as_str)
.map(openai_responses_output_format_to_mime_type)
});
let (media_type, data, url) = if raw_image.starts_with("data:image/") {
split_data_url(Some(raw_image.to_string()), fallback_media_type)
} else if raw_image.starts_with("http://") || raw_image.starts_with("https://") {
(fallback_media_type, None, Some(raw_image.to_string()))
} else if result.is_some() {
(
fallback_media_type.or_else(|| Some("image/png".to_string())),
Some(raw_image.to_string()),
None,
)
} else {
(fallback_media_type, None, Some(raw_image.to_string()))
};
let mut extensions = openai_responses_extensions(
item_object,
&[
"type",
"id",
"status",
"action",
"result",
"url",
"output_format",
"mime_type",
],
);
canonical_extension_object_mut(&mut extensions, OPENAI_RESPONSES_EXTENSION_NAMESPACE).insert(
"item_type".to_string(),
Value::String("image_generation_call".to_string()),
);
Some(CanonicalContentBlock::Image {
data,
url,
media_type,
detail: None,
extensions,
})
}
fn openai_responses_output_format_to_mime_type(output_format: &str) -> String {
match output_format.trim().to_ascii_lowercase().as_str() {
"jpeg" | "jpg" => "image/jpeg",
"webp" => "image/webp",
"gif" => "image/gif",
_ => "image/png",
}
.to_string()
}
pub(crate) fn openai_responses_part_to_canonical_block(
part: &Value,
) -> Option<CanonicalContentBlock> {
@@ -5327,6 +5405,48 @@ mod tests {
assert_eq!(rebuilt["service_tier"], "flex");
}
#[test]
fn openai_responses_image_generation_call_becomes_canonical_image_block() {
let response = json!({
"id": "resp_img",
"model": "gpt-image-2",
"status": "completed",
"output": [{
"id": "ig_1",
"type": "image_generation_call",
"status": "completed",
"output_format": "png",
"result": "aW1hZ2U="
}]
});
let canonical =
from_openai_responses_to_canonical_response(&response).expect("canonical response");
assert!(matches!(
canonical.content[0],
CanonicalContentBlock::Image { ref data, ref media_type, .. }
if data.as_deref() == Some("aW1hZ2U=")
&& media_type.as_deref() == Some("image/png")
));
let rebuilt_chat = canonical_to_openai_chat_response(&canonical);
assert_eq!(
rebuilt_chat["choices"][0]["message"]["content"][0]["type"],
json!("image_url")
);
assert_eq!(
rebuilt_chat["choices"][0]["message"]["content"][0]["image_url"]["url"],
json!("data:image/png;base64,aW1hZ2U=")
);
let rebuilt_responses = canonical_to_openai_responses_response(&canonical, &json!({}));
assert_eq!(
rebuilt_responses["output"][0]["type"],
json!("image_generation_call")
);
assert_eq!(rebuilt_responses["output"][0]["result"], json!("aW1hZ2U="));
}
#[test]
fn claude_request_adapter_preserves_cache_thinking_tools_and_builtin_extensions() {
let request = json!({

View File

@@ -38,6 +38,10 @@ pub enum CanonicalStreamEvent {
ReasoningSummaryDone,
ReasoningSignature(String),
ContentPart(CanonicalContentPart),
ImageGenerationCall {
index: usize,
item: Value,
},
ToolCallStart {
index: usize,
call_id: String,