use std::collections::BTreeMap; use std::io::Error as IoError; use std::sync::Arc; use std::time::{Duration, Instant}; use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope, UPSTREAM_IS_STREAM_KEY}; use aether_contracts::{ ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResponseObservation, ExecutionResult, ExecutionTelemetry, }; use aether_data_contracts::repository::candidates::RequestCandidateStatus; use aether_scheduler_core::{ execution_error_details, parse_request_candidate_report_context, SchedulerRequestCandidateStatusUpdate, }; use aether_usage_runtime::{ build_lifecycle_usage_seed, build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed, build_usage_event_data_seed, UsageEvent, UsageEventType, }; 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; use http_body_util::BodyExt; use serde_json::{json, Value}; use tokio::sync::mpsc; use tokio::sync::Mutex; use tokio::time::MissedTickBehavior; use tracing::{debug, warn}; use crate::ai_serving::api::{ 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, }; use crate::api::response::{ attach_control_metadata_headers, build_client_response, build_client_response_from_parts, build_client_response_from_parts_with_mutator, }; use crate::clock::current_unix_ms as current_request_candidate_unix_ms; use crate::control::GatewayControlDecision; use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync; use crate::execution_runtime::grok::maybe_execute_grok_sync; use crate::execution_runtime::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, }; use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry; #[cfg(test)] use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_runtime; use crate::execution_runtime::submission::{ resolve_local_sync_error_status_code, submit_local_core_error_or_sync_finalize, }; use crate::execution_runtime::transport::{ 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, }; use crate::execution_runtime::windsurf::maybe_execute_windsurf_sync; use crate::execution_runtime::{ 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, }; use crate::log_ids::short_request_id; use crate::orchestration::{ apply_local_execution_effect, build_local_error_flow_metadata, spawn_local_oauth_success_effect, trace_upstream_response_body, with_error_flow_report_context, with_upstream_response_report_context, LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, LocalPoolErrorEffect, }; use crate::provider_pool_demand::acquire_provider_pool_in_flight_guard; use crate::request_candidate_runtime::{ ensure_execution_request_candidate_slot, record_local_request_candidate_extra_data, record_local_request_candidate_status, record_local_request_candidate_status_snapshot, snapshot_local_request_candidate_status, }; 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, }; use crate::usage::{spawn_sync_report, submit_sync_report}; use crate::video_tasks::VideoTaskSyncReportMode; use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError}; #[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, resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessBuild, LocalVideoSyncSuccessOutcome, }; 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); 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."; 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 { let telemetry = result.telemetry.as_ref()?; calibrate_candidate_first_byte_elapsed_ms( elapsed_ms_since(candidate_started_at), telemetry.elapsed_ms, telemetry.ttfb_ms, ) } #[derive(Debug)] struct SyncExecutionFailure { error_type: &'static str, message: String, status_code: Option, latency_ms: Option, fallback_kind: Option, } #[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", } } } struct SyncAttemptTerminalGuard { state: AppState, plan: ExecutionPlan, report_context: Option, request_diagnostics: Option>, candidate_started_unix_ms: u64, candidate_started_at: Instant, armed: bool, } impl SyncAttemptTerminalGuard { fn new( state: &AppState, plan: &ExecutionPlan, report_context: Option, candidate_started_unix_ms: u64, candidate_started_at: Instant, ) -> Self { Self { state: state.clone(), plan: plan.clone(), report_context, request_diagnostics: current_request_diagnostics(), candidate_started_unix_ms, candidate_started_at, 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(), self.request_diagnostics.clone(), self.candidate_started_unix_ms, self.candidate_started_at, 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(); let request_diagnostics = self.request_diagnostics.clone(); let candidate_started_unix_ms = self.candidate_started_unix_ms; let candidate_started_at = self.candidate_started_at; if let Ok(handle) = tokio::runtime::Handle::try_current() { handle.spawn(async move { record_sync_attempt_forced_terminal_state( state, plan, report_context, request_diagnostics, candidate_started_unix_ms, candidate_started_at, 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, request_diagnostics: Option>, candidate_started_unix_ms: u64, candidate_started_at: Instant, usage_event_type: UsageEventType, candidate_status: RequestCandidateStatus, status_code: u16, error_type: &'static str, error_message: impl Into, ) { let error_message = error_message.into(); let report_context = attach_request_diagnostics_to_report_context(report_context, request_diagnostics.as_ref()); let terminal_unix_ms = current_request_candidate_unix_ms(); let latency_ms = elapsed_ms_since(candidate_started_at); 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( state.usage_lifecycle_data_state().as_ref(), UsageEvent::new(usage_event_type, plan.request_id.clone(), usage_data), ) .await; } impl SyncExecutionFailure { fn from_transport(err: ExecutionRuntimeTransportError) -> Self { let fallback_kind = match &err { ExecutionRuntimeTransportError::UpstreamResponseTooLarge { .. } => { Some(SyncExecutionFailureFallbackKind::UpstreamResponseTooLarge) } ExecutionRuntimeTransportError::UpstreamResponseDecode { .. } => { Some(SyncExecutionFailureFallbackKind::UpstreamResponseDecode) } _ => None, }; Self { error_type: fallback_kind .map(SyncExecutionFailureFallbackKind::error_type) .unwrap_or("execution_runtime_unavailable"), message: err.to_string(), status_code: fallback_kind.map(|_| StatusCode::BAD_GATEWAY.as_u16()), latency_ms: None, fallback_kind, } } 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), fallback_kind: None, } } } 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>, 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>>, ) -> 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(()) } 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>, 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) } struct ImplicitSyncFinalizeOutcome { payload: GatewaySyncReportRequest, outcome: LocalCoreSyncFinalizeOutcome, } 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( state.usage_lifecycle_data_state().as_ref(), 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); state.usage_runtime.record_sync_active_immediate_async( state.usage_lifecycle_data_state().as_ref(), lifecycle_seed, ); 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, }, ); } } async fn record_sync_terminal_usage( state: &AppState, plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, candidate_started_at: Instant, candidate_first_byte_elapsed_ms: Option, ) { let report_context_with_diagnostics = attach_current_request_diagnostics_and_candidate_start_timing_to_report_context( report_context, candidate_started_at, candidate_first_byte_elapsed_ms, ); let context_seed = build_terminal_usage_context_seed( plan, report_context_with_diagnostics.as_ref().or(report_context), ); let payload_seed = build_sync_terminal_usage_payload_seed(payload); state .usage_runtime .record_sync_terminal( state.usage_lifecycle_data_state().as_ref(), context_seed, payload_seed, ) .await; } async fn record_sync_terminal_usage_and_disarm_guard( state: &AppState, plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, candidate_started_at: Instant, candidate_first_byte_elapsed_ms: Option, terminal_guard: &mut SyncAttemptTerminalGuard, ) { record_sync_terminal_usage( state, plan, report_context, payload, candidate_started_at, candidate_first_byte_elapsed_ms, ) .await; terminal_guard.disarm(); } fn with_sync_error_trace_context( report_context: Option<&serde_json::Value>, status_code: u16, headers: &BTreeMap, body_json: Option<&serde_json::Value>, body_bytes: &[u8], response_text: Option<&str>, local_failover_analysis: crate::orchestration::LocalFailoverAnalysis, ) -> Option { let body = trace_upstream_response_body(body_json, body_bytes); let upstream_context = with_upstream_response_report_context( report_context, status_code, Some(headers), body.as_ref(), 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), ) } fn build_sync_report_payload( trace_id: &str, report_kind: String, report_context: Option, status_code: u16, headers: BTreeMap, body_json: Option, body_base64: Option, telemetry: Option, ) -> GatewaySyncReportRequest { GatewaySyncReportRequest { trace_id: trace_id.to_string(), report_kind, report_context, status_code, headers, body_json, client_body_json: None, body_base64, telemetry, } } fn seed_kiro_sync_report_context_input_tokens( plan: &ExecutionPlan, report_context: &mut Option, ) { 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, ) { 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, ) { 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, ) -> Option { 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, }, ) } 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; } if !provider_api_format_is_gemini_generate_content(plan, report_context) { 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; } 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); if crate::ai_serving::gemini_generate_content_response_has_visible_output(body_json) { return None; } 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, } } fn build_invalid_provider_success_body( plan: &ExecutionPlan, report_context: Option<&Value>, message: &str, ) -> Option { 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, ) } fn provider_private_error_details(body_json: &Value) -> (Option, Option) { 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>, body_object: Option<&serde_json::Map>, keys: &[&str], ) -> Option { 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 } #[derive(Debug, Clone)] struct OpenAiImageSyncProgressSnapshot { phase: &'static str, upstream_ttfb_ms: Option, upstream_sse_frame_count: u64, last_upstream_event: Option, last_upstream_frame_at_unix_ms: Option, partial_image_count: u64, last_client_visible_event: Option, downstream_heartbeat_count: u64, last_downstream_heartbeat_at_unix_ms: Option, downstream_heartbeat_interval_ms: Option, } struct OpenAiImageSyncProgressRecorder<'a> { state: &'a AppState, plan: &'a ExecutionPlan, report_context: Option<&'a Value>, snapshot: Arc>, buffer: Vec, last_persist_at: Option, } #[derive(Clone)] struct OpenAiImageSyncJsonHeartbeatContext { state: AppState, plan: ExecutionPlan, report_context: Option, snapshot: Arc>, started_at: Instant, trace_id: String, request_id_for_log: String, candidate_id: Option, } #[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>>, ) -> 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, latency_ms: Option, 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::>(); 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; } async fn fail(&mut self, status_code: Option, elapsed_ms: u64) { self.snapshot.lock().await.phase = "failed"; self.persist( RequestCandidateStatus::Failed, status_code, 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 { 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 { 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::(&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, candidate_started_unix_ms: u64, 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>>, ) -> Result { if let Some(result) = maybe_execute_windsurf_sync(state, plan, report_context) .await .map_err(SyncExecutionFailure::from_transport)? { return Ok(result); } if !should_track_openai_image_sync_upstream_sse(plan_kind, plan, report_context) { 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); return DirectSyncExecutionRuntime::new() .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, ); spawn_local_oauth_success_effect( state_for_response_started.clone(), plan, report_context, LocalOAuthSuccessEffect { status_code: event.status_code, request_started_at_unix_ms: Some( event.response_observation.request_started_at_unix_ms, ), request_order_id: Some(&event.response_observation.request_order_id), }, ); }) .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; let status_code = err.status_code; record_openai_image_sync_failed_progress( state, plan, report_context, status_code, 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, Some(StatusCode::GATEWAY_TIMEOUT.as_u16()), 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>>, ) -> Result { 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 request_started_at_unix_ms = current_request_candidate_unix_ms(); let request_order_id = uuid::Uuid::now_v7().to_string(); let response = send_request(plan, request_body) .await .map_err(SyncExecutionFailure::from_transport)?; let ttfb_ms = started_at.elapsed().as_millis() as u64; let response_headers_observed_at_unix_ms = current_request_candidate_unix_ms(); let status_code = response.status_code(); let headers = response.headers(); spawn_local_oauth_success_effect( state.clone(), plan, report_context, LocalOAuthSuccessEffect { status_code, request_started_at_unix_ms: Some(request_started_at_unix_ms), request_order_id: Some(&request_order_id), }, ); progress.record_response_started(status_code, ttfb_ms).await; let mut body_bytes = Vec::new(); 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), ), ) })?; append_upstream_response_body_chunk(&mut body_bytes, &chunk) .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; progress .observe_chunk(&chunk, status_code, elapsed_ms) .await; } } 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, )), ) })?; append_upstream_response_body_chunk(&mut body_bytes, &chunk) .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; progress .observe_chunk(&chunk, status_code, elapsed_ms) .await; } } 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), ), ) })?; append_upstream_response_body_chunk(&mut body_bytes, &chunk) .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; progress .observe_chunk(&chunk, status_code, elapsed_ms) .await; } } } let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes) .map_err(SyncExecutionFailure::from_transport)?; 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; 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)?; Ok(ExecutionResult { request_id: plan.request_id.clone(), candidate_id: plan.candidate_id.clone(), status_code, headers, response_observation: Some(ExecutionResponseObservation { request_started_at_unix_ms, response_headers_observed_at_unix_ms, request_order_id, }), 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>, status_code: Option, elapsed_ms: u64, progress_snapshot: Option>>, ) { let mut progress = OpenAiImageSyncProgressRecorder::new(state, plan, report_context, progress_snapshot); progress.fail(status_code, elapsed_ms).await; } 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 .and_then(|value| value.get(UPSTREAM_IS_STREAM_KEY)) .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( _plan_kind: &str, _plan: &ExecutionPlan, _report_context: Option<&Value>, ) -> bool { false } #[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, report_context: Option, ) -> Result, 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::>(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), None, None, ) .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>, heartbeat_interval: Duration, heartbeat_context: Option, ) -> impl futures_util::Stream> + 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)); } } } } } pub(crate) fn build_sync_json_whitespace_heartbeat_stream( rx: mpsc::Receiver>, ) -> impl futures_util::Stream> + Send + 'static { build_json_whitespace_heartbeat_stream(rx, OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_INTERVAL, None) } pub(crate) fn build_openai_image_sync_json_whitespace_heartbeat_stream( rx: mpsc::Receiver>, ) -> impl futures_util::Stream> + Send + 'static { build_sync_json_whitespace_heartbeat_stream(rx) } async fn openai_image_sync_json_heartbeat_final_bytes( result: Result>, GatewayError>, ) -> Vec { match result { Ok(Some(response)) => match to_bytes( response.into_body(), crate::headers::max_internal_buffered_body_bytes(), ) .await { 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 { serde_json::to_vec(&json!({ "error": { "type": "aether_gateway_error", "message": message, } })) .unwrap_or_else(|_| b"{\"error\":{\"type\":\"aether_gateway_error\"}}".to_vec()) } async fn apply_sync_success_effects( state: &AppState, plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, ) { if let Some(report_context) = report_context { crate::ai_serving::persist_converted_response_history( state.runtime_state(), report_context, payload .client_body_json .as_ref() .or(payload.body_json.as_ref()), ) .await; } 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; } #[cfg(test)] enum RemoteSyncFallbackOutcome { Executed(ExecutionResult), ClientResponse(Response), 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, mut report_context: Option, ) -> Result>, GatewayError> { execute_execution_runtime_sync_impl( state, request_path, plan, trace_id, decision, plan_kind, report_kind, report_context, true, None, None, None, ) .await } #[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, report_context: Option, ) -> Result>, 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, }, }) } #[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, mut report_context: Option, allow_json_heartbeat: bool, progress_snapshot: Option>>, mut retry_scope_out: Option<&mut AiAttemptRetryScope>, mut retry_fallback_out: Option<&mut Option>>, ) -> Result>, 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); } ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await; 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()); 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()); let candidate_started_at = Instant::now(); let candidate_started_unix_secs = current_request_candidate_unix_ms(); let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); let usage_data = state.usage_lifecycle_data_state().as_ref().clone(); state .usage_runtime .record_pending_direct(&usage_data, lifecycle_seed) .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(candidate_started_unix_secs), finished_at_unix_ms: None, }, ) .await; let mut terminal_guard = SyncAttemptTerminalGuard::new( state, &plan, report_context.clone(), candidate_started_unix_secs, candidate_started_at, ); let result = (async { 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; record_sync_execution_active( state, &plan, report_context.as_ref(), candidate_started_unix_secs, ); #[cfg(not(test))] let mut result = { match maybe_execute_grok_sync(&plan, report_context.as_ref()).await { Ok(Some(result)) => result, Ok(None) => { match maybe_execute_chatgpt_web_image_sync(state, &plan, report_context.as_ref()) .await { Ok(Some(result)) => result, Ok(None) => match execute_direct_sync_runtime_candidate( state, &plan, report_context.as_ref(), trace_id, plan_kind, candidate_started_unix_secs, 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) => { 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)); maybe_store_sync_execution_failure_fallback( &err, &plan, trace_id, decision, &mut retry_scope_out, &mut retry_fallback_out, )?; warn!( event_name = "sync_execution_runtime_unavailable", log_type = "ops", trace_id = %trace_id, request_id = %plan_request_id_for_log, candidate_id = ?plan_candidate_id, provider_name, endpoint_id, key_id, model_name, candidate_index = candidate_index.as_str(), error_type = err.error_type, error = %err.message, "gateway in-process sync execution unavailable" ); let terminal_unix_secs = current_request_candidate_unix_ms(); record_local_request_candidate_status( state, &plan, report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, status_code: None, error_type: Some(failure_error_type.to_string()), error_message: Some(err.message), latency_ms: Some(failure_latency_ms), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; 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)); } return Ok(None); } }, Err(err) => { let transport_error_message = err.to_string(); warn!( event_name = "chatgpt_web_image_execution_unavailable", log_type = "ops", trace_id = %trace_id, request_id = %plan_request_id_for_log, candidate_id = ?plan_candidate_id, provider_name, endpoint_id, key_id, model_name, candidate_index = candidate_index.as_str(), error = %err, "gateway ChatGPT-Web image execution unavailable" ); let terminal_unix_secs = current_request_candidate_unix_ms(); record_local_request_candidate_status( state, &plan, report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, status_code: None, error_type: Some( "chatgpt_web_image_execution_unavailable".to_string(), ), error_message: Some(transport_error_message.clone()), latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; 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)); } return Ok(None); } } } Err(err) => { let transport_error_message = err.to_string(); warn!( event_name = "grok_execution_unavailable", log_type = "ops", trace_id = %trace_id, request_id = %plan_request_id_for_log, candidate_id = ?plan_candidate_id, provider_name, endpoint_id, key_id, model_name, candidate_index = candidate_index.as_str(), error = %err, "gateway Grok execution unavailable" ); let terminal_unix_secs = current_request_candidate_unix_ms(); record_local_request_candidate_status( state, &plan, report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("grok_execution_unavailable".to_string()), error_message: Some(transport_error_message.clone()), latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; 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)); } return Ok(None); } } }; #[cfg(test)] let mut result = { if let Some(override_fn) = state.execution_runtime_sync_override.as_ref() { match (override_fn.0)(&plan) { Ok(result) => result, Err(err) => { let transport_error_message = format!("{err:?}"); 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()), error_message: Some(transport_error_message.clone()), latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; 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)); } return Ok(None); } } } else if state .execution_runtime_override_base_url() .unwrap_or_default() .trim() .is_empty() { match maybe_execute_grok_sync(&plan, report_context.as_ref()).await { Ok(Some(result)) => result, Ok(None) => match maybe_execute_chatgpt_web_image_sync( state, &plan, report_context.as_ref(), ) .await { Ok(Some(result)) => result, Ok(None) => match execute_direct_sync_runtime_candidate( state, &plan, report_context.as_ref(), trace_id, plan_kind, candidate_started_unix_secs, 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) => { 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)); maybe_store_sync_execution_failure_fallback( &err, &plan, trace_id, decision, &mut retry_scope_out, &mut retry_fallback_out, )?; warn!( event_name = "sync_execution_runtime_unavailable", log_type = "ops", trace_id = %trace_id, request_id = %plan_request_id_for_log, candidate_id = ?plan_candidate_id, provider_name, endpoint_id, key_id, model_name, candidate_index = candidate_index.as_str(), error_type = err.error_type, error = %err.message, "gateway in-process sync execution unavailable" ); let terminal_unix_secs = current_request_candidate_unix_ms(); record_local_request_candidate_status( state, &plan, report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, status_code: None, error_type: Some(failure_error_type.to_string()), error_message: Some(err.message), latency_ms: Some(failure_latency_ms), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; 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)); } return Ok(None); } }, Err(err) => { let transport_error_message = err.to_string(); warn!( event_name = "chatgpt_web_image_execution_unavailable", log_type = "ops", trace_id = %trace_id, request_id = %plan_request_id_for_log, candidate_id = ?plan_candidate_id, provider_name, endpoint_id, key_id, model_name, candidate_index = candidate_index.as_str(), error = %err, "gateway ChatGPT-Web image execution unavailable" ); let terminal_unix_secs = current_request_candidate_unix_ms(); record_local_request_candidate_status( state, &plan, report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, status_code: None, error_type: Some( "chatgpt_web_image_execution_unavailable".to_string(), ), error_message: Some(transport_error_message.clone()), latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; 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)); } return Ok(None); } }, Err(err) => { let transport_error_message = err.to_string(); warn!( event_name = "grok_execution_unavailable", log_type = "ops", trace_id = %trace_id, request_id = %plan_request_id_for_log, candidate_id = ?plan_candidate_id, provider_name, endpoint_id, key_id, model_name, candidate_index = candidate_index.as_str(), error = %err, "gateway Grok execution unavailable" ); let terminal_unix_secs = current_request_candidate_unix_ms(); record_local_request_candidate_status( state, &plan, report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("grok_execution_unavailable".to_string()), error_message: Some(transport_error_message.clone()), latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; 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)); } return Ok(None); } } } else { let remote_execution_runtime_base_url = state .execution_runtime_override_base_url() .unwrap_or_default(); let remote_outcome = execute_sync_via_remote_execution_runtime( state, remote_execution_runtime_base_url, trace_id, decision, &plan, plan_request_id.as_str(), plan_candidate_id.as_deref(), report_context.as_ref(), candidate_started_unix_secs, candidate_started_at, ) .await?; match remote_outcome { RemoteSyncFallbackOutcome::Executed(result) => result, RemoteSyncFallbackOutcome::ClientResponse(response) => return Ok(Some(response)), RemoteSyncFallbackOutcome::Unavailable => return Ok(None), } } }; let mut candidate_first_byte_elapsed_ms = calibrated_sync_candidate_first_byte_elapsed_ms(candidate_started_at, &result); let initial_response_observed_at_unix_ms = current_request_candidate_unix_ms(); let mut provider_response_observation = result .response_observation .clone() .unwrap_or(ExecutionResponseObservation { request_started_at_unix_ms: candidate_started_unix_secs, response_headers_observed_at_unix_ms: initial_response_observed_at_unix_ms, request_order_id: uuid::Uuid::now_v7().to_string(), }); 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 { spawn_local_oauth_success_effect( state.clone(), &plan, report_context.as_ref(), LocalOAuthSuccessEffect { status_code: result.status_code, request_started_at_unix_ms: Some( provider_response_observation.request_started_at_unix_ms, ), request_order_id: Some(&provider_response_observation.request_order_id), }, ); let result_latency_ms = result .telemetry .as_ref() .and_then(|telemetry| telemetry.elapsed_ms); let mut headers = std::mem::take(&mut result.headers); let (body_bytes, mut body_json, body_base64) = decode_execution_result_body(result.body.take(), &mut headers)?; if let Some(message) = invalid_gemini_provider_success_message( &plan, report_context.as_ref(), result.status_code, body_json.as_ref(), ) .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(), ) }) { result.status_code = StatusCode::BAD_GATEWAY.as_u16(); result.error = Some(invalid_gemini_provider_success_execution_error(message)); 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()); } } let (mut result_error_type, mut result_error_message) = execution_error_details(result.error.as_ref(), body_json.as_ref()); if result.status_code < 400 && body_json.is_none() { 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); } } 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, report_context.as_ref(), Some(provider_response_observation.request_started_at_unix_ms), Some(&provider_response_observation.request_order_id), ) .await { oauth_retry_attempted = true; let retry_started_at_unix_ms = current_request_candidate_unix_ms(); let retry_request_order_id = uuid::Uuid::now_v7().to_string(); match crate::execution_runtime::execute_execution_runtime_sync_plan( state, Some(trace_id), &plan, ) .await { Ok(retry_result) => { let retry_response_observed_at_unix_ms = current_request_candidate_unix_ms(); provider_response_observation = retry_result .response_observation .clone() .unwrap_or(ExecutionResponseObservation { request_started_at_unix_ms: retry_started_at_unix_ms, response_headers_observed_at_unix_ms: retry_response_observed_at_unix_ms, request_order_id: retry_request_order_id, }); candidate_first_byte_elapsed_ms = calibrated_sync_candidate_first_byte_elapsed_ms( candidate_started_at, &retry_result, ); 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, ); }; let mut report_context = attach_provider_response_headers_to_report_context( report_context, &headers, provider_response_observation.request_started_at_unix_ms, provider_response_observation.response_headers_observed_at_unix_ms, &provider_response_observation.request_order_id, ); if result.status_code >= 400 { apply_local_execution_effect( state, 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(), }), ) .await; } if matches!( local_failover_analysis.decision, LocalFailoverDecision::RetryNextCandidate ) { 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(), )?); } } let terminal_unix_secs = current_request_candidate_unix_ms(); let error_trace_report_context = with_sync_error_trace_context( report_context.as_ref(), result.status_code, &headers, body_json.as_ref(), &body_bytes, local_failover_response_text.as_deref(), local_failover_analysis, ); record_local_request_candidate_status( state, &plan, error_trace_report_context .as_ref() .or(report_context.as_ref()), 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, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; warn!( event_name = "local_sync_candidate_retry_scheduled", log_type = "event", trace_id = %trace_id, request_id = %plan_request_id_for_log, status_code = result.status_code, provider_name, endpoint_id, key_id, model_name, candidate_index = candidate_index.as_str(), "gateway local sync decision retrying next candidate after retryable execution runtime result" ); return Ok(None); } let status_code = result.status_code; let has_body_bytes = body_base64.is_some(); if (200..300).contains(&status_code) { seed_kiro_sync_simulated_cache_enabled(state, &plan, &mut report_context).await; if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) { seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context); } seed_kiro_sync_report_context_prompt_cache_usage(state, &plan, &mut report_context).await; } let mut client_headers = headers.clone(); apply_endpoint_response_header_rules(state, &plan, &mut client_headers, body_json.as_ref()) .await?; 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()); let implicit_finalize = if !explicit_finalize && mapped_error_finalize_kind.is_none() { maybe_build_implicit_sync_finalize_outcome( trace_id, decision, plan_kind, &report_context, status_code, &client_headers, &body_json, &body_base64, &result.telemetry, )? } else { None }; 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(), ) { let terminal_unix_secs = current_request_candidate_unix_ms(); let error_trace_report_context = with_sync_error_trace_context( report_context.as_ref(), result.status_code, &headers, body_json.as_ref(), &body_bytes, local_failover_response_text.as_deref(), local_failover_analysis, ); record_local_request_candidate_status( state, &plan, error_trace_report_context .as_ref() .or(report_context.as_ref()), 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, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(None); } let terminal_unix_secs = current_request_candidate_unix_ms(); let error_flow_report_context = (result.status_code >= 400) .then(|| { with_sync_error_trace_context( report_context.as_ref(), result.status_code, &headers, body_json.as_ref(), &body_bytes, local_failover_response_text.as_deref(), local_failover_analysis, ) }) .flatten(); record_local_request_candidate_status( state, &plan, error_flow_report_context .as_ref() .or(report_context.as_ref()), 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, started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; 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()) .or(Some(plan_request_id.as_str())); let request_id_for_log = short_request_id(request_id.unwrap_or("-")); let candidate_id = candidate_id_owned .as_deref() .or(plan_candidate_id.as_deref()); 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( state, &plan, implicit_finalize.payload.report_context.as_ref(), usage_payload, ) .await; record_sync_terminal_usage_and_disarm_guard( state, &plan, implicit_finalize.payload.report_context.as_ref(), usage_payload, candidate_started_at, candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; 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, )?)); } let finalize_report_kind = if explicit_finalize { report_kind.clone() } else { mapped_error_finalize_kind }; if let Some(finalize_report_kind) = finalize_report_kind { let mut payload = build_sync_report_payload( trace_id, finalize_report_kind, report_context, status_code, client_headers, body_json, body_base64, telemetry, ); if let Some(outcome) = maybe_build_sync_finalize_outcome(trace_id, decision, &payload)? { let usage_payload = outcome.background_report.as_ref().unwrap_or(&payload); if status_code < 400 { apply_sync_success_effects( state, &plan, payload.report_context.as_ref(), usage_payload, ) .await; } record_sync_terminal_usage_and_disarm_guard( state, &plan, payload.report_context.as_ref(), usage_payload, candidate_started_at, candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; if let Some(report_payload) = 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 = %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, )?)); } let mut payload = match maybe_build_local_video_success_outcome( trace_id, decision, payload, &state.video_tasks, &plan, )? { 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; record_sync_terminal_usage_and_disarm_guard( state, &plan, original_report_context.as_ref(), &report_payload, candidate_started_at, candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; if let Some(snapshot) = local_task_snapshot { let _ = state.upsert_video_task_snapshot(&snapshot).await?; state.video_tasks.record_snapshot(snapshot); } match report_mode { VideoTaskSyncReportMode::InlineSync => { submit_sync_report(state, report_payload).await?; } VideoTaskSyncReportMode::Background => { spawn_sync_report(state.clone(), report_payload); } } return Ok(Some(attach_control_metadata_headers( response, request_id, candidate_id, )?)); } LocalVideoSyncSuccessBuild::NotHandled(payload) => payload, }; if let Some(response) = maybe_build_local_sync_finalize_response(trace_id, decision, &payload)? { 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; record_sync_terminal_usage_and_disarm_guard( state, &plan, payload.report_context.as_ref(), &payload, candidate_started_at, candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; 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?; } 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); } else { warn!( event_name = "local_video_finalize_missing_success_report_mapping", log_type = "ops", trace_id = %trace_id, request_id = %request_id_for_log, candidate_id = ?candidate_id, 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)? { 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(); } record_sync_terminal_usage_and_disarm_guard( state, &plan, payload.report_context.as_ref(), &payload, candidate_started_at, candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; if background_error_report_kind.is_some() { spawn_sync_report(state.clone(), payload); } else { warn!( event_name = "local_video_finalize_missing_error_report_mapping", log_type = "ops", trace_id = %trace_id, request_id = %request_id_for_log, candidate_id = ?candidate_id, 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, )?)); } record_sync_terminal_usage_and_disarm_guard( state, &plan, payload.report_context.as_ref(), &payload, candidate_started_at, candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; 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, )?)); } let usage_payload = build_sync_report_payload( trace_id, report_kind.unwrap_or_default(), report_context, status_code, client_headers, body_json, body_base64, telemetry, ); if status_code < 400 { apply_sync_success_effects( state, &plan, usage_payload.report_context.as_ref(), &usage_payload, ) .await; } record_sync_terminal_usage_and_disarm_guard( state, &plan, usage_payload.report_context.as_ref(), &usage_payload, candidate_started_at, candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; 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() { 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); } } Ok(Some(response)) }) .await; if let Err(error) = result.as_ref() { terminal_guard.fail_and_disarm(error).await; } else { terminal_guard.disarm(); } result } #[allow(clippy::too_many_arguments)] // mirrors sync execution context fn maybe_build_implicit_sync_finalize_outcome( trace_id: &str, decision: &GatewayControlDecision, plan_kind: &str, report_context: &Option, status_code: u16, headers: &BTreeMap, body_json: &Option, body_base64: &Option, telemetry: &Option, ) -> Result, GatewayError> { if status_code >= 400 || body_json.is_some() || body_base64.is_none() { return Ok(None); } let Some(report_kind) = implicit_sync_finalize_report_kind(plan_kind) else { return Ok(None); }; let payload = GatewaySyncReportRequest { trace_id: trace_id.to_string(), report_kind: report_kind.to_string(), report_context: report_context.clone(), status_code, headers: headers.clone(), body_json: body_json.clone(), client_body_json: None, body_base64: body_base64.clone(), telemetry: telemetry.clone(), }; 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>, candidate_started_unix_secs: u64, candidate_started_at: Instant, ) -> Result { let remote_request_started_at_unix_ms = current_request_candidate_unix_ms(); let remote_request_order_id = uuid::Uuid::now_v7().to_string(); 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!( event_name = "sync_execution_runtime_remote_unavailable", log_type = "ops", trace_id = %trace_id, request_id = %short_request_id(plan_request_id), candidate_id = ?plan_candidate_id, error = ?err, "gateway remote execution runtime sync unavailable" ); 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:?}")), latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(RemoteSyncFallbackOutcome::Unavailable); } }; if response.status() != http::StatusCode::OK { let terminal_unix_secs = current_request_candidate_unix_ms(); record_local_request_candidate_status( state, plan, report_context, 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() )), latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; return Ok(RemoteSyncFallbackOutcome::ClientResponse( attach_control_metadata_headers( build_client_response(response, trace_id, Some(decision))?, Some(plan_request_id), plan_candidate_id, )?, )); } let remote_response_observed_at_unix_ms = current_request_candidate_unix_ms(); let mut result = response .json::() .await .map_err(|err| GatewayError::Internal(err.to_string()))?; result .response_observation .get_or_insert(ExecutionResponseObservation { request_started_at_unix_ms: remote_request_started_at_unix_ms, response_headers_observed_at_unix_ms: remote_response_observed_at_unix_ms, request_order_id: remote_request_order_id, }); Ok(RemoteSyncFallbackOutcome::Executed(result)) } #[cfg(test)] mod tests { use super::*; 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; use futures_util::{pin_mut, StreamExt as _}; use serde_json::json; use std::collections::BTreeMap; use std::sync::Arc; use tokio::io::{AsyncReadExt, AsyncWriteExt}; 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, } } 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 } 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) } #[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"); } 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"}] }) } #[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")); } #[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()); } #[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()); } #[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; let candidate_started_at = Instant::now(); let started_at = current_request_candidate_unix_ms(); state.usage_runtime.record_pending( state.usage_lifecycle_data_state().as_ref(), 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; 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; 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"); assert_eq!(stored_usage.billing_status, "void"); assert_eq!(stored_usage.status_code, Some(499)); assert_eq!(stored_usage.error_category.as_deref(), Some("cancelled")); 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()); 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") ); } #[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\ncontent-type: application/json\r\ncontent-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( state.usage_lifecycle_data_state().as_ref(), 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\ncontent-type: application/json\r\ncontent-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(); } #[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); } #[tokio::test] async fn json_whitespace_heartbeat_stream_prefixes_final_json() { let (tx, rx) = mpsc::channel::>(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::>(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::(&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] fn openai_image_sync_progress_tracks_upstream_stream_without_json_heartbeat_wrapper() { 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), )); assert!(!should_enable_openai_image_sync_json_heartbeat( 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), )); } }