2026-04-03 14:59:58 +08:00
use std ::collections ::BTreeMap ;
2026-05-09 01:21:26 +08:00
use std ::io ::Error as IoError ;
use std ::sync ::Arc ;
use std ::time ::{ Duration , Instant };
2026-04-03 14:59:58 +08:00
2026-07-27 09:36:31 +08:00
use aether_ai_serving ::{ AiAttemptExecutionOutcome , AiAttemptRetryScope , UPSTREAM_IS_STREAM_KEY };
2026-05-18 00:53:34 +00:00
use aether_contracts ::{
ExecutionError , ExecutionErrorKind , ExecutionPhase , ExecutionPlan , ExecutionResult ,
ExecutionTelemetry ,
};
2026-04-07 02:50:19 +08:00
use aether_data_contracts ::repository ::candidates ::RequestCandidateStatus ;
2026-04-11 01:50:24 +08:00
use aether_scheduler_core ::{
execution_error_details , parse_request_candidate_report_context ,
SchedulerRequestCandidateStatusUpdate ,
};
2026-04-13 14:01:22 +08:00
use aether_usage_runtime ::{
build_lifecycle_usage_seed , build_sync_terminal_usage_payload_seed ,
2026-05-18 14:15:22 +00:00
build_terminal_usage_context_seed , build_usage_event_data_seed , UsageEvent , UsageEventType ,
2026-04-13 14:01:22 +08:00
};
2026-05-09 01:21:26 +08:00
use async_stream ::stream ;
use axum ::body ::{ to_bytes , Body , Bytes };
use axum ::http ::header ::{ CACHE_CONTROL , CONTENT_ENCODING , CONTENT_LENGTH , CONTENT_TYPE };
use axum ::http ::{ HeaderName , HeaderValue , Response , StatusCode };
use futures_util ::StreamExt ;
2026-06-25 22:36:27 +08:00
use http_body_util ::BodyExt ;
2026-05-09 01:21:26 +08:00
use serde_json ::{ json , Value };
use tokio ::sync ::mpsc ;
use tokio ::sync ::Mutex ;
use tokio ::time ::MissedTickBehavior ;
use tracing ::{ debug , warn };
2026-04-03 14:59:58 +08:00
2026-05-02 13:23:54 +08:00
use crate ::ai_serving ::api ::{
2026-05-21 00:57:46 +08:00
build_core_error_body_for_client_format , extract_provider_private_stream_error_body ,
implicit_sync_finalize_report_kind , maybe_build_sync_finalize_outcome , LocalCoreSyncErrorKind ,
LocalCoreSyncFinalizeOutcome ,
2026-04-09 00:10:38 +08:00
};
2026-04-05 20:23:16 +08:00
use crate ::api ::response ::{
2026-04-03 16:26:16 +08:00
attach_control_metadata_headers , build_client_response , build_client_response_from_parts ,
2026-05-09 01:21:26 +08:00
build_client_response_from_parts_with_mutator ,
2026-04-03 16:26:16 +08:00
};
2026-04-10 01:46:14 +08:00
use crate ::clock ::current_unix_ms as current_request_candidate_unix_ms ;
2026-04-05 20:23:16 +08:00
use crate ::control ::GatewayControlDecision ;
2026-05-06 02:29:17 +08:00
use crate ::execution_runtime ::chatgpt_web_image ::maybe_execute_chatgpt_web_image_sync ;
2026-05-16 19:23:23 +08:00
use crate ::execution_runtime ::grok ::maybe_execute_grok_sync ;
2026-05-31 22:13:45 +08:00
use crate ::execution_runtime ::kiro_cache ::{
build_kiro_prompt_cache_profile , compute_kiro_prompt_cache_usage ,
estimate_kiro_prompt_input_tokens , kiro_simulated_cache_enabled_from_provider_config ,
kiro_simulated_cache_enabled_from_report_context , KiroPromptCacheUsage ,
KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD ,
};
2026-04-25 21:29:45 +08:00
use crate ::execution_runtime ::oauth_retry ::refresh_oauth_plan_auth_for_retry ;
2026-04-03 14:59:58 +08:00
#[cfg(test)]
2026-04-05 20:23:16 +08:00
use crate ::execution_runtime ::remote_compat ::post_sync_plan_to_remote_execution_runtime ;
2026-05-20 23:42:12 +08:00
use crate ::execution_runtime ::submission ::{
resolve_local_sync_error_status_code , submit_local_core_error_or_sync_finalize ,
2026-05-09 01:21:26 +08:00
};
use crate ::execution_runtime ::transport ::{
2026-07-27 09:36:31 +08:00
append_upstream_response_body_chunk , build_execution_response_body , build_request_body ,
collect_response_headers , decode_response_body_bytes , execution_response_body_mode ,
format_hyper_error_chain , format_upstream_request_error , format_wreq_upstream_request_error ,
response_body_is_json , send_request , DirectHttpResponse , DirectSyncExecutionRuntime ,
ExecutionRuntimeTransportError ,
2026-05-09 01:21:26 +08:00
};
2026-05-20 23:42:12 +08:00
use crate ::execution_runtime ::windsurf ::maybe_execute_windsurf_sync ;
2026-04-07 02:50:19 +08:00
use crate ::execution_runtime ::{
2026-07-27 09:36:31 +08:00
ai_attempt_retry_scope_from_failure_disposition , analyze_local_candidate_failover_sync ,
apply_endpoint_response_header_rules , attach_provider_response_headers_to_report_context ,
local_failover_response_text , resolve_core_sync_error_finalize_report_kind ,
should_fallback_to_control_sync , should_finalize_sync_response , LocalFailoverDecision ,
2026-04-07 02:50:19 +08:00
};
use crate ::log_ids ::short_request_id ;
2026-04-18 11:35:11 +08:00
use crate ::orchestration ::{
2026-05-11 02:30:45 +08:00
apply_local_execution_effect , build_local_error_flow_metadata , trace_upstream_response_body ,
with_error_flow_report_context , with_upstream_response_report_context ,
LocalAdaptiveRateLimitEffect , LocalAdaptiveSuccessEffect , LocalAttemptFailureEffect ,
LocalExecutionEffect , LocalExecutionEffectContext , LocalHealthFailureEffect ,
LocalHealthSuccessEffect , LocalOAuthInvalidationEffect , LocalPoolErrorEffect ,
2026-04-18 11:35:11 +08:00
};
2026-05-15 01:46:24 +08:00
use crate ::provider_pool_demand ::acquire_provider_pool_in_flight_guard ;
2026-04-07 02:50:19 +08:00
use crate ::request_candidate_runtime ::{
2026-05-09 01:21:26 +08:00
ensure_execution_request_candidate_slot , record_local_request_candidate_extra_data ,
2026-07-01 02:15:51 +08:00
record_local_request_candidate_status , record_local_request_candidate_status_snapshot ,
snapshot_local_request_candidate_status ,
2026-04-03 14:59:58 +08:00
};
2026-07-30 01:03:05 +08:00
use crate ::request_diagnostics ::{
attach_current_request_diagnostics_and_candidate_start_timing_to_report_context ,
attach_request_diagnostics_to_report_context , calibrate_candidate_first_byte_elapsed_ms ,
current_request_diagnostics , RequestDiagnostics ,
};
2026-04-05 20:23:16 +08:00
use crate ::usage ::{ spawn_sync_report , submit_sync_report };
use crate ::video_tasks ::VideoTaskSyncReportMode ;
use crate ::{ usage ::GatewaySyncReportRequest , AppState , GatewayError };
2026-04-03 14:59:58 +08:00
#[path = "execution/policy.rs" ]
mod policy ;
#[path = "execution/response.rs" ]
mod response ;
use policy ::decode_execution_result_body ;
pub ( crate ) use response ::{
maybe_build_local_sync_finalize_response , maybe_build_local_video_error_response ,
maybe_build_local_video_success_outcome , resolve_local_sync_error_background_report_kind ,
2026-04-21 16:19:07 +08:00
resolve_local_sync_success_background_report_kind , LocalVideoSyncSuccessBuild ,
LocalVideoSyncSuccessOutcome ,
2026-04-03 14:59:58 +08:00
};
2026-05-09 01:21:26 +08:00
const OPENAI_IMAGE_SYNC_PLAN_KIND : & str = "openai_image_sync" ;
const OPENAI_IMAGE_SYNC_DEFAULT_TOTAL_TIMEOUT_MS : u64 = 900_000 ;
const SYNC_EXECUTION_IDLE_LOG_INTERVAL : Duration = Duration ::from_secs ( 60 );
const OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_INTERVAL : Duration = Duration ::from_secs ( 15 );
const OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_BYTES : & [ u8 ] = b " \n " ;
const OPENAI_IMAGE_SYNC_PROGRESS_WRITE_INTERVAL : Duration = Duration ::from_secs ( 5 );
2026-07-09 23:10:25 +08:00
const INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE : & str = "Provider returned HTTP 200 but the Gemini response did not contain visible model output; refusing to finalize it as a successful response." ;
2026-05-09 01:21:26 +08:00
2026-07-30 01:03:05 +08:00
fn elapsed_ms_since ( started_at : Instant ) -> u64 {
started_at . elapsed (). as_millis (). min ( u128 ::from ( u64 ::MAX )) as u64
}
fn calibrated_sync_candidate_first_byte_elapsed_ms (
candidate_started_at : Instant ,
result : & ExecutionResult ,
) -> Option < u64 > {
let telemetry = result . telemetry . as_ref () ? ;
calibrate_candidate_first_byte_elapsed_ms (
elapsed_ms_since ( candidate_started_at ),
telemetry . elapsed_ms ,
telemetry . ttfb_ms ,
)
}
2026-05-09 01:21:26 +08:00
#[derive(Debug)]
struct SyncExecutionFailure {
error_type : & 'static str ,
message : String ,
status_code : Option < u16 > ,
latency_ms : Option < u64 > ,
2026-07-27 09:36:31 +08:00
fallback_kind : Option < SyncExecutionFailureFallbackKind > ,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SyncExecutionFailureFallbackKind {
UpstreamResponseTooLarge ,
UpstreamResponseDecode ,
}
impl SyncExecutionFailureFallbackKind {
fn error_type ( self ) -> & 'static str {
match self {
Self ::UpstreamResponseTooLarge => "upstream_response_too_large" ,
Self ::UpstreamResponseDecode => "upstream_response_decode_failed" ,
}
}
fn client_message ( self ) -> & 'static str {
match self {
Self ::UpstreamResponseTooLarge => "Upstream response too large" ,
Self ::UpstreamResponseDecode => "Failed to decode upstream response" ,
}
}
2026-05-09 01:21:26 +08:00
}
2026-05-18 14:15:22 +00:00
struct SyncAttemptTerminalGuard {
state : AppState ,
plan : ExecutionPlan ,
report_context : Option < Value > ,
2026-07-30 01:03:05 +08:00
request_diagnostics : Option < Arc < RequestDiagnostics >> ,
2026-05-18 14:15:22 +00:00
candidate_started_unix_ms : u64 ,
2026-07-30 01:03:05 +08:00
candidate_started_at : Instant ,
2026-05-18 14:15:22 +00:00
armed : bool ,
}
impl SyncAttemptTerminalGuard {
fn new (
state : & AppState ,
plan : & ExecutionPlan ,
report_context : Option < Value > ,
candidate_started_unix_ms : u64 ,
2026-07-30 01:03:05 +08:00
candidate_started_at : Instant ,
2026-05-18 14:15:22 +00:00
) -> Self {
Self {
state : state . clone (),
plan : plan . clone (),
report_context ,
2026-07-30 01:03:05 +08:00
request_diagnostics : current_request_diagnostics (),
2026-05-18 14:15:22 +00:00
candidate_started_unix_ms ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
2026-05-18 14:15:22 +00:00
armed : true ,
}
}
fn disarm ( & mut self ) {
self . armed = false ;
}
async fn fail_and_disarm ( & mut self , error : & GatewayError ) {
if ! self . armed {
return ;
}
self . armed = false ;
record_sync_attempt_forced_terminal_state (
self . state . clone (),
self . plan . clone (),
self . report_context . clone (),
2026-07-30 01:03:05 +08:00
self . request_diagnostics . clone (),
2026-05-18 14:15:22 +00:00
self . candidate_started_unix_ms ,
2026-07-30 01:03:05 +08:00
self . candidate_started_at ,
2026-05-18 14:15:22 +00:00
UsageEventType ::Failed ,
RequestCandidateStatus ::Failed ,
StatusCode ::INTERNAL_SERVER_ERROR . as_u16 (),
"local_sync_attempt_aborted" ,
format! ( "Local sync attempt failed before terminal finalization: {error:?} " ),
)
. await ;
}
}
impl Drop for SyncAttemptTerminalGuard {
fn drop ( & mut self ) {
if ! self . armed {
return ;
}
self . armed = false ;
let state = self . state . clone ();
let plan = self . plan . clone ();
let report_context = self . report_context . clone ();
2026-07-30 01:03:05 +08:00
let request_diagnostics = self . request_diagnostics . clone ();
2026-05-18 14:15:22 +00:00
let candidate_started_unix_ms = self . candidate_started_unix_ms ;
2026-07-30 01:03:05 +08:00
let candidate_started_at = self . candidate_started_at ;
2026-05-18 14:15:22 +00:00
if let Ok ( handle ) = tokio ::runtime ::Handle ::try_current () {
handle . spawn ( async move {
record_sync_attempt_forced_terminal_state (
state ,
plan ,
report_context ,
2026-07-30 01:03:05 +08:00
request_diagnostics ,
2026-05-18 14:15:22 +00:00
candidate_started_unix_ms ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
2026-05-18 14:15:22 +00:00
UsageEventType ::Cancelled ,
RequestCandidateStatus ::Cancelled ,
499 ,
"local_sync_attempt_cancelled" ,
"Local sync attempt was dropped before terminal finalization, usually because the client disconnected or the request task was cancelled." ,
)
. await ;
});
} else {
warn! (
event_name = "local_sync_attempt_terminal_guard_no_runtime" ,
log_type = "ops" ,
request_id = % short_request_id ( self . plan . request_id . as_str ()),
candidate_id = ? self . plan . candidate_id ,
"gateway could not finalize dropped local sync attempt because no Tokio runtime is available"
);
}
}
}
#[allow(clippy::too_many_arguments)]
async fn record_sync_attempt_forced_terminal_state (
state : AppState ,
plan : ExecutionPlan ,
report_context : Option < Value > ,
2026-07-30 01:03:05 +08:00
request_diagnostics : Option < Arc < RequestDiagnostics >> ,
2026-05-18 14:15:22 +00:00
candidate_started_unix_ms : u64 ,
2026-07-30 01:03:05 +08:00
candidate_started_at : Instant ,
2026-05-18 14:15:22 +00:00
usage_event_type : UsageEventType ,
candidate_status : RequestCandidateStatus ,
status_code : u16 ,
error_type : & 'static str ,
error_message : impl Into < String > ,
) {
let error_message = error_message . into ();
2026-07-30 01:03:05 +08:00
let report_context =
attach_request_diagnostics_to_report_context ( report_context , request_diagnostics . as_ref ());
2026-05-18 14:15:22 +00:00
let terminal_unix_ms = current_request_candidate_unix_ms ();
2026-07-30 01:03:05 +08:00
let latency_ms = elapsed_ms_since ( candidate_started_at );
2026-05-18 14:15:22 +00:00
record_local_request_candidate_status (
& state ,
& plan ,
report_context . as_ref (),
SchedulerRequestCandidateStatusUpdate {
status : candidate_status ,
status_code : Some ( status_code ),
error_type : Some ( error_type . to_string ()),
error_message : Some ( error_message . clone ()),
latency_ms : Some ( latency_ms ),
started_at_unix_ms : Some ( candidate_started_unix_ms ),
finished_at_unix_ms : Some ( terminal_unix_ms ),
},
)
. await ;
if ! state . usage_runtime . is_enabled () {
return ;
}
let mut usage_data = build_usage_event_data_seed ( & plan , report_context . as_ref ());
usage_data . status_code = Some ( status_code );
usage_data . error_message = Some ( error_message . clone ());
usage_data . error_category = Some (
match usage_event_type {
UsageEventType ::Cancelled => "cancelled" ,
_ => "server_error" ,
}
. to_string (),
);
usage_data . response_time_ms = Some ( latency_ms );
let error_body = json! ({
"error" : {
"type" : error_type ,
"message" : error_message ,
"code" : status_code
}
});
usage_data . response_headers = Some ( json! ({ "content-type" : "application/json" }));
usage_data . response_body = Some ( error_body . clone ());
usage_data . client_response_headers = Some ( json! ({ "content-type" : "application/json" }));
usage_data . client_response_body = Some ( error_body );
state
. usage_runtime
. record_terminal_event_direct (
2026-07-22 02:09:34 +08:00
state . usage_lifecycle_data_state (). as_ref (),
2026-05-18 14:15:22 +00:00
UsageEvent ::new ( usage_event_type , plan . request_id . clone (), usage_data ),
)
. await ;
}
2026-05-09 01:21:26 +08:00
impl SyncExecutionFailure {
fn from_transport ( err : ExecutionRuntimeTransportError ) -> Self {
2026-07-27 09:36:31 +08:00
let fallback_kind = match & err {
ExecutionRuntimeTransportError ::UpstreamResponseTooLarge { .. } => {
Some ( SyncExecutionFailureFallbackKind ::UpstreamResponseTooLarge )
}
ExecutionRuntimeTransportError ::UpstreamResponseDecode { .. } => {
Some ( SyncExecutionFailureFallbackKind ::UpstreamResponseDecode )
}
_ => None ,
};
2026-05-09 01:21:26 +08:00
Self {
2026-07-27 09:36:31 +08:00
error_type : fallback_kind
. map ( SyncExecutionFailureFallbackKind ::error_type )
. unwrap_or ( "execution_runtime_unavailable" ),
2026-05-09 01:21:26 +08:00
message : err . to_string (),
2026-07-27 09:36:31 +08:00
status_code : fallback_kind . map ( | _ | StatusCode ::BAD_GATEWAY . as_u16 ()),
2026-05-09 01:21:26 +08:00
latency_ms : None ,
2026-07-27 09:36:31 +08:00
fallback_kind ,
2026-05-09 01:21:26 +08:00
}
}
fn image_sync_total_timeout ( timeout_ms : u64 , elapsed_ms : u64 ) -> Self {
Self {
error_type : "image_sync_total_timeout" ,
message : format ! (
"OpenAI image sync execution exceeded total timeout of {timeout_ms}ms"
),
status_code : Some ( StatusCode ::GATEWAY_TIMEOUT . as_u16 ()),
latency_ms : Some ( elapsed_ms ),
2026-07-27 09:36:31 +08:00
fallback_kind : None ,
2026-05-09 01:21:26 +08:00
}
}
}
2026-07-27 09:36:31 +08:00
fn build_sync_execution_failure_fallback_body (
client_api_format : & str ,
kind : SyncExecutionFailureFallbackKind ,
) -> Value {
let message = kind . client_message ();
let error_type = kind . error_type ();
match crate ::ai_serving ::normalize_api_format_alias ( client_api_format ). as_str () {
"claude:messages" => json! ({
"type" : "error" ,
"error" : {
"type" : "upstream_error" ,
"message" : message ,
}
}),
"gemini:generate_content" => json! ({
"error" : {
"code" : StatusCode ::BAD_GATEWAY . as_u16 (),
"message" : message ,
"status" : "BAD_GATEWAY" ,
}
}),
_ => json! ({
"error" : {
"type" : "upstream_error" ,
"message" : message ,
"code" : error_type ,
}
}),
}
}
fn build_sync_execution_failure_fallback_response (
failure : & SyncExecutionFailure ,
plan : & ExecutionPlan ,
trace_id : & str ,
decision : & GatewayControlDecision ,
) -> Result < Option < Response < Body >> , GatewayError > {
let Some ( kind ) = failure . fallback_kind else {
return Ok ( None );
};
let body_json = build_sync_execution_failure_fallback_body ( & plan . client_api_format , kind );
let body_bytes = serde_json ::to_vec ( & body_json )
. map_err ( | error | GatewayError ::Internal ( error . to_string ())) ? ;
let headers = BTreeMap ::from ([
( "content-type" . to_string (), "application/json" . to_string ()),
( "content-length" . to_string (), body_bytes . len (). to_string ()),
]);
let response = build_client_response_from_parts (
StatusCode ::BAD_GATEWAY . as_u16 (),
& headers ,
Body ::from ( body_bytes ),
trace_id ,
Some ( decision ),
) ? ;
attach_control_metadata_headers (
response ,
Some ( plan . request_id . as_str ()),
plan . candidate_id . as_deref (),
)
. map ( Some )
}
fn maybe_store_sync_execution_failure_fallback (
failure : & SyncExecutionFailure ,
plan : & ExecutionPlan ,
trace_id : & str ,
decision : & GatewayControlDecision ,
retry_scope_out : & mut Option <& mut AiAttemptRetryScope > ,
retry_fallback_out : & mut Option <& mut Option < Response < Body >>> ,
) -> Result < (), GatewayError > {
if failure . fallback_kind . is_none () {
return Ok (());
}
if let Some ( retry_scope ) = retry_scope_out . as_deref_mut () {
* retry_scope = AiAttemptRetryScope ::Candidate ;
}
if let Some ( retry_fallback ) = retry_fallback_out . as_deref_mut () {
* retry_fallback =
build_sync_execution_failure_fallback_response ( failure , plan , trace_id , decision ) ? ;
}
Ok (())
}
2026-07-30 01:03:05 +08:00
async fn maybe_build_sync_transport_error_stop_response (
state : & AppState ,
plan : & ExecutionPlan ,
report_context : Option <& Value > ,
trace_id : & str ,
decision : & GatewayControlDecision ,
error_type : & str ,
error_message : & str ,
elapsed_ms : u64 ,
) -> Result < Option < Response < Body >> , GatewayError > {
let analysis = crate ::orchestration ::resolve_local_transport_failover_analysis_for_attempt (
state ,
plan ,
report_context ,
)
. await ;
if ! matches! ( analysis . decision , LocalFailoverDecision ::StopLocalFailover ) {
return Ok ( None );
}
crate ::execution_runtime ::build_transport_error_stop_response (
state ,
plan ,
report_context ,
trace_id ,
decision ,
StatusCode ::BAD_GATEWAY . as_u16 (),
error_type ,
error_message ,
elapsed_ms ,
)
. await
. map ( Some )
}
2026-04-03 14:59:58 +08:00
struct ImplicitSyncFinalizeOutcome {
payload : GatewaySyncReportRequest ,
2026-04-09 00:10:38 +08:00
outcome : LocalCoreSyncFinalizeOutcome ,
2026-04-03 14:59:58 +08:00
}
2026-07-01 02:15:51 +08:00
fn spawn_sync_candidate_status_update (
state : AppState ,
snapshot : crate ::request_candidate_runtime ::LocalRequestCandidateStatusSnapshot ,
status_update : SchedulerRequestCandidateStatusUpdate ,
) {
tokio ::spawn ( async move {
record_local_request_candidate_status_snapshot ( & state , & snapshot , status_update ). await ;
});
}
fn record_sync_response_started (
state : & AppState ,
lifecycle_seed : aether_usage_runtime ::LifecycleUsageSeed ,
request_candidate_status_snapshot : Option <
crate ::request_candidate_runtime ::LocalRequestCandidateStatusSnapshot ,
> ,
candidate_started_unix_ms : u64 ,
status_code : u16 ,
ttfb_ms : u64 ,
) {
state . usage_runtime . record_stream_started_immediate_async (
2026-07-22 02:09:34 +08:00
state . usage_lifecycle_data_state (). as_ref (),
2026-07-01 02:15:51 +08:00
lifecycle_seed ,
status_code ,
Some ( ExecutionTelemetry {
ttfb_ms : Some ( ttfb_ms ),
elapsed_ms : Some ( ttfb_ms ),
upstream_bytes : None ,
}),
);
if let Some ( snapshot ) = request_candidate_status_snapshot {
spawn_sync_candidate_status_update (
state . clone (),
snapshot ,
SchedulerRequestCandidateStatusUpdate {
status : RequestCandidateStatus ::Streaming ,
status_code : Some ( status_code ),
error_type : None ,
error_message : None ,
latency_ms : Some ( ttfb_ms ),
started_at_unix_ms : Some ( candidate_started_unix_ms ),
finished_at_unix_ms : None ,
},
);
}
}
fn record_sync_execution_active (
state : & AppState ,
plan : & ExecutionPlan ,
report_context : Option <& Value > ,
candidate_started_unix_ms : u64 ,
) {
let lifecycle_seed = build_lifecycle_usage_seed ( plan , report_context );
2026-07-22 02:09:34 +08:00
state . usage_runtime . record_sync_active_immediate_async (
state . usage_lifecycle_data_state (). as_ref (),
lifecycle_seed ,
);
2026-07-01 02:15:51 +08:00
if let Some ( snapshot ) = snapshot_local_request_candidate_status ( plan , report_context ) {
spawn_sync_candidate_status_update (
state . clone (),
snapshot ,
SchedulerRequestCandidateStatusUpdate {
status : RequestCandidateStatus ::Streaming ,
status_code : None ,
error_type : None ,
error_message : None ,
latency_ms : None ,
started_at_unix_ms : Some ( candidate_started_unix_ms ),
finished_at_unix_ms : None ,
},
);
}
}
2026-07-22 02:09:34 +08:00
async fn record_sync_terminal_usage (
2026-04-03 14:59:58 +08:00
state : & AppState ,
plan : & ExecutionPlan ,
report_context : Option <& serde_json ::Value > ,
payload : & GatewaySyncReportRequest ,
2026-07-30 01:03:05 +08:00
candidate_started_at : Instant ,
candidate_first_byte_elapsed_ms : Option < u64 > ,
2026-04-03 14:59:58 +08:00
) {
2026-06-24 01:47:26 +08:00
let report_context_with_diagnostics =
2026-07-30 01:03:05 +08:00
attach_current_request_diagnostics_and_candidate_start_timing_to_report_context (
report_context ,
candidate_started_at ,
candidate_first_byte_elapsed_ms ,
);
2026-06-24 01:47:26 +08:00
let context_seed = build_terminal_usage_context_seed (
plan ,
report_context_with_diagnostics . as_ref (). or ( report_context ),
);
2026-04-13 14:01:22 +08:00
let payload_seed = build_sync_terminal_usage_payload_seed ( payload );
2026-04-03 14:59:58 +08:00
state
. usage_runtime
2026-07-22 02:09:34 +08:00
. record_sync_terminal (
state . usage_lifecycle_data_state (). as_ref (),
context_seed ,
payload_seed ,
)
. await ;
2026-04-21 16:19:07 +08:00
}
2026-07-22 02:09:34 +08:00
async fn record_sync_terminal_usage_and_disarm_guard (
2026-05-18 14:15:22 +00:00
state : & AppState ,
plan : & ExecutionPlan ,
report_context : Option <& serde_json ::Value > ,
payload : & GatewaySyncReportRequest ,
2026-07-30 01:03:05 +08:00
candidate_started_at : Instant ,
candidate_first_byte_elapsed_ms : Option < u64 > ,
2026-05-18 14:15:22 +00:00
terminal_guard : & mut SyncAttemptTerminalGuard ,
) {
2026-07-30 01:03:05 +08:00
record_sync_terminal_usage (
state ,
plan ,
report_context ,
payload ,
candidate_started_at ,
candidate_first_byte_elapsed_ms ,
)
. await ;
2026-05-18 14:15:22 +00:00
terminal_guard . disarm ();
}
2026-05-09 10:46:04 +08:00
fn with_sync_error_trace_context (
report_context : Option <& serde_json ::Value > ,
status_code : u16 ,
headers : & BTreeMap < String , String > ,
2026-05-11 02:30:45 +08:00
body_json : Option <& serde_json ::Value > ,
body_bytes : & [ u8 ],
2026-05-09 10:46:04 +08:00
response_text : Option <& str > ,
local_failover_analysis : crate ::orchestration ::LocalFailoverAnalysis ,
) -> Option < serde_json ::Value > {
2026-05-11 02:30:45 +08:00
let body = trace_upstream_response_body ( body_json , body_bytes );
2026-05-09 10:46:04 +08:00
let upstream_context = with_upstream_response_report_context (
report_context ,
status_code ,
Some ( headers ),
2026-05-11 02:30:45 +08:00
body . as_ref (),
2026-05-09 10:46:04 +08:00
None ,
None ,
);
with_error_flow_report_context (
upstream_context . as_ref (). or ( report_context ),
build_local_error_flow_metadata ( status_code , response_text , local_failover_analysis ),
)
}
2026-04-21 16:19:07 +08:00
fn build_sync_report_payload (
trace_id : & str ,
report_kind : String ,
report_context : Option < serde_json ::Value > ,
status_code : u16 ,
headers : BTreeMap < String , String > ,
body_json : Option < serde_json ::Value > ,
body_base64 : Option < String > ,
telemetry : Option < ExecutionTelemetry > ,
) -> GatewaySyncReportRequest {
GatewaySyncReportRequest {
trace_id : trace_id . to_string (),
report_kind ,
report_context ,
status_code ,
headers ,
body_json ,
client_body_json : None ,
body_base64 ,
telemetry ,
}
}
2026-05-31 22:13:45 +08:00
fn seed_kiro_sync_report_context_input_tokens (
plan : & ExecutionPlan ,
report_context : & mut Option < Value > ,
) {
if ! plan
. provider_name
. as_deref ()
. is_some_and ( | provider_name | provider_name . eq_ignore_ascii_case ( "Kiro" ))
{
return ;
}
let Some ( context ) = report_context . as_mut (). and_then ( Value ::as_object_mut ) else {
return ;
};
if context
. get ( "input_tokens" )
. and_then ( Value ::as_u64 )
. is_some_and ( | input_tokens | input_tokens > 0 )
{
return ;
}
let Some ( original_request_body ) = context . get ( "original_request_body" ). cloned () else {
return ;
};
let estimated_input_tokens = estimate_kiro_prompt_input_tokens ( & original_request_body );
context . insert (
"input_tokens" . to_string (),
Value ::from ( estimated_input_tokens ),
);
}
async fn seed_kiro_sync_simulated_cache_enabled (
state : & AppState ,
plan : & ExecutionPlan ,
report_context : & mut Option < Value > ,
) {
if ! plan
. provider_name
. as_deref ()
. is_some_and ( | provider_name | provider_name . eq_ignore_ascii_case ( "Kiro" ))
{
return ;
}
let enabled = match state
. read_provider_catalog_providers_by_ids ( std ::slice ::from_ref ( & plan . provider_id ))
. await
{
Ok ( providers ) => providers
. iter ()
. find ( | provider | provider . id == plan . provider_id )
. filter ( | provider | provider . provider_type . eq_ignore_ascii_case ( "kiro" ))
. is_some_and ( | provider | {
kiro_simulated_cache_enabled_from_provider_config ( provider . config . as_ref ())
}),
Err ( err ) => {
warn! (
event_name = "kiro_simulated_cache_config_read_failed" ,
log_type = "event" ,
request_id = % plan . request_id ,
provider_id = % plan . provider_id ,
error = ? err ,
"failed to read Kiro simulated cache provider config; defaulting disabled"
);
false
}
};
let Some ( context ) = report_context . as_mut (). and_then ( Value ::as_object_mut ) else {
return ;
};
if enabled {
context . insert (
KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD . to_string (),
Value ::Bool ( true ),
);
} else {
context . remove ( KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD );
}
}
async fn seed_kiro_sync_report_context_prompt_cache_usage (
state : & AppState ,
plan : & ExecutionPlan ,
report_context : & mut Option < Value > ,
) {
if ! plan
. provider_name
. as_deref ()
. is_some_and ( | provider_name | provider_name . eq_ignore_ascii_case ( "Kiro" ))
{
return ;
}
let simulated_cache_enabled =
kiro_simulated_cache_enabled_from_report_context ( report_context . as_ref ());
let Some ( context ) = report_context . as_mut (). and_then ( Value ::as_object_mut ) else {
return ;
};
if context
. get ( "kiro_web_search_mcp" )
. and_then ( Value ::as_bool )
. unwrap_or ( false )
{
return ;
}
if ! simulated_cache_enabled {
return ;
}
if kiro_cache_usage_from_context_object ( context ). is_some () {
return ;
}
let Some ( original_request_body ) = context . get ( "original_request_body" ). cloned () else {
return ;
};
let input_tokens = context
. get ( "input_tokens" )
. and_then ( Value ::as_u64 )
. filter ( | value | * value > 0 )
. unwrap_or_else ( || {
let estimated = estimate_kiro_prompt_input_tokens ( & original_request_body );
context . insert ( "input_tokens" . to_string (), Value ::from ( estimated ));
estimated
});
let Some ( profile ) = build_kiro_prompt_cache_profile ( & original_request_body , input_tokens )
else {
return ;
};
let cache_usage = compute_kiro_prompt_cache_usage (
state . runtime_state (),
kiro_sync_cache_credential_id ( plan ),
& profile ,
)
. await ;
if cache_usage . cache_creation_input_tokens == 0 && cache_usage . cache_read_input_tokens == 0 {
return ;
}
context . insert (
"cache_creation_input_tokens" . to_string (),
Value ::from ( cache_usage . cache_creation_input_tokens ),
);
context . insert (
"cache_read_input_tokens" . to_string (),
Value ::from ( cache_usage . cache_read_input_tokens ),
);
}
fn kiro_sync_cache_credential_id ( plan : & ExecutionPlan ) -> String {
format! ( " {} : {} : {} " , plan . provider_id , plan . endpoint_id , plan . key_id )
}
fn kiro_cache_usage_from_context_object (
context : & serde_json ::Map < String , Value > ,
) -> Option < KiroPromptCacheUsage > {
let cache_creation_input_tokens = context
. get ( "cache_creation_input_tokens" )
. and_then ( Value ::as_u64 )
. unwrap_or ( 0 );
let cache_read_input_tokens = context
. get ( "cache_read_input_tokens" )
. and_then ( Value ::as_u64 )
. unwrap_or ( 0 );
( cache_creation_input_tokens > 0 || cache_read_input_tokens > 0 ). then_some (
KiroPromptCacheUsage {
cache_creation_input_tokens ,
cache_read_input_tokens ,
},
)
}
2026-05-18 00:53:34 +00:00
fn invalid_gemini_provider_success_message (
plan : & ExecutionPlan ,
report_context : Option <& Value > ,
status_code : u16 ,
body_json : Option <& Value > ,
) -> Option <& 'static str > {
if status_code >= 400 {
return None ;
}
2026-07-09 23:10:25 +08:00
if ! provider_api_format_is_gemini_generate_content ( plan , report_context ) {
2026-05-18 00:53:34 +00:00
return None ;
}
let body_json = body_json ? ;
if body_json
. as_object ()
. is_some_and ( | object | object . get ( "error" ). is_some_and ( | error | ! error . is_null ()))
{
return None ;
}
2026-05-21 15:38:10 +08:00
let normalized_body_json = report_context
. filter ( | context | {
context
. get ( "has_envelope" )
. and_then ( Value ::as_bool )
. unwrap_or ( false )
})
. and_then ( | context | {
crate ::ai_serving ::normalize_provider_private_response_value ( body_json . clone (), context )
});
let body_json = normalized_body_json . as_ref (). unwrap_or ( body_json );
2026-05-19 10:18:25 +08:00
if crate ::ai_serving ::gemini_generate_content_response_has_visible_output ( body_json ) {
2026-05-18 00:53:34 +00:00
return None ;
}
2026-07-09 23:10:25 +08:00
Some ( INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE )
}
fn invalid_gemini_provider_stream_success_message (
plan : & ExecutionPlan ,
report_context : Option <& Value > ,
status_code : u16 ,
body_json : Option <& Value > ,
body_bytes : & [ u8 ],
has_body_bytes : bool ,
) -> Option <& 'static str > {
if status_code >= 400 || body_json . is_some () || ! has_body_bytes {
return None ;
}
if ! provider_api_format_is_gemini_generate_content ( plan , report_context ) {
return None ;
}
let Some ( body_json ) = crate ::ai_serving ::aggregate_gemini_stream_sync_response ( body_bytes )
else {
return Some ( INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE );
};
if crate ::ai_serving ::gemini_generate_content_response_has_visible_output ( & body_json ) {
return None ;
}
Some ( INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE )
}
fn provider_api_format_is_gemini_generate_content (
plan : & ExecutionPlan ,
report_context : Option <& Value > ,
) -> bool {
let provider_api_format = report_context
. and_then ( | value | value . get ( "provider_api_format" ))
. and_then ( Value ::as_str )
. unwrap_or ( plan . provider_api_format . as_str ());
crate ::ai_serving ::normalize_api_format_alias ( provider_api_format ) == "gemini:generate_content"
}
fn invalid_gemini_provider_success_execution_error ( message : & str ) -> ExecutionError {
ExecutionError {
kind : ExecutionErrorKind ::Upstream5xx ,
phase : ExecutionPhase ::Finalize ,
message : message . to_string (),
upstream_status : Some ( StatusCode ::OK . as_u16 ()),
retryable : true ,
failover_recommended : true ,
}
2026-05-18 00:53:34 +00:00
}
fn build_invalid_provider_success_body (
plan : & ExecutionPlan ,
report_context : Option <& Value > ,
message : & str ,
) -> Option < Value > {
let client_api_format = report_context
. and_then ( | value | value . get ( "client_api_format" ))
. and_then ( Value ::as_str )
. unwrap_or ( plan . client_api_format . as_str ());
build_core_error_body_for_client_format (
client_api_format ,
message ,
Some ( "invalid_provider_success_response" ),
LocalCoreSyncErrorKind ::ServerError ,
)
}
2026-05-20 23:42:12 +08:00
fn provider_private_error_details ( body_json : & Value ) -> ( Option < String > , Option < String > ) {
let body_object = body_json . as_object ();
let error_object = body_object
. and_then ( | object | object . get ( "error" ))
. and_then ( Value ::as_object );
let error_type =
first_non_empty_error_text ( error_object , body_object , & [ "type" , "code" , "status" ]);
let error_message = first_non_empty_error_text (
error_object ,
body_object ,
& [ "message" , "detail" , "reason" , "status" , "type" , "code" ],
);
( error_type , error_message )
}
fn first_non_empty_error_text (
error_object : Option <& serde_json ::Map < String , Value >> ,
body_object : Option <& serde_json ::Map < String , Value >> ,
keys : & [ & str ],
) -> Option < String > {
for object in [ error_object , body_object ]. into_iter (). flatten () {
for key in keys {
let Some ( value ) = object . get ( * key ) else {
continue ;
};
match value {
Value ::String ( text ) if ! text . trim (). is_empty () => {
return Some ( text . trim (). to_string ());
}
Value ::Number ( number ) => return Some ( number . to_string ()),
_ => {}
}
}
}
None
}
2026-05-09 01:21:26 +08:00
#[derive(Debug, Clone)]
struct OpenAiImageSyncProgressSnapshot {
phase : & 'static str ,
upstream_ttfb_ms : Option < u64 > ,
upstream_sse_frame_count : u64 ,
last_upstream_event : Option < String > ,
last_upstream_frame_at_unix_ms : Option < u64 > ,
partial_image_count : u64 ,
last_client_visible_event : Option < String > ,
downstream_heartbeat_count : u64 ,
last_downstream_heartbeat_at_unix_ms : Option < u64 > ,
downstream_heartbeat_interval_ms : Option < u64 > ,
}
struct OpenAiImageSyncProgressRecorder < 'a > {
state : & 'a AppState ,
plan : & 'a ExecutionPlan ,
report_context : Option <& 'a Value > ,
snapshot : Arc < Mutex < OpenAiImageSyncProgressSnapshot >> ,
buffer : Vec < u8 > ,
last_persist_at : Option < Instant > ,
}
#[derive(Clone)]
struct OpenAiImageSyncJsonHeartbeatContext {
state : AppState ,
plan : ExecutionPlan ,
report_context : Option < Value > ,
snapshot : Arc < Mutex < OpenAiImageSyncProgressSnapshot >> ,
started_at : Instant ,
trace_id : String ,
request_id_for_log : String ,
candidate_id : Option < String > ,
}
#[derive(Debug)]
struct OpenAiImageSyncSseFrame {
event_name : String ,
is_partial_image : bool ,
is_completed : bool ,
is_failed : bool ,
client_visible_event : Option <& 'static str > ,
}
impl OpenAiImageSyncProgressSnapshot {
fn new () -> Self {
Self {
phase : "upstream_connecting" ,
upstream_ttfb_ms : None ,
upstream_sse_frame_count : 0 ,
last_upstream_event : None ,
last_upstream_frame_at_unix_ms : None ,
partial_image_count : 0 ,
last_client_visible_event : None ,
downstream_heartbeat_count : 0 ,
last_downstream_heartbeat_at_unix_ms : None ,
downstream_heartbeat_interval_ms : None ,
}
}
fn to_json ( & self ) -> Value {
json! ({
"phase" : self . phase ,
"upstream_ttfb_ms" : self . upstream_ttfb_ms ,
"upstream_sse_frame_count" : self . upstream_sse_frame_count ,
"last_upstream_event" : self . last_upstream_event ,
"last_upstream_frame_at_unix_ms" : self . last_upstream_frame_at_unix_ms ,
"partial_image_count" : self . partial_image_count ,
"last_client_visible_event" : self . last_client_visible_event ,
"downstream_heartbeat_count" : self . downstream_heartbeat_count ,
"last_downstream_heartbeat_at_unix_ms" : self . last_downstream_heartbeat_at_unix_ms ,
"downstream_heartbeat_interval_ms" : self . downstream_heartbeat_interval_ms ,
})
}
}
impl < 'a > OpenAiImageSyncProgressRecorder < 'a > {
fn new (
state : & 'a AppState ,
plan : & 'a ExecutionPlan ,
report_context : Option <& 'a Value > ,
snapshot : Option < Arc < Mutex < OpenAiImageSyncProgressSnapshot >>> ,
) -> Self {
Self {
state ,
plan ,
report_context ,
snapshot : snapshot
. unwrap_or_else ( || Arc ::new ( Mutex ::new ( OpenAiImageSyncProgressSnapshot ::new ()))),
buffer : Vec ::new (),
last_persist_at : None ,
}
}
async fn persist (
& mut self ,
status : RequestCandidateStatus ,
status_code : Option < u16 > ,
latency_ms : Option < u64 > ,
force : bool ,
) {
let now = Instant ::now ();
if ! force
&& self . last_persist_at . is_some_and ( | last | {
now . duration_since ( last ) < OPENAI_IMAGE_SYNC_PROGRESS_WRITE_INTERVAL
})
{
return ;
}
let snapshot = self . snapshot . lock (). await . clone ();
let extra_data = json! ({
"image_progress" : snapshot . to_json (),
});
record_local_request_candidate_extra_data (
self . state ,
self . plan ,
self . report_context ,
status ,
status_code ,
latency_ms ,
extra_data ,
)
. await ;
self . last_persist_at = Some ( now );
}
async fn record_connecting ( & mut self ) {
self . snapshot . lock (). await . phase = "upstream_connecting" ;
self . persist ( RequestCandidateStatus ::Pending , None , None , true )
. await ;
}
async fn record_response_started ( & mut self , status_code : u16 , ttfb_ms : u64 ) {
{
let mut snapshot = self . snapshot . lock (). await ;
snapshot . phase = if status_code >= 400 {
"failed"
} else {
"upstream_streaming"
};
snapshot . upstream_ttfb_ms = Some ( ttfb_ms );
}
self . persist (
if status_code >= 400 {
RequestCandidateStatus ::Failed
} else {
RequestCandidateStatus ::Streaming
},
Some ( status_code ),
Some ( ttfb_ms ),
true ,
)
. await ;
}
async fn observe_chunk ( & mut self , chunk : & [ u8 ], status_code : u16 , elapsed_ms : u64 ) {
if chunk . is_empty () {
return ;
}
self . buffer . extend_from_slice ( chunk );
let mut force_persist = false ;
while let Some ( block_end ) = find_sse_block_end ( & self . buffer ) {
let block = self . buffer . drain ( .. block_end ). collect ::< Vec < _ >> ();
let Some ( frame ) = parse_openai_image_sync_sse_frame ( & block ) else {
continue ;
};
{
let mut snapshot = self . snapshot . lock (). await ;
snapshot . upstream_sse_frame_count =
snapshot . upstream_sse_frame_count . saturating_add ( 1 );
snapshot . last_upstream_event = Some ( frame . event_name );
snapshot . last_upstream_frame_at_unix_ms = Some ( current_request_candidate_unix_ms ());
if frame . is_partial_image {
snapshot . partial_image_count = snapshot . partial_image_count . saturating_add ( 1 );
}
if let Some ( client_visible_event ) = frame . client_visible_event {
snapshot . last_client_visible_event = Some ( client_visible_event . to_string ());
force_persist = true ;
}
if frame . is_failed || status_code >= 400 {
snapshot . phase = "failed" ;
force_persist = true ;
} else if frame . is_completed {
snapshot . phase = "upstream_completed" ;
force_persist = true ;
} else {
snapshot . phase = "upstream_streaming" ;
}
}
}
let phase = self . snapshot . lock (). await . phase ;
self . persist (
if phase == "failed" {
RequestCandidateStatus ::Failed
} else {
RequestCandidateStatus ::Streaming
},
Some ( status_code ),
Some ( elapsed_ms ),
force_persist ,
)
. await ;
}
async fn finish ( & mut self , status_code : u16 , elapsed_ms : u64 ) {
{
let mut snapshot = self . snapshot . lock (). await ;
if status_code >= 400 || snapshot . phase == "failed" {
snapshot . phase = "failed" ;
} else {
snapshot . phase = "upstream_completed" ;
}
}
self . persist (
if status_code >= 400 {
RequestCandidateStatus ::Failed
} else {
RequestCandidateStatus ::Streaming
},
Some ( status_code ),
Some ( elapsed_ms ),
true ,
)
. await ;
}
2026-05-10 02:10:23 +08:00
async fn fail ( & mut self , status_code : Option < u16 > , elapsed_ms : u64 ) {
2026-05-09 01:21:26 +08:00
self . snapshot . lock (). await . phase = "failed" ;
self . persist (
RequestCandidateStatus ::Failed ,
2026-05-10 02:10:23 +08:00
status_code ,
2026-05-09 01:21:26 +08:00
Some ( elapsed_ms ),
true ,
)
. await ;
}
}
impl OpenAiImageSyncJsonHeartbeatContext {
async fn record_heartbeat ( & self , heartbeat_kind : & 'static str , heartbeat_interval : Duration ) {
let now_unix_ms = current_request_candidate_unix_ms ();
let elapsed_ms = self . started_at . elapsed (). as_millis () as u64 ;
let interval_ms = heartbeat_interval
. as_millis ()
. try_into ()
. unwrap_or ( u64 ::MAX );
let ( count , phase , progress_json ) = {
let mut snapshot = self . snapshot . lock (). await ;
snapshot . downstream_heartbeat_count =
snapshot . downstream_heartbeat_count . saturating_add ( 1 );
snapshot . last_downstream_heartbeat_at_unix_ms = Some ( now_unix_ms );
snapshot . downstream_heartbeat_interval_ms = Some ( interval_ms );
(
snapshot . downstream_heartbeat_count ,
snapshot . phase ,
snapshot . to_json (),
)
};
let status = match phase {
"failed" => RequestCandidateStatus ::Failed ,
"upstream_connecting" => RequestCandidateStatus ::Pending ,
_ => RequestCandidateStatus ::Streaming ,
};
record_local_request_candidate_extra_data (
& self . state ,
& self . plan ,
self . report_context . as_ref (),
status ,
None ,
Some ( elapsed_ms ),
json! ({ "image_progress" : progress_json }),
)
. await ;
debug! (
event_name = "openai_image_sync_json_heartbeat_sent" ,
log_type = "event" ,
trace_id = % self . trace_id ,
request_id = % self . request_id_for_log ,
candidate_id = self . candidate_id . as_deref (). unwrap_or ( "-" ),
heartbeat_kind ,
heartbeat_count = count ,
heartbeat_interval_ms = interval_ms ,
elapsed_ms ,
phase ,
"gateway emitted OpenAI image sync JSON whitespace heartbeat"
);
}
}
fn find_sse_block_end ( buffer : & [ u8 ]) -> Option < usize > {
buffer
. windows ( 2 )
. position ( | window | window == b " \n\n " )
. map ( | index | index + 2 )
. or_else ( || {
buffer
. windows ( 4 )
. position ( | window | window == b " \r\n\r\n " )
. map ( | index | index + 4 )
})
}
fn parse_openai_image_sync_sse_frame ( block : & [ u8 ]) -> Option < OpenAiImageSyncSseFrame > {
let text = std ::str ::from_utf8 ( block ). ok () ? . trim ();
if text . is_empty () {
return None ;
}
let mut event_name = None ;
let mut data_lines = Vec ::new ();
for line in text . lines () {
let line = line . trim_end_matches ( '\r' ). trim ();
if line . is_empty () || line . starts_with ( ':' ) {
continue ;
}
if let Some ( value ) = line . strip_prefix ( "event:" ) {
event_name = Some ( value . trim (). to_string ());
} else if let Some ( value ) = line . strip_prefix ( "data:" ) {
data_lines . push ( value . trim (). to_string ());
}
}
let data_text = data_lines . join ( " \n " );
if data_text . trim (). eq ( "[DONE]" ) {
let event_name = event_name . unwrap_or_else ( || "done" . to_string ());
return Some ( OpenAiImageSyncSseFrame {
event_name ,
is_partial_image : false ,
is_completed : true ,
is_failed : false ,
client_visible_event : None ,
});
}
let data_event_name = serde_json ::from_str ::< Value > ( & data_text )
. ok ()
. and_then ( | value | {
value
. get ( "type" )
. and_then ( Value ::as_str )
. map ( str ::to_string )
. or_else ( || {
value
. get ( "error" )
. and_then ( Value ::as_object )
. map ( | _ | "error" . to_string ())
})
});
let event_name = event_name . or ( data_event_name ) ? ;
let is_partial_image = event_name == "response.image_generation_call.partial_image" ;
let is_completed = event_name == "response.completed" ;
let is_failed = event_name == "response.failed"
|| event_name == "response.error"
|| event_name == "error"
|| event_name . ends_with ( ".failed" );
let client_visible_event = if is_partial_image {
Some ( "image_generation.partial_image" )
} else if is_completed {
Some ( "image_generation.completed" )
} else if is_failed {
Some ( "image_generation.failed" )
} else {
None
};
Some ( OpenAiImageSyncSseFrame {
event_name ,
is_partial_image ,
is_completed ,
is_failed ,
client_visible_event ,
})
}
#[allow(clippy::too_many_arguments)]
async fn execute_direct_sync_runtime_candidate (
state : & AppState ,
plan : & ExecutionPlan ,
report_context : Option <& Value > ,
trace_id : & str ,
plan_kind : & str ,
2026-07-01 02:15:51 +08:00
candidate_started_unix_ms : u64 ,
2026-05-09 01:21:26 +08:00
request_id_for_log : & str ,
candidate_id : Option <& str > ,
provider_name : & str ,
endpoint_id : & str ,
key_id : & str ,
model_name : & str ,
candidate_index : & str ,
progress_snapshot : Option < Arc < Mutex < OpenAiImageSyncProgressSnapshot >>> ,
) -> Result < ExecutionResult , SyncExecutionFailure > {
2026-05-20 23:42:12 +08:00
if let Some ( result ) = maybe_execute_windsurf_sync ( state , plan , report_context )
. await
. map_err ( SyncExecutionFailure ::from_transport ) ?
{
return Ok ( result );
}
2026-05-09 01:21:26 +08:00
if ! should_track_openai_image_sync_upstream_sse ( plan_kind , plan , report_context ) {
2026-07-01 02:15:51 +08:00
let state_for_response_started = state . clone ();
let response_started_lifecycle_seed = build_lifecycle_usage_seed ( plan , report_context );
let response_started_candidate_snapshot =
snapshot_local_request_candidate_status ( plan , report_context );
2026-05-09 01:21:26 +08:00
return DirectSyncExecutionRuntime ::new ()
2026-07-01 02:15:51 +08:00
. execute_sync_with_response_started ( plan , move | event | {
record_sync_response_started (
& state_for_response_started ,
response_started_lifecycle_seed ,
response_started_candidate_snapshot ,
candidate_started_unix_ms ,
event . status_code ,
event . ttfb_ms ,
)
})
2026-05-09 01:21:26 +08:00
. await
. map_err ( SyncExecutionFailure ::from_transport );
}
let started_at = Instant ::now ();
let timeout_ms = resolve_openai_image_sync_total_timeout_ms ( plan );
let mut execution = Box ::pin ( execute_openai_image_sync_upstream_sse_candidate (
state ,
plan ,
report_context ,
progress_snapshot . clone (),
));
let mut idle_interval = tokio ::time ::interval ( SYNC_EXECUTION_IDLE_LOG_INTERVAL );
idle_interval . set_missed_tick_behavior ( tokio ::time ::MissedTickBehavior ::Delay );
idle_interval . tick (). await ;
let mut total_timeout = Box ::pin ( tokio ::time ::sleep ( Duration ::from_millis ( timeout_ms )));
loop {
tokio ::select! {
result = execution . as_mut () => {
match result {
Ok ( result ) => return Ok ( result ),
Err ( err ) => {
let elapsed_ms = started_at . elapsed (). as_millis () as u64 ;
2026-05-10 02:10:23 +08:00
let status_code = err . status_code ;
2026-05-09 01:21:26 +08:00
record_openai_image_sync_failed_progress (
state ,
plan ,
report_context ,
2026-05-10 02:10:23 +08:00
status_code ,
2026-05-09 01:21:26 +08:00
elapsed_ms ,
progress_snapshot . clone (),
)
. await ;
return Err ( err );
}
}
}
_ = idle_interval . tick () => {
warn! (
event_name = "openai_image_sync_execution_idle" ,
log_type = "ops" ,
trace_id = % trace_id ,
request_id = % request_id_for_log ,
candidate_id = candidate_id . unwrap_or ( "-" ),
provider_name ,
endpoint_id ,
key_id ,
model_name ,
candidate_index ,
elapsed_ms = started_at . elapsed (). as_millis () as u64 ,
timeout_ms ,
"gateway OpenAI image sync execution still waiting for upstream response"
);
}
_ = total_timeout . as_mut () => {
let elapsed_ms = started_at . elapsed (). as_millis () as u64 ;
record_openai_image_sync_failed_progress (
state ,
plan ,
report_context ,
2026-05-10 02:10:23 +08:00
Some ( StatusCode ::GATEWAY_TIMEOUT . as_u16 ()),
2026-05-09 01:21:26 +08:00
elapsed_ms ,
progress_snapshot . clone (),
)
. await ;
warn! (
event_name = "openai_image_sync_total_timeout" ,
log_type = "ops" ,
trace_id = % trace_id ,
request_id = % request_id_for_log ,
candidate_id = candidate_id . unwrap_or ( "-" ),
provider_name ,
endpoint_id ,
key_id ,
model_name ,
candidate_index ,
elapsed_ms ,
timeout_ms ,
"gateway OpenAI image sync execution exceeded total timeout"
);
return Err ( SyncExecutionFailure ::image_sync_total_timeout (
timeout_ms ,
elapsed_ms ,
));
}
}
}
}
async fn execute_openai_image_sync_upstream_sse_candidate (
state : & AppState ,
plan : & ExecutionPlan ,
report_context : Option <& Value > ,
progress_snapshot : Option < Arc < Mutex < OpenAiImageSyncProgressSnapshot >>> ,
) -> Result < ExecutionResult , SyncExecutionFailure > {
let request_body = build_request_body ( plan ). map_err ( SyncExecutionFailure ::from_transport ) ? ;
let started_at = Instant ::now ();
let mut progress =
OpenAiImageSyncProgressRecorder ::new ( state , plan , report_context , progress_snapshot );
progress . record_connecting (). await ;
let response = send_request ( plan , request_body )
. await
. map_err ( SyncExecutionFailure ::from_transport ) ? ;
let ttfb_ms = started_at . elapsed (). as_millis () as u64 ;
2026-05-16 19:23:23 +08:00
let status_code = response . status_code ();
let headers = response . headers ();
2026-05-09 01:21:26 +08:00
progress . record_response_started ( status_code , ttfb_ms ). await ;
let mut body_bytes = Vec ::new ();
2026-05-16 19:23:23 +08:00
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 ),
),
)
}) ? ;
2026-07-27 09:36:31 +08:00
append_upstream_response_body_chunk ( & mut body_bytes , & chunk )
. map_err ( SyncExecutionFailure ::from_transport ) ? ;
2026-05-16 19:23:23 +08:00
let elapsed_ms = started_at . elapsed (). as_millis () as u64 ;
progress
. observe_chunk ( & chunk , status_code , elapsed_ms )
. await ;
}
}
2026-06-25 22:36:27 +08:00
DirectHttpResponse ::HyperH2c ( response ) => {
let mut upstream_stream = response . into_body (). into_data_stream ();
while let Some ( chunk ) = upstream_stream . next (). await {
let chunk = chunk . map_err ( | err | {
SyncExecutionFailure ::from_transport (
ExecutionRuntimeTransportError ::UpstreamRequest ( format_hyper_error_chain (
& err ,
)),
)
}) ? ;
2026-07-27 09:36:31 +08:00
append_upstream_response_body_chunk ( & mut body_bytes , & chunk )
. map_err ( SyncExecutionFailure ::from_transport ) ? ;
2026-06-25 22:36:27 +08:00
let elapsed_ms = started_at . elapsed (). as_millis () as u64 ;
progress
. observe_chunk ( & chunk , status_code , elapsed_ms )
. await ;
}
}
2026-05-16 19:23:23 +08:00
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 ),
),
)
}) ? ;
2026-07-27 09:36:31 +08:00
append_upstream_response_body_chunk ( & mut body_bytes , & chunk )
. map_err ( SyncExecutionFailure ::from_transport ) ? ;
2026-05-16 19:23:23 +08:00
let elapsed_ms = started_at . elapsed (). as_millis () as u64 ;
progress
. observe_chunk ( & chunk , status_code , elapsed_ms )
. await ;
}
}
2026-05-09 01:21:26 +08:00
}
2026-07-27 09:36:31 +08:00
let decoded_body_bytes = decode_response_body_bytes ( & headers , & body_bytes )
. map_err ( SyncExecutionFailure ::from_transport ) ? ;
2026-05-09 01:21:26 +08:00
let elapsed_ms = started_at . elapsed (). as_millis () as u64 ;
let upstream_bytes = body_bytes . len () as u64 ;
progress . finish ( status_code , elapsed_ms ). await ;
2026-07-27 09:36:31 +08:00
let body = build_execution_response_body (
& headers ,
& body_bytes ,
decoded_body_bytes . as_ref (),
plan . stream ,
execution_response_body_mode ( plan ),
)
. map_err ( SyncExecutionFailure ::from_transport ) ? ;
2026-05-09 01:21:26 +08:00
Ok ( ExecutionResult {
request_id : plan . request_id . clone (),
candidate_id : plan . candidate_id . clone (),
status_code ,
headers ,
body ,
telemetry : Some ( ExecutionTelemetry {
ttfb_ms : Some ( ttfb_ms ),
elapsed_ms : Some ( elapsed_ms ),
upstream_bytes : Some ( upstream_bytes ),
}),
error : None ,
})
}
async fn record_openai_image_sync_failed_progress (
state : & AppState ,
plan : & ExecutionPlan ,
report_context : Option <& Value > ,
2026-05-10 02:10:23 +08:00
status_code : Option < u16 > ,
2026-05-09 01:21:26 +08:00
elapsed_ms : u64 ,
progress_snapshot : Option < Arc < Mutex < OpenAiImageSyncProgressSnapshot >>> ,
) {
let mut progress =
OpenAiImageSyncProgressRecorder ::new ( state , plan , report_context , progress_snapshot );
2026-05-10 02:10:23 +08:00
progress . fail ( status_code , elapsed_ms ). await ;
2026-05-09 01:21:26 +08:00
}
fn resolve_openai_image_sync_total_timeout_ms ( plan : & ExecutionPlan ) -> u64 {
plan . timeouts
. as_ref ()
. and_then ( | timeouts | timeouts . total_ms )
. unwrap_or ( OPENAI_IMAGE_SYNC_DEFAULT_TOTAL_TIMEOUT_MS )
. max ( 1 )
}
fn report_context_upstream_is_stream ( report_context : Option <& Value > ) -> bool {
report_context
2026-05-21 14:34:51 +08:00
. and_then ( | value | value . get ( UPSTREAM_IS_STREAM_KEY ))
2026-05-09 01:21:26 +08:00
. and_then ( Value ::as_bool )
. unwrap_or ( false )
}
fn should_track_openai_image_sync_upstream_sse (
plan_kind : & str ,
plan : & ExecutionPlan ,
report_context : Option <& Value > ,
) -> bool {
plan_kind == OPENAI_IMAGE_SYNC_PLAN_KIND
&& ( plan . stream || report_context_upstream_is_stream ( report_context ))
}
fn should_enable_openai_image_sync_json_heartbeat (
2026-05-10 02:10:23 +08:00
_plan_kind : & str ,
_plan : & ExecutionPlan ,
_report_context : Option <& Value > ,
2026-05-09 01:21:26 +08:00
) -> bool {
2026-05-10 02:10:23 +08:00
false
2026-05-09 01:21:26 +08:00
}
#[allow(clippy::too_many_arguments)]
fn build_openai_image_sync_json_heartbeat_response (
state : AppState ,
request_path : String ,
plan : ExecutionPlan ,
trace_id : String ,
decision : GatewayControlDecision ,
plan_kind : String ,
report_kind : Option < String > ,
report_context : Option < Value > ,
) -> Result < Response < Body > , GatewayError > {
let request_id = plan . request_id . clone ();
let candidate_id = plan . candidate_id . clone ();
let trace_id_for_response = trace_id . clone ();
let decision_for_response = decision . clone ();
let progress_snapshot = Arc ::new ( Mutex ::new ( OpenAiImageSyncProgressSnapshot ::new ()));
let heartbeat_context = OpenAiImageSyncJsonHeartbeatContext {
state : state . clone (),
plan : plan . clone (),
report_context : report_context . clone (),
snapshot : progress_snapshot . clone (),
started_at : Instant ::now (),
trace_id : trace_id . clone (),
request_id_for_log : short_request_id ( request_id . as_str ()),
candidate_id : candidate_id . clone (),
};
let ( tx , rx ) = mpsc ::channel ::< Result < Bytes , IoError >> ( 1 );
tokio ::spawn ( async move {
let bytes = openai_image_sync_json_heartbeat_final_bytes (
execute_execution_runtime_sync_impl (
& state ,
request_path . as_str (),
plan ,
trace_id . as_str (),
& decision ,
plan_kind . as_str (),
report_kind ,
report_context ,
false ,
Some ( progress_snapshot ),
2026-07-27 09:36:31 +08:00
None ,
None ,
2026-05-09 01:21:26 +08:00
)
. await ,
)
. await ;
let _ = tx . send ( Ok ( Bytes ::from ( bytes ))). await ;
});
let headers = BTreeMap ::from ([(
CONTENT_TYPE . as_str (). to_string (),
"application/json" . to_string (),
)]);
let response = build_client_response_from_parts_with_mutator (
StatusCode ::OK . as_u16 (),
& headers ,
Body ::from_stream ( build_json_whitespace_heartbeat_stream (
rx ,
OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_INTERVAL ,
Some ( heartbeat_context ),
)),
trace_id_for_response . as_str (),
Some ( & decision_for_response ),
| headers | {
headers . remove ( CONTENT_LENGTH );
headers . remove ( CONTENT_ENCODING );
headers . insert (
CACHE_CONTROL ,
HeaderValue ::from_static ( "no-cache, no-transform" ),
);
headers . insert (
HeaderName ::from_static ( "x-accel-buffering" ),
HeaderValue ::from_static ( "no" ),
);
Ok (())
},
) ? ;
attach_control_metadata_headers ( response , Some ( request_id . as_str ()), candidate_id . as_deref ())
}
fn build_json_whitespace_heartbeat_stream (
mut rx : mpsc ::Receiver < Result < Bytes , IoError >> ,
heartbeat_interval : Duration ,
heartbeat_context : Option < OpenAiImageSyncJsonHeartbeatContext > ,
) -> impl futures_util ::Stream < Item = Result < Bytes , IoError >> + Send + 'static {
stream! {
if let Some ( context ) = heartbeat_context . as_ref () {
context . record_heartbeat ( "initial" , heartbeat_interval ). await ;
}
yield Ok ( Bytes ::from_static ( OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_BYTES ));
let mut heartbeat = tokio ::time ::interval ( heartbeat_interval );
heartbeat . set_missed_tick_behavior ( MissedTickBehavior ::Delay );
heartbeat . tick (). await ;
loop {
tokio ::select! {
biased ;
item = rx . recv () => {
let Some ( item ) = item else {
break ;
};
yield item ;
break ;
}
_ = heartbeat . tick () => {
if let Some ( context ) = heartbeat_context . as_ref () {
context . record_heartbeat ( "interval" , heartbeat_interval ). await ;
}
yield Ok ( Bytes ::from_static ( OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_BYTES ));
}
}
}
}
}
2026-05-29 14:35:16 +08:00
pub ( crate ) fn build_sync_json_whitespace_heartbeat_stream (
2026-05-10 12:40:52 +08:00
rx : mpsc ::Receiver < Result < Bytes , IoError >> ,
) -> impl futures_util ::Stream < Item = Result < Bytes , IoError >> + Send + 'static {
build_json_whitespace_heartbeat_stream ( rx , OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_INTERVAL , None )
}
2026-05-29 14:35:16 +08:00
pub ( crate ) fn build_openai_image_sync_json_whitespace_heartbeat_stream (
rx : mpsc ::Receiver < Result < Bytes , IoError >> ,
) -> impl futures_util ::Stream < Item = Result < Bytes , IoError >> + Send + 'static {
build_sync_json_whitespace_heartbeat_stream ( rx )
}
2026-05-09 01:21:26 +08:00
async fn openai_image_sync_json_heartbeat_final_bytes (
result : Result < Option < Response < Body >> , GatewayError > ,
) -> Vec < u8 > {
match result {
2026-07-12 14:10:54 +08:00
Ok ( Some ( response )) => match to_bytes (
response . into_body (),
crate ::headers ::max_internal_buffered_body_bytes (),
)
. await
{
2026-05-09 01:21:26 +08:00
Ok ( bytes ) if ! bytes . is_empty () => bytes . to_vec (),
Ok ( _ ) => openai_image_sync_json_heartbeat_error_body ( "empty sync image response" ),
Err ( err ) => openai_image_sync_json_heartbeat_error_body ( & err . to_string ()),
},
Ok ( None ) => openai_image_sync_json_heartbeat_error_body (
"sync image execution ended without a local response" ,
),
Err ( err ) => openai_image_sync_json_heartbeat_error_body ( & format! ( " {err:?} " )),
}
}
fn openai_image_sync_json_heartbeat_error_body ( message : & str ) -> Vec < u8 > {
serde_json ::to_vec ( & json! ({
"error" : {
"type" : "aether_gateway_error" ,
"message" : message ,
}
}))
. unwrap_or_else ( | _ | b "{ \" error \" :{ \" type \" : \" aether_gateway_error \" }}" . to_vec ())
}
2026-04-21 16:19:07 +08:00
async fn apply_sync_success_effects (
state : & AppState ,
plan : & ExecutionPlan ,
report_context : Option <& serde_json ::Value > ,
payload : & GatewaySyncReportRequest ,
) {
apply_local_execution_effect (
state ,
LocalExecutionEffectContext {
plan ,
report_context ,
},
LocalExecutionEffect ::HealthSuccess ( LocalHealthSuccessEffect ),
)
. await ;
apply_local_execution_effect (
state ,
LocalExecutionEffectContext {
plan ,
report_context ,
},
LocalExecutionEffect ::AdaptiveSuccess ( LocalAdaptiveSuccessEffect ),
)
. await ;
apply_local_execution_effect (
state ,
LocalExecutionEffectContext {
plan ,
report_context ,
},
LocalExecutionEffect ::PoolSuccessSync { payload },
)
. await ;
2026-04-03 14:59:58 +08:00
}
#[cfg(test)]
enum RemoteSyncFallbackOutcome {
Executed ( ExecutionResult ),
ClientResponse ( Response < Body > ),
Unavailable ,
}
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
pub ( crate ) async fn execute_execution_runtime_sync (
state : & AppState ,
request_path : & str ,
mut plan : ExecutionPlan ,
trace_id : & str ,
decision : & GatewayControlDecision ,
plan_kind : & str ,
report_kind : Option < String > ,
mut report_context : Option < serde_json ::Value > ,
) -> Result < Option < Response < Body >> , GatewayError > {
2026-05-09 01:21:26 +08:00
execute_execution_runtime_sync_impl (
state ,
request_path ,
plan ,
trace_id ,
decision ,
plan_kind ,
report_kind ,
report_context ,
true ,
None ,
2026-07-27 09:36:31 +08:00
None ,
None ,
2026-05-09 01:21:26 +08:00
)
. await
}
2026-07-27 09:36:31 +08:00
#[allow(clippy::too_many_arguments)]
pub ( crate ) async fn execute_execution_runtime_sync_with_retry_scope (
state : & AppState ,
request_path : & str ,
plan : ExecutionPlan ,
trace_id : & str ,
decision : & GatewayControlDecision ,
plan_kind : & str ,
report_kind : Option < String > ,
report_context : Option < serde_json ::Value > ,
) -> Result < AiAttemptExecutionOutcome < Response < Body >> , GatewayError > {
let mut retry_scope = AiAttemptRetryScope ::Candidate ;
let mut fallback_response = None ;
let response = execute_execution_runtime_sync_impl (
state ,
request_path ,
plan ,
trace_id ,
decision ,
plan_kind ,
report_kind ,
report_context ,
true ,
None ,
Some ( & mut retry_scope ),
Some ( & mut fallback_response ),
)
. await ? ;
Ok ( match response {
Some ( response ) => AiAttemptExecutionOutcome ::Responded ( response ),
None => AiAttemptExecutionOutcome ::Retry {
scope : retry_scope ,
fallback_response ,
},
})
}
2026-05-09 01:21:26 +08:00
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
async fn execute_execution_runtime_sync_impl (
state : & AppState ,
request_path : & str ,
mut plan : ExecutionPlan ,
trace_id : & str ,
decision : & GatewayControlDecision ,
plan_kind : & str ,
report_kind : Option < String > ,
mut report_context : Option < serde_json ::Value > ,
allow_json_heartbeat : bool ,
progress_snapshot : Option < Arc < Mutex < OpenAiImageSyncProgressSnapshot >>> ,
2026-07-27 09:36:31 +08:00
mut retry_scope_out : Option <& mut AiAttemptRetryScope > ,
mut retry_fallback_out : Option <& mut Option < Response < Body >>> ,
2026-05-09 01:21:26 +08:00
) -> Result < Option < Response < Body >> , GatewayError > {
if allow_json_heartbeat
&& should_enable_openai_image_sync_json_heartbeat ( plan_kind , & plan , report_context . as_ref ())
{
return build_openai_image_sync_json_heartbeat_response (
state . clone (),
request_path . to_string (),
plan ,
trace_id . to_string (),
decision . clone (),
plan_kind . to_string (),
report_kind ,
report_context ,
)
. map ( Some );
}
2026-04-03 14:59:58 +08:00
ensure_execution_request_candidate_slot ( state , & mut plan , & mut report_context ). await ;
2026-04-25 21:29:45 +08:00
let plan_request_id = plan . request_id . clone ();
let plan_request_id_for_log = short_request_id ( plan_request_id . as_str ());
let plan_candidate_id = plan . candidate_id . clone ();
let provider_name = plan
. provider_name
. clone ()
. unwrap_or_else ( || "-" . to_string ());
let endpoint_id = plan . endpoint_id . clone ();
let key_id = plan . key_id . clone ();
let model_name = plan . model_name . clone (). unwrap_or_else ( || "-" . to_string ());
2026-04-11 01:50:24 +08:00
let candidate_index = parse_request_candidate_report_context ( report_context . as_ref ())
. and_then ( | context | context . candidate_index )
. map ( | value | value . to_string ())
. unwrap_or_else ( || "-" . to_string ());
2026-07-30 01:03:05 +08:00
let candidate_started_at = Instant ::now ();
2026-04-10 01:46:14 +08:00
let candidate_started_unix_secs = current_request_candidate_unix_ms ();
2026-04-13 14:01:22 +08:00
let lifecycle_seed = build_lifecycle_usage_seed ( & plan , report_context . as_ref ());
2026-07-22 02:09:34 +08:00
let usage_data = state . usage_lifecycle_data_state (). as_ref (). clone ();
2026-04-10 20:58:01 +08:00
state
. usage_runtime
2026-06-30 16:54:58 +08:00
. record_pending_direct ( & usage_data , lifecycle_seed )
2026-06-25 22:36:27 +08:00
. await ;
2026-04-14 01:27:04 +08:00
record_local_request_candidate_status (
state ,
& plan ,
report_context . as_ref (),
SchedulerRequestCandidateStatusUpdate {
status : RequestCandidateStatus ::Pending ,
status_code : None ,
error_type : None ,
error_message : None ,
latency_ms : None ,
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : None ,
},
)
. await ;
2026-05-18 14:15:22 +00:00
let mut terminal_guard = SyncAttemptTerminalGuard ::new (
state ,
& plan ,
report_context . clone (),
candidate_started_unix_secs ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
2026-05-18 14:15:22 +00:00
);
let result = ( async {
2026-05-15 01:46:24 +08:00
let _provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard (
state . runtime_state . clone (),
& plan . provider_id ,
plan_request_id . as_str (),
plan_candidate_id . as_deref (),
key_id . as_str (),
)
. await ;
2026-07-01 02:15:51 +08:00
record_sync_execution_active (
state ,
& plan ,
report_context . as_ref (),
candidate_started_unix_secs ,
);
2026-04-03 14:59:58 +08:00
#[cfg(not(test))]
2026-04-21 16:19:07 +08:00
let mut result = {
2026-05-16 19:23:23 +08:00
match maybe_execute_grok_sync ( & plan , report_context . as_ref ()). await {
2026-05-06 02:29:17 +08:00
Ok ( Some ( result )) => result ,
2026-05-16 19:23:23 +08:00
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 (
2026-05-06 02:29:17 +08:00
state ,
& plan ,
report_context . as_ref (),
2026-05-16 19:23:23 +08:00
trace_id ,
plan_kind ,
2026-07-01 02:15:51 +08:00
candidate_started_unix_secs ,
2026-05-16 19:23:23 +08:00
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 (),
2026-05-06 02:29:17 +08:00
)
2026-05-16 19:23:23 +08:00
. await
{
Ok ( result ) => result ,
Err ( err ) => {
2026-07-30 01:03:05 +08:00
let failure_error_type = err . error_type ;
let failure_message = err . message . clone ();
let failure_latency_ms = err
. latency_ms
. unwrap_or_else ( || elapsed_ms_since ( candidate_started_at ));
2026-07-27 09:36:31 +08:00
maybe_store_sync_execution_failure_fallback (
& err ,
& plan ,
trace_id ,
decision ,
& mut retry_scope_out ,
& mut retry_fallback_out ,
) ? ;
2026-05-16 19:23:23 +08:00
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 ,
2026-07-30 01:03:05 +08:00
status_code : None ,
error_type : Some ( failure_error_type . to_string ()),
2026-05-16 19:23:23 +08:00
error_message : Some ( err . message ),
2026-07-30 01:03:05 +08:00
latency_ms : Some ( failure_latency_ms ),
2026-05-16 19:23:23 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
},
)
. await ;
2026-07-30 01:03:05 +08:00
if let Some ( response ) = maybe_build_sync_transport_error_stop_response (
state ,
& plan ,
report_context . as_ref (),
trace_id ,
decision ,
failure_error_type ,
failure_message . as_str (),
failure_latency_ms ,
)
. await ?
{
return Ok ( Some ( response ));
}
2026-05-16 19:23:23 +08:00
return Ok ( None );
}
},
Err ( err ) => {
2026-07-30 01:03:05 +08:00
let transport_error_message = err . to_string ();
2026-05-16 19:23:23 +08:00
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 (),
),
2026-07-30 01:03:05 +08:00
error_message : Some ( transport_error_message . clone ()),
latency_ms : Some ( elapsed_ms_since ( candidate_started_at )),
2026-05-16 19:23:23 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
},
)
. await ;
2026-07-30 01:03:05 +08:00
if let Some ( response ) = maybe_build_sync_transport_error_stop_response (
state ,
& plan ,
report_context . as_ref (),
trace_id ,
decision ,
"chatgpt_web_image_execution_unavailable" ,
transport_error_message . as_str (),
elapsed_ms_since ( candidate_started_at ),
)
. await ?
{
return Ok ( Some ( response ));
}
2026-05-16 19:23:23 +08:00
return Ok ( None );
}
2026-05-06 02:29:17 +08:00
}
2026-05-16 19:23:23 +08:00
}
2026-04-03 14:59:58 +08:00
Err ( err ) => {
2026-07-30 01:03:05 +08:00
let transport_error_message = err . to_string ();
2026-04-03 14:59:58 +08:00
warn! (
2026-05-16 19:23:23 +08:00
event_name = "grok_execution_unavailable" ,
2026-04-05 20:23:16 +08:00
log_type = "ops" ,
2026-04-03 14:59:58 +08:00
trace_id = % trace_id ,
2026-04-07 02:50:19 +08:00
request_id = % plan_request_id_for_log ,
2026-04-03 14:59:58 +08:00
candidate_id = ? plan_candidate_id ,
2026-04-11 01:50:24 +08:00
provider_name ,
endpoint_id ,
key_id ,
model_name ,
candidate_index = candidate_index . as_str (),
2026-04-03 14:59:58 +08:00
error = % err ,
2026-05-16 19:23:23 +08:00
"gateway Grok execution unavailable"
2026-04-03 14:59:58 +08:00
);
2026-04-14 01:27:04 +08:00
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 ,
2026-05-16 19:23:23 +08:00
error_type : Some ( "grok_execution_unavailable" . to_string ()),
2026-07-30 01:03:05 +08:00
error_message : Some ( transport_error_message . clone ()),
latency_ms : Some ( elapsed_ms_since ( candidate_started_at )),
2026-04-14 01:27:04 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
},
)
. await ;
2026-07-30 01:03:05 +08:00
if let Some ( response ) = maybe_build_sync_transport_error_stop_response (
state ,
& plan ,
report_context . as_ref (),
trace_id ,
decision ,
"grok_execution_unavailable" ,
transport_error_message . as_str (),
elapsed_ms_since ( candidate_started_at ),
)
. await ?
{
return Ok ( Some ( response ));
}
2026-04-03 14:59:58 +08:00
return Ok ( None );
}
}
};
#[cfg(test)]
2026-04-21 16:19:07 +08:00
let mut result = {
2026-04-20 16:59:18 +08:00
if let Some ( override_fn ) = state . execution_runtime_sync_override . as_ref () {
match ( override_fn . 0 )( & plan ) {
Ok ( result ) => result ,
Err ( err ) => {
2026-07-30 01:03:05 +08:00
let transport_error_message = format! ( " {err:?} " );
2026-04-20 16:59:18 +08:00
warn! (
event_name = "sync_execution_runtime_test_override_failed" ,
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 test sync execution override failed"
);
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 ( "execution_runtime_unavailable" . to_string ()),
2026-07-30 01:03:05 +08:00
error_message : Some ( transport_error_message . clone ()),
latency_ms : Some ( elapsed_ms_since ( candidate_started_at )),
2026-04-20 16:59:18 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
},
)
. await ;
2026-07-30 01:03:05 +08:00
if let Some ( response ) = maybe_build_sync_transport_error_stop_response (
state ,
& plan ,
report_context . as_ref (),
trace_id ,
decision ,
"execution_runtime_unavailable" ,
transport_error_message . as_str (),
elapsed_ms_since ( candidate_started_at ),
)
. await ?
{
return Ok ( Some ( response ));
}
2026-04-20 16:59:18 +08:00
return Ok ( None );
}
}
} else if state
2026-04-04 01:40:24 +08:00
. execution_runtime_override_base_url ()
2026-04-20 16:59:18 +08:00
. unwrap_or_default ()
. trim ()
. is_empty ()
{
2026-05-16 19:23:23 +08:00
match maybe_execute_grok_sync ( & plan , report_context . as_ref ()). await {
2026-05-06 02:29:17 +08:00
Ok ( Some ( result )) => result ,
2026-05-16 19:23:23 +08:00
Ok ( None ) => match maybe_execute_chatgpt_web_image_sync (
2026-05-09 01:21:26 +08:00
state ,
& plan ,
report_context . as_ref (),
)
. await
{
2026-05-16 19:23:23 +08:00
Ok ( Some ( result )) => result ,
Ok ( None ) => match execute_direct_sync_runtime_candidate (
state ,
& plan ,
report_context . as_ref (),
trace_id ,
plan_kind ,
2026-07-01 02:15:51 +08:00
candidate_started_unix_secs ,
2026-05-16 19:23:23 +08:00
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 ) => {
2026-07-30 01:03:05 +08:00
let failure_error_type = err . error_type ;
let failure_message = err . message . clone ();
let failure_latency_ms = err
. latency_ms
. unwrap_or_else ( || elapsed_ms_since ( candidate_started_at ));
2026-07-27 09:36:31 +08:00
maybe_store_sync_execution_failure_fallback (
& err ,
& plan ,
trace_id ,
decision ,
& mut retry_scope_out ,
& mut retry_fallback_out ,
) ? ;
2026-05-16 19:23:23 +08:00
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 ,
2026-07-30 01:03:05 +08:00
status_code : None ,
error_type : Some ( failure_error_type . to_string ()),
2026-05-16 19:23:23 +08:00
error_message : Some ( err . message ),
2026-07-30 01:03:05 +08:00
latency_ms : Some ( failure_latency_ms ),
2026-05-16 19:23:23 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
},
)
. await ;
2026-07-30 01:03:05 +08:00
if let Some ( response ) = maybe_build_sync_transport_error_stop_response (
state ,
& plan ,
report_context . as_ref (),
trace_id ,
decision ,
failure_error_type ,
failure_message . as_str (),
failure_latency_ms ,
)
. await ?
{
return Ok ( Some ( response ));
}
2026-05-16 19:23:23 +08:00
return Ok ( None );
}
},
2026-05-06 02:29:17 +08:00
Err ( err ) => {
2026-07-30 01:03:05 +08:00
let transport_error_message = err . to_string ();
2026-05-06 02:29:17 +08:00
warn! (
2026-05-16 19:23:23 +08:00
event_name = "chatgpt_web_image_execution_unavailable" ,
2026-05-06 02:29:17 +08:00
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 (),
2026-05-16 19:23:23 +08:00
error = % err ,
"gateway ChatGPT-Web image execution unavailable"
2026-05-06 02:29:17 +08:00
);
let terminal_unix_secs = current_request_candidate_unix_ms ();
record_local_request_candidate_status (
state ,
& plan ,
report_context . as_ref (),
SchedulerRequestCandidateStatusUpdate {
status : RequestCandidateStatus ::Failed ,
2026-05-16 19:23:23 +08:00
status_code : None ,
error_type : Some (
"chatgpt_web_image_execution_unavailable" . to_string (),
),
2026-07-30 01:03:05 +08:00
error_message : Some ( transport_error_message . clone ()),
latency_ms : Some ( elapsed_ms_since ( candidate_started_at )),
2026-05-06 02:29:17 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
},
)
. await ;
2026-07-30 01:03:05 +08:00
if let Some ( response ) = maybe_build_sync_transport_error_stop_response (
state ,
& plan ,
report_context . as_ref (),
trace_id ,
decision ,
"chatgpt_web_image_execution_unavailable" ,
transport_error_message . as_str (),
elapsed_ms_since ( candidate_started_at ),
)
. await ?
{
return Ok ( Some ( response ));
}
2026-05-06 02:29:17 +08:00
return Ok ( None );
}
},
2026-04-03 14:59:58 +08:00
Err ( err ) => {
2026-07-30 01:03:05 +08:00
let transport_error_message = err . to_string ();
2026-04-03 14:59:58 +08:00
warn! (
2026-05-16 19:23:23 +08:00
event_name = "grok_execution_unavailable" ,
2026-04-05 20:23:16 +08:00
log_type = "ops" ,
2026-04-03 14:59:58 +08:00
trace_id = % trace_id ,
2026-04-07 02:50:19 +08:00
request_id = % plan_request_id_for_log ,
2026-04-03 14:59:58 +08:00
candidate_id = ? plan_candidate_id ,
2026-04-11 01:50:24 +08:00
provider_name ,
endpoint_id ,
key_id ,
model_name ,
candidate_index = candidate_index . as_str (),
2026-04-03 14:59:58 +08:00
error = % err ,
2026-05-16 19:23:23 +08:00
"gateway Grok execution unavailable"
2026-04-03 14:59:58 +08:00
);
2026-04-14 01:27:04 +08:00
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 ,
2026-05-16 19:23:23 +08:00
error_type : Some ( "grok_execution_unavailable" . to_string ()),
2026-07-30 01:03:05 +08:00
error_message : Some ( transport_error_message . clone ()),
latency_ms : Some ( elapsed_ms_since ( candidate_started_at )),
2026-04-14 01:27:04 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
},
)
. await ;
2026-07-30 01:03:05 +08:00
if let Some ( response ) = maybe_build_sync_transport_error_stop_response (
state ,
& plan ,
report_context . as_ref (),
trace_id ,
decision ,
"grok_execution_unavailable" ,
transport_error_message . as_str (),
elapsed_ms_since ( candidate_started_at ),
)
. await ?
{
return Ok ( Some ( response ));
}
2026-04-03 14:59:58 +08:00
return Ok ( None );
}
}
} else {
2026-04-20 16:59:18 +08:00
let remote_execution_runtime_base_url = state
. execution_runtime_override_base_url ()
. unwrap_or_default ();
2026-04-03 14:59:58 +08:00
let remote_outcome = execute_sync_via_remote_execution_runtime (
state ,
remote_execution_runtime_base_url ,
trace_id ,
decision ,
& plan ,
2026-04-25 21:29:45 +08:00
plan_request_id . as_str (),
plan_candidate_id . as_deref (),
2026-04-03 14:59:58 +08:00
report_context . as_ref (),
2026-04-10 01:46:14 +08:00
candidate_started_unix_secs ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
2026-04-03 14:59:58 +08:00
)
. await ? ;
match remote_outcome {
RemoteSyncFallbackOutcome ::Executed ( result ) => result ,
RemoteSyncFallbackOutcome ::ClientResponse ( response ) => return Ok ( Some ( response )),
RemoteSyncFallbackOutcome ::Unavailable => return Ok ( None ),
}
}
};
2026-07-30 01:03:05 +08:00
let mut candidate_first_byte_elapsed_ms =
calibrated_sync_candidate_first_byte_elapsed_ms ( candidate_started_at , & result );
2026-04-25 21:29:45 +08:00
let mut oauth_retry_attempted = false ;
let (
result_error_type ,
result_error_message ,
result_latency_ms ,
headers ,
body_bytes ,
body_json ,
body_base64 ,
local_failover_response_text ,
local_failover_analysis ,
) = loop {
let result_latency_ms = result
. telemetry
. as_ref ()
. and_then ( | telemetry | telemetry . elapsed_ms );
let mut headers = std ::mem ::take ( & mut result . headers );
2026-05-18 00:53:34 +00:00
let ( body_bytes , mut body_json , body_base64 ) =
2026-04-25 21:29:45 +08:00
decode_execution_result_body ( result . body . take (), & mut headers ) ? ;
2026-05-18 00:53:34 +00:00
if let Some ( message ) = invalid_gemini_provider_success_message (
& plan ,
report_context . as_ref (),
result . status_code ,
body_json . as_ref (),
2026-07-09 23:10:25 +08:00
)
. or_else ( || {
invalid_gemini_provider_stream_success_message (
& plan ,
report_context . as_ref (),
result . status_code ,
body_json . as_ref (),
& body_bytes ,
body_base64 . is_some (),
)
}) {
2026-05-18 00:53:34 +00:00
result . status_code = StatusCode ::BAD_GATEWAY . as_u16 ();
2026-07-09 23:10:25 +08:00
result . error = Some ( invalid_gemini_provider_success_execution_error ( message ));
2026-05-18 00:53:34 +00:00
if let Some ( error_body ) =
build_invalid_provider_success_body ( & plan , report_context . as_ref (), message )
{
body_json = Some ( error_body );
headers . insert ( "content-type" . to_string (), "application/json" . to_string ());
}
}
2026-05-20 23:42:12 +08:00
let ( mut result_error_type , mut result_error_message ) =
2026-05-18 00:53:34 +00:00
execution_error_details ( result . error . as_ref (), body_json . as_ref ());
2026-05-21 02:21:35 +08:00
if result . status_code < 400 && body_json . is_none () {
2026-05-20 23:42:12 +08:00
if let Some ( error_body_json ) =
extract_provider_private_stream_error_body ( report_context . as_ref (), & body_bytes )
{
result . status_code =
resolve_local_sync_error_status_code ( result . status_code , & error_body_json );
let ( private_error_type , private_error_message ) =
provider_private_error_details ( & error_body_json );
result_error_type = private_error_type . or ( result_error_type );
result_error_message = private_error_message . or ( result_error_message );
body_json = Some ( error_body_json );
}
}
2026-04-25 21:29:45 +08:00
let local_failover_response_text = local_failover_response_text (
body_json . as_ref (),
& body_bytes ,
result . error . as_ref (). map ( | error | error . message . as_str ()),
);
if result . status_code >= 400
&& ! oauth_retry_attempted
&& refresh_oauth_plan_auth_for_retry (
state ,
& mut plan ,
result . status_code ,
local_failover_response_text . as_deref (),
trace_id ,
)
. await
{
oauth_retry_attempted = true ;
match crate ::execution_runtime ::execute_execution_runtime_sync_plan (
state ,
Some ( trace_id ),
& plan ,
)
. await
{
Ok ( retry_result ) => {
2026-07-30 01:03:05 +08:00
candidate_first_byte_elapsed_ms =
calibrated_sync_candidate_first_byte_elapsed_ms (
candidate_started_at ,
& retry_result ,
);
2026-04-25 21:29:45 +08:00
result = retry_result ;
continue ;
}
Err ( err ) => {
warn! (
event_name = "local_sync_oauth_retry_execution_failed" ,
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 oauth retry sync execution failed"
);
}
}
}
let local_failover_analysis = analyze_local_candidate_failover_sync (
state ,
& plan ,
plan_kind ,
report_context . as_ref (),
& result ,
local_failover_response_text . as_deref (),
)
. await ;
break (
result_error_type ,
result_error_message ,
result_latency_ms ,
headers ,
body_bytes ,
body_json ,
body_base64 ,
local_failover_response_text ,
local_failover_analysis ,
);
};
2026-04-17 23:05:49 +08:00
if result . status_code >= 400 {
2026-04-18 11:35:11 +08:00
apply_local_execution_effect (
2026-04-17 23:05:49 +08:00
state ,
2026-04-18 11:35:11 +08:00
LocalExecutionEffectContext {
plan : & plan ,
report_context : report_context . as_ref (),
},
LocalExecutionEffect ::AttemptFailure ( LocalAttemptFailureEffect {
status_code : result . status_code ,
classification : local_failover_analysis . classification ,
}),
)
. await ;
apply_local_execution_effect (
state ,
LocalExecutionEffectContext {
plan : & plan ,
report_context : report_context . as_ref (),
},
LocalExecutionEffect ::AdaptiveRateLimit ( LocalAdaptiveRateLimitEffect {
status_code : result . status_code ,
classification : local_failover_analysis . classification ,
headers : Some ( & headers ),
}),
)
. await ;
apply_local_execution_effect (
state ,
LocalExecutionEffectContext {
plan : & plan ,
report_context : report_context . as_ref (),
},
LocalExecutionEffect ::HealthFailure ( LocalHealthFailureEffect {
status_code : result . status_code ,
classification : local_failover_analysis . classification ,
}),
)
. await ;
apply_local_execution_effect (
state ,
LocalExecutionEffectContext {
plan : & plan ,
report_context : report_context . as_ref (),
},
LocalExecutionEffect ::OauthInvalidation ( LocalOAuthInvalidationEffect {
status_code : result . status_code ,
response_text : local_failover_response_text . as_deref (),
}),
)
. await ;
apply_local_execution_effect (
state ,
LocalExecutionEffectContext {
plan : & plan ,
report_context : report_context . as_ref (),
},
LocalExecutionEffect ::PoolError ( LocalPoolErrorEffect {
status_code : result . status_code ,
classification : local_failover_analysis . classification ,
headers : & headers ,
error_body : local_failover_response_text . as_deref (),
}),
2026-04-17 23:05:49 +08:00
)
. await ;
}
2026-04-18 11:35:11 +08:00
if matches! (
local_failover_analysis . decision ,
LocalFailoverDecision ::RetryNextCandidate
) {
2026-07-27 09:36:31 +08:00
let failure_disposition = crate ::orchestration ::classify_failure_disposition (
& plan . provider_api_format ,
local_failover_analysis . classification ,
result . status_code ,
);
if let Some ( retry_scope ) = retry_scope_out . as_deref_mut () {
* retry_scope =
ai_attempt_retry_scope_from_failure_disposition ( failure_disposition );
}
if failure_disposition . preserve_upstream_error {
if let Some ( retry_fallback ) = retry_fallback_out . as_deref_mut () {
let mut fallback_headers = headers . clone ();
apply_endpoint_response_header_rules (
state ,
& plan ,
& mut fallback_headers ,
body_json . as_ref (),
)
. await ? ;
* retry_fallback = Some ( attach_control_metadata_headers (
build_client_response_from_parts (
result . status_code ,
& fallback_headers ,
Body ::from ( body_bytes . clone ()),
trace_id ,
Some ( decision ),
) ? ,
Some ( plan . request_id . as_str ()),
plan . candidate_id . as_deref (),
) ? );
}
}
2026-04-10 01:46:14 +08:00
let terminal_unix_secs = current_request_candidate_unix_ms ();
2026-05-09 10:46:04 +08:00
let error_trace_report_context = with_sync_error_trace_context (
2026-04-25 16:58:17 +08:00
report_context . as_ref (),
2026-05-09 10:46:04 +08:00
result . status_code ,
& headers ,
2026-05-11 02:30:45 +08:00
body_json . as_ref (),
& body_bytes ,
2026-05-09 10:46:04 +08:00
local_failover_response_text . as_deref (),
local_failover_analysis ,
2026-04-25 16:58:17 +08:00
);
2026-04-03 14:59:58 +08:00
record_local_request_candidate_status (
state ,
& plan ,
2026-05-09 10:46:04 +08:00
error_trace_report_context
2026-04-25 16:58:17 +08:00
. as_ref ()
. or ( report_context . as_ref ()),
2026-04-09 13:51:50 +08:00
SchedulerRequestCandidateStatusUpdate {
status : RequestCandidateStatus ::Failed ,
status_code : Some ( result . status_code ),
error_type : result_error_type . clone (),
error_message : result_error_message . clone (),
latency_ms : result_latency_ms ,
2026-04-10 01:46:14 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
2026-04-09 13:51:50 +08:00
},
2026-04-03 14:59:58 +08:00
)
. await ;
warn! (
2026-04-05 20:23:16 +08:00
event_name = "local_sync_candidate_retry_scheduled" ,
log_type = "event" ,
2026-04-03 14:59:58 +08:00
trace_id = % trace_id ,
2026-04-07 02:50:19 +08:00
request_id = % plan_request_id_for_log ,
2026-04-03 14:59:58 +08:00
status_code = result . status_code ,
2026-04-11 01:50:24 +08:00
provider_name ,
endpoint_id ,
key_id ,
model_name ,
candidate_index = candidate_index . as_str (),
2026-04-03 14:59:58 +08:00
"gateway local sync decision retrying next candidate after retryable execution runtime result"
);
return Ok ( None );
}
2026-04-21 16:19:07 +08:00
let status_code = result . status_code ;
2026-04-03 14:59:58 +08:00
let has_body_bytes = body_base64 . is_some ();
2026-05-31 22:13:45 +08:00
let mut report_context =
2026-05-06 01:29:42 +08:00
attach_provider_response_headers_to_report_context ( report_context , & headers );
2026-05-31 22:13:45 +08:00
if ( 200 .. 300 ). contains ( & status_code ) {
seed_kiro_sync_simulated_cache_enabled ( state , & plan , & mut report_context ). await ;
2026-06-02 23:06:29 +08:00
if kiro_simulated_cache_enabled_from_report_context ( report_context . as_ref ()) {
seed_kiro_sync_report_context_input_tokens ( & plan , & mut report_context );
}
2026-05-31 22:13:45 +08:00
seed_kiro_sync_report_context_prompt_cache_usage ( state , & plan , & mut report_context ). await ;
}
2026-05-06 01:29:42 +08:00
let mut client_headers = headers . clone ();
apply_endpoint_response_header_rules ( state , & plan , & mut client_headers , body_json . as_ref ())
. await ? ;
2026-04-03 14:59:58 +08:00
let explicit_finalize = should_finalize_sync_response ( report_kind . as_deref ());
let mapped_error_finalize_kind =
resolve_core_sync_error_finalize_report_kind ( plan_kind , & result , body_json . as_ref ());
2026-04-21 16:19:07 +08:00
let implicit_finalize = if ! explicit_finalize && mapped_error_finalize_kind . is_none () {
2026-04-03 14:59:58 +08:00
maybe_build_implicit_sync_finalize_outcome (
trace_id ,
decision ,
plan_kind ,
2026-04-21 16:19:07 +08:00
& report_context ,
status_code ,
2026-05-06 01:29:42 +08:00
& client_headers ,
2026-04-21 16:19:07 +08:00
& body_json ,
& body_base64 ,
& result . telemetry ,
2026-04-03 14:59:58 +08:00
) ?
} else {
2026-04-21 16:19:07 +08:00
None
2026-04-03 14:59:58 +08:00
};
2026-04-18 11:35:11 +08:00
if ! matches! (
local_failover_analysis . decision ,
LocalFailoverDecision ::StopLocalFailover
) && should_fallback_to_control_sync (
plan_kind ,
& result ,
body_json . as_ref (),
has_body_bytes ,
explicit_finalize || implicit_finalize . is_some (),
mapped_error_finalize_kind . is_some (),
) {
2026-04-10 01:46:14 +08:00
let terminal_unix_secs = current_request_candidate_unix_ms ();
2026-05-09 10:46:04 +08:00
let error_trace_report_context = with_sync_error_trace_context (
2026-04-25 16:58:17 +08:00
report_context . as_ref (),
2026-05-09 10:46:04 +08:00
result . status_code ,
& headers ,
2026-05-11 02:30:45 +08:00
body_json . as_ref (),
& body_bytes ,
2026-05-09 10:46:04 +08:00
local_failover_response_text . as_deref (),
local_failover_analysis ,
2026-04-25 16:58:17 +08:00
);
2026-04-03 14:59:58 +08:00
record_local_request_candidate_status (
state ,
& plan ,
2026-05-09 10:46:04 +08:00
error_trace_report_context
2026-04-25 16:58:17 +08:00
. as_ref ()
. or ( report_context . as_ref ()),
2026-04-09 13:51:50 +08:00
SchedulerRequestCandidateStatusUpdate {
status : RequestCandidateStatus ::Failed ,
status_code : Some ( result . status_code ),
error_type : result_error_type . clone (),
error_message : result_error_message . clone (),
latency_ms : result_latency_ms ,
2026-04-10 01:46:14 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
2026-04-09 13:51:50 +08:00
},
2026-04-03 14:59:58 +08:00
)
. await ;
return Ok ( None );
}
2026-04-10 01:46:14 +08:00
let terminal_unix_secs = current_request_candidate_unix_ms ();
2026-04-25 16:58:17 +08:00
let error_flow_report_context = ( result . status_code >= 400 )
. then ( || {
2026-05-09 10:46:04 +08:00
with_sync_error_trace_context (
2026-04-25 16:58:17 +08:00
report_context . as_ref (),
2026-05-09 10:46:04 +08:00
result . status_code ,
& headers ,
2026-05-11 02:30:45 +08:00
body_json . as_ref (),
& body_bytes ,
2026-05-09 10:46:04 +08:00
local_failover_response_text . as_deref (),
local_failover_analysis ,
2026-04-25 16:58:17 +08:00
)
})
. flatten ();
2026-04-03 14:59:58 +08:00
record_local_request_candidate_status (
state ,
& plan ,
2026-04-25 16:58:17 +08:00
error_flow_report_context
. as_ref ()
. or ( report_context . as_ref ()),
2026-04-09 13:51:50 +08:00
SchedulerRequestCandidateStatusUpdate {
status : if result . status_code >= 400 {
RequestCandidateStatus ::Failed
} else {
RequestCandidateStatus ::Success
},
status_code : Some ( result . status_code ),
error_type : result_error_type . clone (),
error_message : result_error_message . clone (),
latency_ms : result_latency_ms ,
2026-04-10 01:46:14 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
2026-04-03 14:59:58 +08:00
},
)
. await ;
2026-04-21 16:19:07 +08:00
let request_id_owned = result . request_id ;
let candidate_id_owned = result . candidate_id ;
let request_id = ( ! request_id_owned . trim (). is_empty ())
. then_some ( request_id_owned . as_str ())
2026-04-25 21:29:45 +08:00
. or ( Some ( plan_request_id . as_str ()));
2026-04-21 16:19:07 +08:00
let request_id_for_log = short_request_id ( request_id . unwrap_or ( "-" ));
2026-04-25 21:29:45 +08:00
let candidate_id = candidate_id_owned
. as_deref ()
. or ( plan_candidate_id . as_deref ());
2026-04-21 16:19:07 +08:00
let report_context = report_context ;
let body_json = body_json ;
let telemetry = result . telemetry ;
if let Some ( implicit_finalize ) = implicit_finalize {
let usage_payload = implicit_finalize
. outcome
. background_report
. as_ref ()
. unwrap_or ( & implicit_finalize . payload );
apply_sync_success_effects (
2026-04-17 23:05:49 +08:00
state ,
2026-04-21 16:19:07 +08:00
& plan ,
implicit_finalize . payload . report_context . as_ref (),
usage_payload ,
2026-04-18 11:35:11 +08:00
)
. await ;
2026-05-18 14:15:22 +00:00
record_sync_terminal_usage_and_disarm_guard (
2026-04-18 11:35:11 +08:00
state ,
2026-04-21 16:19:07 +08:00
& plan ,
implicit_finalize . payload . report_context . as_ref (),
usage_payload ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
candidate_first_byte_elapsed_ms ,
2026-05-18 14:15:22 +00:00
& mut terminal_guard ,
2026-07-22 02:09:34 +08:00
)
. await ;
2026-04-21 16:19:07 +08:00
if let Some ( report_payload ) = implicit_finalize . outcome . background_report {
spawn_sync_report ( state . clone (), report_payload );
} else {
warn! (
event_name = "local_core_finalize_missing_success_report_mapping" ,
log_type = "event" ,
trace_id = % trace_id ,
report_kind = % implicit_finalize . payload . report_kind ,
"gateway implicit local core finalize produced response without background success report mapping"
);
}
return Ok ( Some ( attach_control_metadata_headers (
implicit_finalize . outcome . response ,
request_id ,
candidate_id ,
) ? ));
2026-04-17 23:05:49 +08:00
}
2026-04-03 14:59:58 +08:00
2026-04-21 16:19:07 +08:00
let finalize_report_kind = if explicit_finalize {
report_kind . clone ()
} else {
mapped_error_finalize_kind
};
2026-04-03 14:59:58 +08:00
2026-04-21 16:19:07 +08:00
if let Some ( finalize_report_kind ) = finalize_report_kind {
let mut payload = build_sync_report_payload (
trace_id ,
finalize_report_kind ,
2026-04-03 14:59:58 +08:00
report_context ,
2026-04-21 16:19:07 +08:00
status_code ,
2026-05-06 01:29:42 +08:00
client_headers ,
2026-04-21 16:19:07 +08:00
body_json ,
body_base64 ,
telemetry ,
);
2026-04-03 14:59:58 +08:00
if let Some ( outcome ) = maybe_build_sync_finalize_outcome ( trace_id , decision , & payload ) ? {
let usage_payload = outcome . background_report . as_ref (). unwrap_or ( & payload );
2026-04-21 16:19:07 +08:00
if status_code < 400 {
apply_sync_success_effects (
state ,
& plan ,
payload . report_context . as_ref (),
usage_payload ,
)
. await ;
}
2026-05-18 14:15:22 +00:00
record_sync_terminal_usage_and_disarm_guard (
2026-04-03 14:59:58 +08:00
state ,
& plan ,
payload . report_context . as_ref (),
usage_payload ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
candidate_first_byte_elapsed_ms ,
2026-05-18 14:15:22 +00:00
& mut terminal_guard ,
2026-07-22 02:09:34 +08:00
)
. await ;
2026-04-03 14:59:58 +08:00
if let Some ( report_payload ) = outcome . background_report {
2026-04-21 16:19:07 +08:00
spawn_sync_report ( state . clone (), report_payload );
2026-04-03 14:59:58 +08:00
} else {
warn! (
2026-04-05 20:23:16 +08:00
event_name = "local_core_finalize_missing_success_report_mapping" ,
log_type = "event" ,
2026-04-03 14:59:58 +08:00
trace_id = % trace_id ,
report_kind = % payload . report_kind ,
"gateway local core finalize produced response without background success report mapping"
);
}
return Ok ( Some ( attach_control_metadata_headers (
outcome . response ,
request_id ,
candidate_id ,
) ? ));
}
2026-04-21 16:19:07 +08:00
let mut payload = match maybe_build_local_video_success_outcome (
2026-04-03 14:59:58 +08:00
trace_id ,
decision ,
2026-04-21 16:19:07 +08:00
payload ,
2026-04-03 14:59:58 +08:00
& state . video_tasks ,
& plan ,
) ? {
2026-04-21 16:19:07 +08:00
LocalVideoSyncSuccessBuild ::Handled ( outcome ) => {
let LocalVideoSyncSuccessOutcome {
response ,
report_payload ,
original_report_context ,
report_mode ,
local_task_snapshot ,
} = outcome ;
apply_sync_success_effects (
state ,
& plan ,
original_report_context . as_ref (),
& report_payload ,
)
. await ;
2026-05-18 14:15:22 +00:00
record_sync_terminal_usage_and_disarm_guard (
2026-04-21 16:19:07 +08:00
state ,
& plan ,
original_report_context . as_ref (),
& report_payload ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
candidate_first_byte_elapsed_ms ,
2026-05-18 14:15:22 +00:00
& mut terminal_guard ,
2026-07-22 02:09:34 +08:00
)
. await ;
2026-04-21 16:19:07 +08:00
if let Some ( snapshot ) = local_task_snapshot {
let _ = state . upsert_video_task_snapshot ( & snapshot ). await ? ;
state . video_tasks . record_snapshot ( snapshot );
2026-04-03 14:59:58 +08:00
}
2026-04-21 16:19:07 +08:00
match report_mode {
VideoTaskSyncReportMode ::InlineSync => {
submit_sync_report ( state , report_payload ). await ? ;
}
VideoTaskSyncReportMode ::Background => {
spawn_sync_report ( state . clone (), report_payload );
}
2026-04-03 14:59:58 +08:00
}
2026-04-21 16:19:07 +08:00
return Ok ( Some ( attach_control_metadata_headers (
response ,
request_id ,
candidate_id ,
) ? ));
2026-04-03 14:59:58 +08:00
}
2026-04-21 16:19:07 +08:00
LocalVideoSyncSuccessBuild ::NotHandled ( payload ) => payload ,
};
2026-04-03 14:59:58 +08:00
if let Some ( response ) =
maybe_build_local_sync_finalize_response ( trace_id , decision , & payload ) ?
{
2026-04-21 16:19:07 +08:00
let background_success_report_kind =
resolve_local_sync_success_background_report_kind ( payload . report_kind . as_str ());
apply_sync_success_effects ( state , & plan , payload . report_context . as_ref (), & payload )
. await ;
2026-05-18 14:15:22 +00:00
record_sync_terminal_usage_and_disarm_guard (
state ,
& plan ,
payload . report_context . as_ref (),
& payload ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
candidate_first_byte_elapsed_ms ,
2026-05-18 14:15:22 +00:00
& mut terminal_guard ,
2026-07-22 02:09:34 +08:00
)
. await ;
2026-04-03 14:59:58 +08:00
state
. video_tasks
. apply_finalize_mutation ( request_path , payload . report_kind . as_str ());
if let Some ( snapshot ) = state
. video_tasks
. snapshot_for_route ( decision . route_family . as_deref (), request_path )
{
let _ = state . upsert_video_task_snapshot ( & snapshot ). await ? ;
}
2026-04-21 16:19:07 +08:00
if let Some ( success_report_kind ) = background_success_report_kind {
payload . report_kind = success_report_kind . to_string ();
}
if background_success_report_kind . is_some () {
spawn_sync_report ( state . clone (), payload );
2026-04-03 14:59:58 +08:00
} else {
warn! (
2026-04-05 20:23:16 +08:00
event_name = "local_video_finalize_missing_success_report_mapping" ,
log_type = "ops" ,
2026-04-03 14:59:58 +08:00
trace_id = % trace_id ,
2026-04-07 02:50:19 +08:00
request_id = % request_id_for_log ,
2026-04-05 20:23:16 +08:00
candidate_id = ? candidate_id ,
2026-04-03 14:59:58 +08:00
report_kind = % payload . report_kind ,
"gateway local video finalize produced response without background success report mapping"
);
}
return Ok ( Some ( attach_control_metadata_headers (
response ,
request_id ,
candidate_id ,
) ? ));
}
if let Some ( response ) =
maybe_build_local_video_error_response ( trace_id , decision , & payload ) ?
{
2026-04-21 16:19:07 +08:00
let background_error_report_kind =
resolve_local_sync_error_background_report_kind ( payload . report_kind . as_str ());
if let Some ( error_report_kind ) = background_error_report_kind {
payload . report_kind = error_report_kind . to_string ();
}
2026-05-18 14:15:22 +00:00
record_sync_terminal_usage_and_disarm_guard (
state ,
& plan ,
payload . report_context . as_ref (),
& payload ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
candidate_first_byte_elapsed_ms ,
2026-05-18 14:15:22 +00:00
& mut terminal_guard ,
2026-07-22 02:09:34 +08:00
)
. await ;
2026-04-21 16:19:07 +08:00
if background_error_report_kind . is_some () {
spawn_sync_report ( state . clone (), payload );
2026-04-03 14:59:58 +08:00
} else {
warn! (
2026-04-05 20:23:16 +08:00
event_name = "local_video_finalize_missing_error_report_mapping" ,
log_type = "ops" ,
2026-04-03 14:59:58 +08:00
trace_id = % trace_id ,
2026-04-07 02:50:19 +08:00
request_id = % request_id_for_log ,
2026-04-05 20:23:16 +08:00
candidate_id = ? candidate_id ,
2026-04-03 14:59:58 +08:00
report_kind = % payload . report_kind ,
"gateway local video finalize produced response without background error report mapping"
);
}
return Ok ( Some ( attach_control_metadata_headers (
response ,
request_id ,
candidate_id ,
) ? ));
}
2026-05-18 14:15:22 +00:00
record_sync_terminal_usage_and_disarm_guard (
state ,
& plan ,
payload . report_context . as_ref (),
& payload ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
candidate_first_byte_elapsed_ms ,
2026-05-18 14:15:22 +00:00
& mut terminal_guard ,
2026-07-22 02:09:34 +08:00
)
. await ;
2026-04-03 14:59:58 +08:00
let response =
submit_local_core_error_or_sync_finalize ( state , trace_id , decision , payload ). await ? ;
return Ok ( Some ( attach_control_metadata_headers (
response ,
request_id ,
candidate_id ,
) ? ));
}
2026-04-21 16:19:07 +08:00
let usage_payload = build_sync_report_payload (
2026-04-03 14:59:58 +08:00
trace_id ,
2026-04-21 16:19:07 +08:00
report_kind . unwrap_or_default (),
report_context ,
status_code ,
2026-05-06 01:29:42 +08:00
client_headers ,
2026-04-21 16:19:07 +08:00
body_json ,
body_base64 ,
telemetry ,
);
if status_code < 400 {
apply_sync_success_effects (
state ,
& plan ,
usage_payload . report_context . as_ref (),
& usage_payload ,
)
. await ;
}
2026-05-18 14:15:22 +00:00
record_sync_terminal_usage_and_disarm_guard (
2026-04-21 16:19:07 +08:00
state ,
& plan ,
usage_payload . report_context . as_ref (),
& usage_payload ,
2026-07-30 01:03:05 +08:00
candidate_started_at ,
candidate_first_byte_elapsed_ms ,
2026-05-18 14:15:22 +00:00
& mut terminal_guard ,
2026-07-22 02:09:34 +08:00
)
. await ;
2026-04-21 16:19:07 +08:00
let response = attach_control_metadata_headers (
build_client_response_from_parts (
status_code ,
& usage_payload . headers ,
Body ::from ( body_bytes ),
trace_id ,
Some ( decision ),
) ? ,
request_id ,
candidate_id ,
) ? ;
if ! usage_payload . report_kind . trim (). is_empty () {
2026-07-08 23:49:18 +08:00
if status_code >= 400 {
let report_kind = usage_payload . report_kind . clone ();
if let Err ( err ) = submit_sync_report ( state , usage_payload ). await {
warn! (
event_name = "local_sync_error_report_submit_failed" ,
log_type = "ops" ,
trace_id = % trace_id ,
report_kind = % report_kind ,
"gateway failed to submit local sync error report before returning response: {err:?}"
);
}
} else {
spawn_sync_report ( state . clone (), usage_payload );
}
2026-04-21 16:19:07 +08:00
}
Ok ( Some ( response ))
2026-05-18 14:15:22 +00:00
})
. await ;
if let Err ( error ) = result . as_ref () {
terminal_guard . fail_and_disarm ( error ). await ;
} else {
terminal_guard . disarm ();
}
result
2026-04-03 14:59:58 +08:00
}
#[allow(clippy::too_many_arguments)] // mirrors sync execution context
fn maybe_build_implicit_sync_finalize_outcome (
trace_id : & str ,
decision : & GatewayControlDecision ,
plan_kind : & str ,
2026-04-21 16:19:07 +08:00
report_context : & Option < serde_json ::Value > ,
2026-04-03 14:59:58 +08:00
status_code : u16 ,
2026-04-21 16:19:07 +08:00
headers : & BTreeMap < String , String > ,
body_json : & Option < serde_json ::Value > ,
body_base64 : & Option < String > ,
telemetry : & Option < ExecutionTelemetry > ,
2026-04-03 14:59:58 +08:00
) -> Result < Option < ImplicitSyncFinalizeOutcome > , GatewayError > {
if status_code >= 400 || body_json . is_some () || body_base64 . is_none () {
return Ok ( None );
}
2026-04-03 16:26:16 +08:00
let Some ( report_kind ) = implicit_sync_finalize_report_kind ( plan_kind ) else {
2026-04-03 14:59:58 +08:00
return Ok ( None );
};
let payload = GatewaySyncReportRequest {
trace_id : trace_id . to_string (),
report_kind : report_kind . to_string (),
2026-04-21 16:19:07 +08:00
report_context : report_context . clone (),
2026-04-03 14:59:58 +08:00
status_code ,
2026-04-21 16:19:07 +08:00
headers : headers . clone (),
body_json : body_json . clone (),
2026-04-03 14:59:58 +08:00
client_body_json : None ,
2026-04-21 16:19:07 +08:00
body_base64 : body_base64 . clone (),
telemetry : telemetry . clone (),
2026-04-03 14:59:58 +08:00
};
let Some ( outcome ) = maybe_build_sync_finalize_outcome ( trace_id , decision , & payload ) ? else {
return Ok ( None );
};
Ok ( Some ( ImplicitSyncFinalizeOutcome { payload , outcome }))
}
#[allow(clippy::too_many_arguments)] // internal helper mirroring execute path context
#[cfg(test)]
async fn execute_sync_via_remote_execution_runtime (
state : & AppState ,
remote_execution_runtime_base_url : & str ,
trace_id : & str ,
decision : & GatewayControlDecision ,
plan : & ExecutionPlan ,
plan_request_id : & str ,
plan_candidate_id : Option <& str > ,
report_context : Option <& serde_json ::Value > ,
2026-04-10 01:46:14 +08:00
candidate_started_unix_secs : u64 ,
2026-07-30 01:03:05 +08:00
candidate_started_at : Instant ,
2026-04-03 14:59:58 +08:00
) -> Result < RemoteSyncFallbackOutcome , GatewayError > {
let response = match post_sync_plan_to_remote_execution_runtime (
state ,
remote_execution_runtime_base_url ,
Some ( trace_id ),
plan ,
)
. await
{
Ok ( response ) => response ,
Err ( err ) => {
warn! (
2026-04-05 20:23:16 +08:00
event_name = "sync_execution_runtime_remote_unavailable" ,
log_type = "ops" ,
2026-04-03 14:59:58 +08:00
trace_id = % trace_id ,
2026-04-07 02:50:19 +08:00
request_id = % short_request_id ( plan_request_id ),
2026-04-03 14:59:58 +08:00
candidate_id = ? plan_candidate_id ,
error = ? err ,
"gateway remote execution runtime sync unavailable"
);
2026-04-14 01:27:04 +08:00
let terminal_unix_secs = current_request_candidate_unix_ms ();
record_local_request_candidate_status (
state ,
plan ,
report_context ,
SchedulerRequestCandidateStatusUpdate {
status : RequestCandidateStatus ::Failed ,
status_code : None ,
error_type : Some ( "execution_runtime_unavailable" . to_string ()),
error_message : Some ( format! ( " {err:?} " )),
2026-07-30 01:03:05 +08:00
latency_ms : Some ( elapsed_ms_since ( candidate_started_at )),
2026-04-14 01:27:04 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
},
)
. await ;
2026-04-03 14:59:58 +08:00
return Ok ( RemoteSyncFallbackOutcome ::Unavailable );
}
};
if response . status () != http ::StatusCode ::OK {
2026-04-10 01:46:14 +08:00
let terminal_unix_secs = current_request_candidate_unix_ms ();
2026-04-03 14:59:58 +08:00
record_local_request_candidate_status (
state ,
plan ,
report_context ,
2026-04-09 13:51:50 +08:00
SchedulerRequestCandidateStatusUpdate {
status : RequestCandidateStatus ::Failed ,
status_code : Some ( response . status (). as_u16 ()),
error_type : Some ( "execution_runtime_http_error" . to_string ()),
error_message : Some ( format! (
"execution runtime returned HTTP {} " ,
response . status ()
)),
2026-07-30 01:03:05 +08:00
latency_ms : Some ( elapsed_ms_since ( candidate_started_at )),
2026-04-10 01:46:14 +08:00
started_at_unix_ms : Some ( candidate_started_unix_secs ),
finished_at_unix_ms : Some ( terminal_unix_secs ),
2026-04-09 13:51:50 +08:00
},
2026-04-03 14:59:58 +08:00
)
. await ;
return Ok ( RemoteSyncFallbackOutcome ::ClientResponse (
attach_control_metadata_headers (
build_client_response ( response , trace_id , Some ( decision )) ? ,
Some ( plan_request_id ),
plan_candidate_id ,
) ? ,
));
}
response
. json ()
. await
. map ( RemoteSyncFallbackOutcome ::Executed )
. map_err ( | err | GatewayError ::Internal ( err . to_string ()))
}
2026-05-09 01:21:26 +08:00
#[cfg(test)]
mod tests {
use super ::* ;
2026-05-18 14:15:22 +00:00
use aether_data ::repository ::candidates ::InMemoryRequestCandidateRepository ;
use aether_data ::repository ::usage ::InMemoryUsageReadRepository ;
use aether_data_contracts ::repository ::candidates ::RequestCandidateReadRepository ;
use aether_data_contracts ::repository ::usage ::UsageReadRepository ;
use aether_usage_runtime ::UsageRuntimeConfig ;
2026-05-09 01:21:26 +08:00
use futures_util ::{ pin_mut , StreamExt as _ };
use serde_json ::json ;
use std ::collections ::BTreeMap ;
2026-05-18 14:15:22 +00:00
use std ::sync ::Arc ;
2026-07-01 02:15:51 +08:00
use tokio ::io ::{ AsyncReadExt , AsyncWriteExt };
2026-05-09 01:21:26 +08:00
fn test_openai_image_plan ( stream : bool ) -> ExecutionPlan {
ExecutionPlan {
request_id : "req-1" . to_string (),
candidate_id : Some ( "candidate-1" . to_string ()),
provider_name : Some ( "OpenAI" . to_string ()),
provider_id : "provider-1" . to_string (),
endpoint_id : "endpoint-1" . to_string (),
key_id : "key-1" . to_string (),
method : "POST" . to_string (),
url : "https://chatgpt.com/backend-api/codex/responses" . to_string (),
headers : BTreeMap ::new (),
content_type : Some ( "application/json" . to_string ()),
content_encoding : None ,
body : aether_contracts ::RequestBody ::from_json ( json! ({ "stream" : true })),
stream ,
client_api_format : "openai:image" . to_string (),
provider_api_format : "openai:image" . to_string (),
model_name : Some ( "gpt-image-2" . to_string ()),
proxy : None ,
transport_profile : None ,
timeouts : None ,
}
}
2026-05-18 00:53:34 +00:00
fn test_gemini_chat_plan () -> ExecutionPlan {
let mut plan = test_openai_image_plan ( false );
plan . client_api_format = "openai:chat" . to_string ();
plan . provider_api_format = "gemini:generate_content" . to_string ();
plan . model_name = Some ( "gemini-3-flash-preview" . to_string ());
plan
}
2026-07-01 02:15:51 +08:00
fn test_decision () -> GatewayControlDecision {
GatewayControlDecision ::synthetic (
"/v1/chat/completions" ,
Some ( "ai_public" . to_string ()),
Some ( "openai" . to_string ()),
Some ( "chat" . to_string ()),
Some ( "openai:chat" . to_string ()),
)
. with_execution_runtime_candidate ( true )
2026-07-27 09:36:31 +08:00
}
#[tokio::test]
async fn oversized_upstream_response_builds_claude_502_retry_fallback () {
let mut plan = test_openai_image_plan ( false );
plan . client_api_format = "claude:messages" . to_string ();
plan . provider_api_format = "claude:messages" . to_string ();
let decision = GatewayControlDecision ::synthetic (
"/v1/messages" ,
Some ( "ai_public" . to_string ()),
Some ( "claude" . to_string ()),
Some ( "messages" . to_string ()),
Some ( "claude:messages" . to_string ()),
)
. with_execution_runtime_candidate ( true );
let failure = SyncExecutionFailure ::from_transport (
ExecutionRuntimeTransportError ::UpstreamResponseTooLarge {
phase : crate ::execution_runtime ::transport ::UpstreamResponseBodyPhase ::Wire ,
limit_bytes : 8 ,
},
);
assert_eq! ( failure . status_code , Some ( StatusCode ::BAD_GATEWAY . as_u16 ()));
assert_eq! (
failure . fallback_kind ,
Some ( SyncExecutionFailureFallbackKind ::UpstreamResponseTooLarge )
);
let mut retry_scope = AiAttemptRetryScope ::Provider ;
let mut retry_fallback = None ;
{
let mut retry_scope_out = Some ( & mut retry_scope );
let mut retry_fallback_out = Some ( & mut retry_fallback );
maybe_store_sync_execution_failure_fallback (
& failure ,
& plan ,
"trace-too-large" ,
& decision ,
& mut retry_scope_out ,
& mut retry_fallback_out ,
)
. expect ( "fallback response should build" );
}
assert_eq! ( retry_scope , AiAttemptRetryScope ::Candidate );
let response = retry_fallback . expect ( "oversized response should provide a fallback" );
assert_eq! ( response . status (), StatusCode ::BAD_GATEWAY );
let body = to_bytes ( response . into_body (), 1024 )
. await
. expect ( "fallback body should read" );
let body : Value = serde_json ::from_slice ( & body ). expect ( "fallback body should be json" );
assert_eq! ( body [ "type" ], "error" );
assert_eq! ( body [ "error" ][ "type" ], "upstream_error" );
assert_eq! ( body [ "error" ][ "message" ], "Upstream response too large" );
2026-07-01 02:15:51 +08:00
}
2026-05-31 22:13:45 +08:00
fn test_kiro_sync_plan () -> ExecutionPlan {
ExecutionPlan {
request_id : "req-kiro-sync-cache-1" . to_string (),
candidate_id : Some ( "candidate-kiro-sync-cache-1" . to_string ()),
provider_name : Some ( "Kiro" . to_string ()),
provider_id : "provider-kiro-sync-1" . to_string (),
endpoint_id : "endpoint-kiro-sync-1" . to_string (),
key_id : "key-kiro-sync-1" . to_string (),
method : "POST" . to_string (),
url : "https://kiro.example/generateAssistantResponse" . to_string (),
headers : BTreeMap ::new (),
content_type : Some ( "application/json" . to_string ()),
content_encoding : None ,
body : aether_contracts ::RequestBody ::from_json ( json! ({
"model" : "claude-sonnet-4" ,
"messages" : [{ "role" : "user" , "content" : "hello kiro" }],
})),
stream : false ,
client_api_format : "claude:messages" . to_string (),
provider_api_format : "claude:messages" . to_string (),
model_name : Some ( "claude-sonnet-4" . to_string ()),
proxy : None ,
transport_profile : None ,
timeouts : None ,
}
}
fn test_kiro_sync_cacheable_request_body () -> serde_json ::Value {
json! ({
"model" : "claude-sonnet-4" ,
"system" : [{
"type" : "text" ,
"text" : format ! ( "sync cacheable prompt {}" , "cacheable prompt chunk " . repeat ( 300 )),
"cache_control" : { "type" : "ephemeral" }
}],
"messages" : [{ "role" : "user" , "content" : "reuse this Kiro prompt" }]
})
}
2026-05-18 00:53:34 +00:00
#[test]
fn invalid_gemini_provider_success_uses_plan_format_when_context_is_missing () {
let plan = test_gemini_chat_plan ();
let body = json! ({
"candidates" : [{
"content" : { "role" : "model" },
"finishReason" : "MAX_TOKENS"
}],
"usageMetadata" : {
"promptTokenCount" : 8 ,
"candidatesTokenCount" : 1 ,
"thoughtsTokenCount" : 25 ,
"totalTokenCount" : 34
}
});
let message = invalid_gemini_provider_success_message (
& plan ,
None ,
StatusCode ::OK . as_u16 (),
Some ( & body ),
)
. expect ( "empty Gemini 200 response should be rejected from plan format" );
assert! ( message . contains ( "visible model output" ));
}
2026-07-09 23:10:25 +08:00
#[test]
fn invalid_gemini_provider_success_error_is_retryable_candidate_failure () {
let error = invalid_gemini_provider_success_execution_error (
INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE ,
);
assert_eq! ( error . kind , ExecutionErrorKind ::Upstream5xx );
assert_eq! ( error . phase , ExecutionPhase ::Finalize );
assert_eq! ( error . upstream_status , Some ( StatusCode ::OK . as_u16 ()));
assert! ( error . retryable );
assert! ( error . failover_recommended );
}
#[test]
fn invalid_gemini_provider_success_accepts_antigravity_chunks_with_visible_output () {
let plan = test_gemini_chat_plan ();
let report_context = json! ({
"has_envelope" : true ,
"envelope_name" : "antigravity:v1internal" ,
"provider_api_format" : "gemini:generate_content" ,
});
let body = json! ({
"chunks" : [{
"response" : {
"responseId" : "resp_antigravity_chunks_123" ,
"candidates" : [{
"content" : {
"parts" : [{ "text" : "Hello Gemini" }],
"role" : "model"
},
"finishReason" : "STOP" ,
"index" : 0
}],
"modelVersion" : "gemini-3-flash-agent" ,
"usageMetadata" : {
"promptTokenCount" : 2 ,
"candidatesTokenCount" : 2 ,
"totalTokenCount" : 4
}
},
"traceId" : "trace-antigravity-chunks"
}],
"metadata" : {
"stream" : true ,
"stored_chunks" : 1 ,
"total_chunks" : 1
}
});
let message = invalid_gemini_provider_success_message (
& plan ,
Some ( & report_context ),
StatusCode ::OK . as_u16 (),
Some ( & body ),
);
assert! ( message . is_none ());
}
2026-05-21 15:38:10 +08:00
#[test]
fn invalid_gemini_provider_success_unwraps_gemini_cli_v1internal_envelope () {
let plan = test_gemini_chat_plan ();
let report_context = json! ({
"has_envelope" : true ,
"envelope_name" : "gemini_cli:v1internal" ,
"provider_api_format" : "gemini:generate_content" ,
});
let body = json! ({
"response" : {
"candidates" : [{
"content" : {
"role" : "model" ,
"parts" : [{ "text" : "Hello from Gemini CLI" }]
},
"finishReason" : "STOP"
}]
},
"remainingCredits" : 41 ,
"consumedCredits" : 1 ,
"traceId" : "trace-upstream-sync-1"
});
let message = invalid_gemini_provider_success_message (
& plan ,
Some ( & report_context ),
StatusCode ::OK . as_u16 (),
Some ( & body ),
);
assert! ( message . is_none ());
}
2026-05-18 14:15:22 +00:00
#[tokio::test]
async fn sync_attempt_terminal_guard_marks_dropped_pending_attempt_cancelled () {
let usage_repository = Arc ::new ( InMemoryUsageReadRepository ::default ());
let request_candidate_repository = Arc ::new ( InMemoryRequestCandidateRepository ::default ());
let state = AppState ::new ()
. expect ( "gateway state should build" )
. with_data_state_for_tests (
crate ::data ::GatewayDataState ::with_request_candidate_and_usage_repository_for_tests (
Arc ::clone ( & request_candidate_repository ),
Arc ::clone ( & usage_repository ),
),
)
. with_usage_runtime_for_tests ( UsageRuntimeConfig {
enabled : true ,
.. UsageRuntimeConfig ::default ()
});
let mut plan = test_openai_image_plan ( false );
plan . request_id = "sync-cancel-guard-request" . to_string ();
plan . candidate_id = None ;
let mut report_context = Some ( json! ({
"candidate_index" : 0 ,
"retry_index" : 0 ,
"user_id" : "user-cancel" ,
"api_key_id" : "api-key-cancel" ,
"client_api_format" : "openai:image" ,
"provider_api_format" : "openai:image" ,
"request_path" : "/v1/images/generations" ,
"request_path_and_query" : "/v1/images/generations" ,
"upstream_url" : "https://example.test/v1/images/generations" ,
"mapped_model" : "gpt-image-2" ,
}));
ensure_execution_request_candidate_slot ( & state , & mut plan , & mut report_context ). await ;
2026-07-30 01:03:05 +08:00
let candidate_started_at = Instant ::now ();
2026-05-18 14:15:22 +00:00
let started_at = current_request_candidate_unix_ms ();
state . usage_runtime . record_pending (
2026-07-22 02:09:34 +08:00
state . usage_lifecycle_data_state (). as_ref (),
2026-05-18 14:15:22 +00:00
build_lifecycle_usage_seed ( & plan , report_context . as_ref ()),
);
record_local_request_candidate_status (
& state ,
& plan ,
report_context . as_ref (),
SchedulerRequestCandidateStatusUpdate {
status : RequestCandidateStatus ::Pending ,
status_code : None ,
error_type : None ,
error_message : None ,
latency_ms : None ,
started_at_unix_ms : Some ( started_at ),
finished_at_unix_ms : None ,
},
)
. await ;
2026-07-30 01:03:05 +08:00
crate ::request_diagnostics ::scope_request_diagnostics ( async {
crate ::request_diagnostics ::record_request_accepted_at (
Instant ::now () - Duration ::from_millis ( 25 ),
);
let _guard = SyncAttemptTerminalGuard ::new (
& state ,
& plan ,
report_context . clone (),
started_at ,
candidate_started_at ,
);
})
. await ;
2026-05-18 14:15:22 +00:00
let mut stored_usage = None ;
for _ in 0 .. 50 {
if let Some ( usage ) = usage_repository
. find_by_request_id ( "sync-cancel-guard-request" )
. await
. expect ( "usage should read" )
{
if usage . status == "cancelled" {
stored_usage = Some ( usage );
break ;
}
}
tokio ::time ::sleep ( Duration ::from_millis ( 10 )). await ;
}
let stored_usage = stored_usage . expect ( "cancelled usage should be recorded" );
assert_eq! ( stored_usage . status , "cancelled" );
2026-06-29 00:30:41 +08:00
assert_eq! ( stored_usage . billing_status , "void" );
2026-05-18 14:15:22 +00:00
assert_eq! ( stored_usage . status_code , Some ( 499 ));
assert_eq! ( stored_usage . error_category . as_deref (), Some ( "cancelled" ));
2026-07-30 01:03:05 +08:00
let request_metadata = stored_usage
. request_metadata
. as_ref ()
. expect ( "cancelled usage should retain request diagnostics" );
assert! ( request_metadata
. get ( "end_to_end_time_ms" )
. and_then ( Value ::as_u64 )
. is_some ());
assert! ( request_metadata
. get ( "end_to_end_first_byte_time_ms" )
. is_none ());
2026-05-18 14:15:22 +00:00
let stored_candidates = request_candidate_repository
. list_by_request_id ( "sync-cancel-guard-request" )
. await
. expect ( "request candidates should read" );
assert_eq! ( stored_candidates . len (), 1 );
assert_eq! (
stored_candidates [ 0 ]. status ,
RequestCandidateStatus ::Cancelled
);
assert_eq! ( stored_candidates [ 0 ]. status_code , Some ( 499 ));
assert_eq! (
stored_candidates [ 0 ]. error_type . as_deref (),
Some ( "local_sync_attempt_cancelled" )
);
}
2026-07-01 02:15:51 +08:00
#[tokio::test]
async fn sync_direct_response_start_marks_usage_and_candidate_active_before_body_finishes () {
let usage_repository = Arc ::new ( InMemoryUsageReadRepository ::default ());
let request_candidate_repository = Arc ::new ( InMemoryRequestCandidateRepository ::default ());
let state = AppState ::new ()
. expect ( "gateway state should build" )
. with_data_state_for_tests (
crate ::data ::GatewayDataState ::with_request_candidate_and_usage_repository_for_tests (
Arc ::clone ( & request_candidate_repository ),
Arc ::clone ( & usage_repository ),
),
)
. with_usage_runtime_for_tests ( UsageRuntimeConfig {
enabled : true ,
.. UsageRuntimeConfig ::default ()
});
let listener = crate ::test_support ::bind_loopback_listener ()
. await
. expect ( "listener should bind" );
let addr = listener . local_addr (). expect ( "local addr should resolve" );
let ( headers_tx , headers_rx ) = tokio ::sync ::oneshot ::channel ();
let ( body_tx , body_rx ) = tokio ::sync ::oneshot ::channel ();
let server = tokio ::spawn ( async move {
let ( mut socket , _ ) = listener . accept (). await . expect ( "client should connect" );
let mut request = [ 0_ u8 ; 4096 ];
let _ = socket
. read ( & mut request )
. await
. expect ( "request should read" );
socket
. write_all (
b "HTTP/1.1 200 OK \r\n content-type: application/json \r\n content-length: 11 \r\n\r\n " ,
)
. await
. expect ( "headers should write" );
socket . flush (). await . expect ( "headers should flush" );
let _ = headers_tx . send (());
let _ = body_rx . await ;
socket
. write_all ( br #"{"ok":true}"# )
. await
. expect ( "body should write" );
});
let mut plan = test_gemini_chat_plan ();
plan . request_id = "sync-response-start-active-request" . to_string ();
plan . candidate_id = Some ( "sync-response-start-active-candidate" . to_string ());
plan . url = format! ( "http:// {addr} /chat" );
plan . provider_api_format = "openai:chat" . to_string ();
plan . model_name = Some ( "gpt-5" . to_string ());
plan . body = aether_contracts ::RequestBody ::from_json ( json! ({
"model" : "gpt-5" ,
"messages" : [{ "role" : "user" , "content" : "slow body" }],
}));
let report_context = Some ( json! ({
"candidate_index" : 0 ,
"retry_index" : 0 ,
"user_id" : "user-active" ,
"api_key_id" : "api-key-active" ,
"client_api_format" : "openai:chat" ,
"provider_api_format" : "openai:chat" ,
"request_path" : "/v1/chat/completions" ,
"request_path_and_query" : "/v1/chat/completions" ,
"upstream_url" : plan . url . clone (),
"mapped_model" : "gpt-5" ,
}));
let started_at = current_request_candidate_unix_ms ();
state
. usage_runtime
. record_pending_direct (
2026-07-22 02:09:34 +08:00
state . usage_lifecycle_data_state (). as_ref (),
2026-07-01 02:15:51 +08:00
build_lifecycle_usage_seed ( & plan , report_context . as_ref ()),
)
. await ;
record_local_request_candidate_status (
& state ,
& plan ,
report_context . as_ref (),
SchedulerRequestCandidateStatusUpdate {
status : RequestCandidateStatus ::Pending ,
status_code : None ,
error_type : None ,
error_message : None ,
latency_ms : None ,
started_at_unix_ms : Some ( started_at ),
finished_at_unix_ms : None ,
},
)
. await ;
let state_for_exec = state . clone ();
let plan_for_exec = plan . clone ();
let report_context_for_exec = report_context . clone ();
let exec = tokio ::spawn ( async move {
execute_direct_sync_runtime_candidate (
& state_for_exec ,
& plan_for_exec ,
report_context_for_exec . as_ref (),
"trace-response-start-active" ,
"openai_chat_sync" ,
started_at ,
"sync-response-start-active-request" ,
plan_for_exec . candidate_id . as_deref (),
"openai" ,
"endpoint-1" ,
"key-1" ,
"gpt-5" ,
"0" ,
None ,
)
. await
});
headers_rx . await . expect ( "headers should be written" );
let mut active_usage = None ;
for _ in 0 .. 50 {
if let Some ( usage ) = usage_repository
. find_by_request_id ( "sync-response-start-active-request" )
. await
. expect ( "usage should read" )
{
if usage . status == "streaming" {
active_usage = Some ( usage );
break ;
}
}
tokio ::time ::sleep ( Duration ::from_millis ( 10 )). await ;
}
let active_usage = active_usage . expect ( "usage should become active before body finishes" );
assert_eq! ( active_usage . status_code , Some ( 200 ));
assert! ( active_usage . first_byte_time_ms . is_some ());
assert! ( active_usage . response_time_ms . is_some ());
let stored_candidates = request_candidate_repository
. list_by_request_id ( "sync-response-start-active-request" )
. await
. expect ( "candidate should read" );
let active_candidate = stored_candidates
. iter ()
. find ( | candidate | candidate . id == "sync-response-start-active-candidate" )
. expect ( "candidate should exist" );
assert_eq! ( active_candidate . status , RequestCandidateStatus ::Streaming );
assert_eq! ( active_candidate . status_code , Some ( 200 ));
let _ = body_tx . send (());
let result = tokio ::time ::timeout ( Duration ::from_secs ( 2 ), exec )
. await
. expect ( "sync execution should finish" )
. expect ( "sync execution task should not panic" )
. expect ( "sync execution should succeed" );
assert_eq! ( result . status_code , 200 );
server . abort ();
}
#[tokio::test]
async fn sync_execution_active_marks_usage_before_response_headers () {
let usage_repository = Arc ::new ( InMemoryUsageReadRepository ::default ());
let request_candidate_repository = Arc ::new ( InMemoryRequestCandidateRepository ::default ());
let state = AppState ::new ()
. expect ( "gateway state should build" )
. with_data_state_for_tests (
crate ::data ::GatewayDataState ::with_request_candidate_and_usage_repository_for_tests (
Arc ::clone ( & request_candidate_repository ),
Arc ::clone ( & usage_repository ),
),
)
. with_usage_runtime_for_tests ( UsageRuntimeConfig {
enabled : true ,
.. UsageRuntimeConfig ::default ()
});
let listener = crate ::test_support ::bind_loopback_listener ()
. await
. expect ( "listener should bind" );
let addr = listener . local_addr (). expect ( "local addr should resolve" );
let ( request_seen_tx , request_seen_rx ) = tokio ::sync ::oneshot ::channel ();
let ( finish_tx , finish_rx ) = tokio ::sync ::oneshot ::channel ();
let server = tokio ::spawn ( async move {
let ( mut socket , _ ) = listener . accept (). await . expect ( "client should connect" );
let mut request = [ 0_ u8 ; 4096 ];
let _ = socket
. read ( & mut request )
. await
. expect ( "request should read" );
let _ = request_seen_tx . send (());
let _ = finish_rx . await ;
socket
. write_all (
b "HTTP/1.1 200 OK \r\n content-type: application/json \r\n content-length: 11 \r\n\r\n { \" ok \" :true}" ,
)
. await
. expect ( "response should write" );
});
let mut plan = test_gemini_chat_plan ();
plan . request_id = "sync-active-before-headers-request" . to_string ();
plan . candidate_id = Some ( "sync-active-before-headers-candidate" . to_string ());
plan . url = format! ( "http:// {addr} /chat" );
plan . provider_name = Some ( "OpenAI" . to_string ());
plan . provider_api_format = "openai:chat" . to_string ();
plan . model_name = Some ( "gpt-5" . to_string ());
plan . body = aether_contracts ::RequestBody ::from_json ( json! ({
"model" : "gpt-5" ,
"messages" : [{ "role" : "user" , "content" : "slow headers" }],
}));
let report_context = Some ( json! ({
"candidate_index" : 0 ,
"retry_index" : 0 ,
"user_id" : "user-active-before-headers" ,
"api_key_id" : "api-key-active-before-headers" ,
"candidate_id" : "sync-active-before-headers-candidate" ,
"provider_id" : "provider-1" ,
"endpoint_id" : "endpoint-1" ,
"key_id" : "key-1" ,
"provider_name" : "OpenAI" ,
"client_api_format" : "openai:chat" ,
"provider_api_format" : "openai:chat" ,
"request_path" : "/v1/chat/completions" ,
"request_path_and_query" : "/v1/chat/completions" ,
"upstream_url" : plan . url . clone (),
"mapped_model" : "gpt-5" ,
}));
let state_for_exec = state . clone ();
let plan_for_exec = plan . clone ();
let report_context_for_exec = report_context . clone ();
let exec = tokio ::spawn ( async move {
execute_execution_runtime_sync (
& state_for_exec ,
"/v1/chat/completions" ,
plan_for_exec ,
"trace-active-before-headers" ,
& test_decision (),
"openai_chat_sync" ,
Some ( "openai_chat_sync" . to_string ()),
report_context_for_exec ,
)
. await
});
request_seen_rx
. await
. expect ( "upstream request should be observed" );
let mut active_usage = None ;
for _ in 0 .. 50 {
if let Some ( usage ) = usage_repository
. find_by_request_id ( "sync-active-before-headers-request" )
. await
. expect ( "usage should read" )
{
if usage . status == "streaming" {
active_usage = Some ( usage );
break ;
}
}
tokio ::time ::sleep ( Duration ::from_millis ( 10 )). await ;
}
let active_usage =
active_usage . expect ( "usage should become active before upstream headers" );
assert_eq! ( active_usage . status_code , None );
assert_eq! ( active_usage . first_byte_time_ms , None );
assert_eq! ( active_usage . response_time_ms , None );
let stored_candidates = request_candidate_repository
. list_by_request_id ( "sync-active-before-headers-request" )
. await
. expect ( "candidate should read" );
let active_candidate = stored_candidates
. iter ()
. find ( | candidate | candidate . id == "sync-active-before-headers-candidate" )
. expect ( "candidate should exist" );
assert_eq! ( active_candidate . status , RequestCandidateStatus ::Streaming );
assert_eq! ( active_candidate . status_code , None );
assert! ( active_candidate . started_at_unix_ms . is_some ());
assert! ( active_candidate . finished_at_unix_ms . is_none ());
let _ = finish_tx . send (());
let response = tokio ::time ::timeout ( Duration ::from_secs ( 2 ), exec )
. await
. expect ( "sync execution should finish" )
. expect ( "sync execution task should not panic" )
. expect ( "sync execution should succeed" )
. expect ( "sync execution should produce a response" );
assert_eq! ( response . status (), StatusCode ::OK );
server . abort ();
}
2026-05-31 22:13:45 +08:00
#[test]
fn kiro_sync_report_context_seeds_input_tokens_from_original_request_body () {
let plan = test_kiro_sync_plan ();
let mut report_context = Some ( json! ({
"original_request_body" : test_kiro_sync_cacheable_request_body (),
}));
seed_kiro_sync_report_context_input_tokens ( & plan , & mut report_context );
assert! ( report_context
. as_ref ()
. and_then ( | value | value . get ( "input_tokens" ))
. and_then ( Value ::as_u64 )
. is_some_and ( | tokens | tokens > 0 ));
}
#[tokio::test]
async fn kiro_sync_report_context_applies_prompt_cache_usage_from_tracker () {
let state = AppState ::new (). expect ( "gateway state should build" );
let plan = test_kiro_sync_plan ();
let mut first_report_context = Some ( json! ({
"original_request_body" : test_kiro_sync_cacheable_request_body (),
"kiro_simulated_cache_enabled" : true ,
}));
seed_kiro_sync_report_context_input_tokens ( & plan , & mut first_report_context );
seed_kiro_sync_report_context_prompt_cache_usage ( & state , & plan , & mut first_report_context )
. await ;
let first_creation = first_report_context
. as_ref ()
. and_then ( | value | value . get ( "cache_creation_input_tokens" ))
. and_then ( Value ::as_u64 )
. unwrap_or_default ();
let first_read = first_report_context
. as_ref ()
. and_then ( | value | value . get ( "cache_read_input_tokens" ))
. and_then ( Value ::as_u64 )
. unwrap_or_default ();
assert! ( first_creation > 0 );
assert_eq! ( first_read , 0 );
let mut second_report_context = Some ( json! ({
"original_request_body" : test_kiro_sync_cacheable_request_body (),
"kiro_simulated_cache_enabled" : true ,
}));
seed_kiro_sync_report_context_input_tokens ( & plan , & mut second_report_context );
seed_kiro_sync_report_context_prompt_cache_usage ( & state , & plan , & mut second_report_context )
. await ;
let second_creation = second_report_context
. as_ref ()
. and_then ( | value | value . get ( "cache_creation_input_tokens" ))
. and_then ( Value ::as_u64 )
. unwrap_or_default ();
let second_read = second_report_context
. as_ref ()
. and_then ( | value | value . get ( "cache_read_input_tokens" ))
. and_then ( Value ::as_u64 )
. unwrap_or_default ();
assert_eq! ( second_creation , 0 );
assert! ( second_read > 0 );
}
2026-05-09 01:21:26 +08:00
#[tokio::test]
async fn json_whitespace_heartbeat_stream_prefixes_final_json () {
let ( tx , rx ) = mpsc ::channel ::< Result < Bytes , IoError >> ( 1 );
tx . send ( Ok ( Bytes ::from_static ( br #"{"data":[]}"# )))
. await
. expect ( "final body should send" );
drop ( tx );
let body = to_bytes (
Body ::from_stream ( build_json_whitespace_heartbeat_stream (
rx ,
Duration ::from_secs ( 60 ),
None ,
)),
usize ::MAX ,
)
. await
. expect ( "body should collect" );
assert! ( body . starts_with ( b " \n " ));
let parsed : Value =
serde_json ::from_slice ( & body ). expect ( "leading whitespace is valid JSON" );
assert_eq! ( parsed , json! ({ "data" : []}));
}
#[tokio::test]
async fn json_whitespace_heartbeat_stream_emits_interval_whitespace () {
let ( tx , rx ) = mpsc ::channel ::< Result < Bytes , IoError >> ( 1 );
let stream = build_json_whitespace_heartbeat_stream ( rx , Duration ::from_millis ( 5 ), None );
pin_mut! ( stream );
let first = stream
. next ()
. await
. expect ( "initial whitespace" )
. expect ( "initial whitespace ok" );
assert_eq! ( first , Bytes ::from_static ( b " \n " ));
let second = tokio ::time ::timeout ( Duration ::from_millis ( 100 ), stream . next ())
. await
. expect ( "interval heartbeat" )
. expect ( "interval heartbeat item" )
. expect ( "interval heartbeat ok" );
assert_eq! ( second , Bytes ::from_static ( b " \n " ));
tx . send ( Ok ( Bytes ::from_static ( br #"{"data":[{"b64_json":"x"}]}"# )))
. await
. expect ( "final body should send" );
let final_body = tokio ::time ::timeout ( Duration ::from_millis ( 100 ), stream . next ())
. await
. expect ( "final body" )
. expect ( "final body item" )
. expect ( "final body ok" );
assert_eq! (
serde_json ::from_slice ::< Value > ( & final_body ). expect ( "final body json" ),
json! ({ "data" : [{ "b64_json" : "x" }]})
);
}
#[test]
fn openai_image_sync_sse_parser_tracks_partial_and_completed_frames () {
let partial = parse_openai_image_sync_sse_frame (
concat! (
"event: response.image_generation_call.partial_image \n " ,
"data: { \" type \" : \" response.image_generation_call.partial_image \" , \" partial_image_index \" :0} \n\n "
)
. as_bytes (),
)
. expect ( "partial frame" );
assert_eq! (
partial . event_name ,
"response.image_generation_call.partial_image"
);
assert! ( partial . is_partial_image );
assert_eq! (
partial . client_visible_event ,
Some ( "image_generation.partial_image" )
);
let completed = parse_openai_image_sync_sse_frame (
b "data: { \" type \" : \" response.completed \" , \" response \" :{}} \n\n " ,
)
. expect ( "completed frame" );
assert_eq! ( completed . event_name , "response.completed" );
assert! ( completed . is_completed );
assert_eq! (
completed . client_visible_event ,
Some ( "image_generation.completed" )
);
}
#[test]
2026-05-10 02:10:23 +08:00
fn openai_image_sync_progress_tracks_upstream_stream_without_json_heartbeat_wrapper () {
2026-05-09 01:21:26 +08:00
let plan = test_openai_image_plan ( false );
let report_context = json! ({ "upstream_is_stream" : true });
assert! ( should_track_openai_image_sync_upstream_sse (
OPENAI_IMAGE_SYNC_PLAN_KIND ,
& plan ,
Some ( & report_context ),
));
2026-05-10 02:10:23 +08:00
assert! ( ! should_enable_openai_image_sync_json_heartbeat (
2026-05-09 01:21:26 +08:00
OPENAI_IMAGE_SYNC_PLAN_KIND ,
& plan ,
Some ( & report_context ),
));
}
#[test]
fn openai_image_sync_progress_ignores_non_stream_upstream () {
let plan = test_openai_image_plan ( false );
let report_context = json! ({ "upstream_is_stream" : false });
assert! ( ! should_track_openai_image_sync_upstream_sse (
OPENAI_IMAGE_SYNC_PLAN_KIND ,
& plan ,
Some ( & report_context ),
));
assert! ( ! should_enable_openai_image_sync_json_heartbeat (
OPENAI_IMAGE_SYNC_PLAN_KIND ,
& plan ,
Some ( & report_context ),
));
}
}