use std::collections::{BTreeMap, VecDeque}; use std::future::Future; use std::io::Error as IoError; use std::pin::Pin; use std::time::Instant; 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 serde_json::{json, Value}; use tokio::sync::mpsc; use crate::ai_serving::api::{ build_core_error_body_for_client_format, build_local_gemini_files_stream_attempt_source_for_kind, build_local_gemini_files_sync_attempt_source_for_kind, build_local_image_stream_attempt_source_for_kind, build_local_image_sync_attempt_source_for_kind, build_local_openai_chat_stream_attempt_source_for_kind, build_local_openai_chat_stream_plan_and_reports_for_kind, build_local_openai_chat_sync_attempt_source_for_kind, build_local_openai_chat_sync_plan_and_reports_for_kind, build_local_openai_responses_stream_attempt_source_for_kind, build_local_openai_responses_stream_plan_and_reports_for_kind, build_local_openai_responses_sync_attempt_source_for_kind, build_local_openai_responses_sync_plan_and_reports_for_kind, build_local_same_format_stream_attempt_source, build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_attempt_source, build_local_same_format_sync_plan_and_reports, build_local_video_sync_attempt_source_for_kind, build_standard_family_stream_attempt_source, build_standard_family_sync_attempt_source, parse_direct_request_body, resolve_claude_stream_spec, resolve_claude_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec, resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec, set_local_openai_chat_execution_exhausted_diagnostic, set_local_openai_image_execution_exhausted_diagnostic, AiStreamAttempt, AiSyncAttempt, LocalCoreSyncErrorKind, LocalStandardSpec, EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION, }; use crate::ai_serving::LocalExecutionAttemptSource; use crate::api::response::{ attach_control_metadata_headers, build_client_response_from_parts_with_mutator, }; use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS}; use crate::control::GatewayControlDecision; use crate::execution_runtime::sync::{ build_openai_image_sync_json_whitespace_heartbeat_stream, build_sync_json_whitespace_heartbeat_stream, }; use crate::executor::candidate_loop::{ execute_stream_attempt_source_with_transfer_tracker, execute_sync_attempt_source, execute_sync_attempt_source_with_transfer_tracker, execute_sync_plan_and_reports_with_transfer_tracker, ProviderTransferTracker, }; use crate::executor::{ record_failed_usage_for_exhausted_request, LocalExecutionExhaustion, LocalExecutionRequestOutcome, }; use crate::request_diagnostics::{current_request_diagnostics, scope_request_diagnostics_with}; use crate::stage_metrics::observe_gateway_stage_ms; use crate::{AiExecutionDecision, AppState, GatewayError}; const OPENAI_IMAGE_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS: u16 = 502; const OPENAI_IMAGE_SYNC_HEARTBEAT_EXHAUSTED_STATUS: u16 = 503; const OPENAI_IMAGE_SYNC_HEARTBEAT_ERROR_MESSAGE_LIMIT: usize = 4096; const STANDARD_TEXT_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS: u16 = 502; const STANDARD_TEXT_SYNC_HEARTBEAT_EXHAUSTED_STATUS: u16 = 503; const STANDARD_TEXT_SYNC_HEARTBEAT_ERROR_MESSAGE_LIMIT: usize = 4096; pub(crate) async fn maybe_execute_sync_local_path( state: &AppState, parts: &http::request::Parts, body_bytes: &axum::body::Bytes, trace_id: &str, decision: &GatewayControlDecision, ) -> Result { super::maybe_execute_via_sync_decision_path(state, parts, body_bytes, trace_id, decision).await } pub(crate) async fn maybe_execute_stream_local_path( state: &AppState, parts: &http::request::Parts, body_bytes: &axum::body::Bytes, trace_id: &str, decision: &GatewayControlDecision, ) -> Result { super::maybe_execute_via_stream_decision_path(state, parts, body_bytes, trace_id, decision) .await } pub(crate) async fn maybe_execute_sync_via_local_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some((attempt_source, candidate_count)) = build_local_openai_chat_sync_attempt_source_for_kind( state, parts, trace_id, decision, body_json, plan_kind, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; if standard_text_sync_heartbeat_should_wrap( plan_kind, attempt_source.routing_execution_policy(), ) { let parts_for_task = parts.clone(); let body_json_for_task = body_json.clone(); let transfer_tracker_for_task = transfer_tracker.clone(); return Ok(LocalExecutionRequestOutcome::responded( build_standard_text_sync_heartbeat_shell_response( state.clone(), parts_for_task, trace_id.to_string(), decision.clone(), plan_kind.to_string(), move |state, parts, trace_id, decision, plan_kind, started_at| async move { let Some((attempt_source, candidate_count)) = build_local_openai_chat_sync_attempt_source_for_kind( &state, &parts, trace_id.as_str(), &decision, &body_json_for_task, plan_kind.as_str(), ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; let outcome = execute_sync_attempt_source_with_transfer_tracker::( &state, &parts, trace_id.as_str(), &decision, plan_kind.as_str(), attempt_source, &transfer_tracker_for_task, ) .await?; match outcome { LocalExecutionRequestOutcome::Exhausted(exhaustion) => { set_local_openai_chat_execution_exhausted_diagnostic( &state, trace_id.as_str(), &decision, plan_kind.as_str(), &body_json_for_task, candidate_count, ); record_standard_text_sync_heartbeat_exhaustion( &state, exhaustion, &started_at, ) .await; Ok(LocalExecutionRequestOutcome::NoPath) } outcome => Ok(outcome), } }, )?, )); } let outcome = execute_sync_attempt_source_with_transfer_tracker::( state, parts, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await?; if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome { set_local_openai_chat_execution_exhausted_diagnostic( state, trace_id, decision, plan_kind, body_json, candidate_count, ); } Ok(outcome) } pub(crate) async fn maybe_execute_stream_via_local_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let attempt_source_started_at = std::time::Instant::now(); let attempt_source = build_local_openai_chat_stream_attempt_source_for_kind( state, parts, trace_id, decision, body_json, plan_kind, ) .await; observe_gateway_stage_ms( "stream_openai_chat_attempt_source_init", attempt_source_started_at.elapsed().as_millis() as u64, ); let Some((attempt_source, candidate_count)) = attempt_source? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; let attempt_source_execute_started_at = std::time::Instant::now(); let outcome = execute_stream_attempt_source_with_transfer_tracker::( state, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await; observe_gateway_stage_ms( "stream_openai_chat_attempt_source_execute", attempt_source_execute_started_at.elapsed().as_millis() as u64, ); let outcome = outcome?; if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome { set_local_openai_chat_execution_exhausted_diagnostic( state, trace_id, decision, plan_kind, body_json, candidate_count, ); } Ok(outcome) } pub(crate) async fn maybe_execute_sync_via_local_openai_responses_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some((attempt_source, _candidate_count)) = build_local_openai_responses_sync_attempt_source_for_kind( state, parts, trace_id, decision, body_json, plan_kind, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; if standard_text_sync_heartbeat_should_wrap( plan_kind, attempt_source.routing_execution_policy(), ) { let parts_for_task = parts.clone(); let body_json_for_task = body_json.clone(); let transfer_tracker_for_task = transfer_tracker.clone(); return Ok(LocalExecutionRequestOutcome::responded( build_standard_text_sync_heartbeat_shell_response( state.clone(), parts_for_task, trace_id.to_string(), decision.clone(), plan_kind.to_string(), move |state, parts, trace_id, decision, plan_kind, started_at| async move { let Some((attempt_source, _candidate_count)) = build_local_openai_responses_sync_attempt_source_for_kind( &state, &parts, trace_id.as_str(), &decision, &body_json_for_task, plan_kind.as_str(), ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; let outcome = execute_sync_attempt_source_with_transfer_tracker::( &state, &parts, trace_id.as_str(), &decision, plan_kind.as_str(), attempt_source, &transfer_tracker_for_task, ) .await?; match outcome { LocalExecutionRequestOutcome::Exhausted(exhaustion) => { record_standard_text_sync_heartbeat_exhaustion( &state, exhaustion, &started_at, ) .await; Ok(LocalExecutionRequestOutcome::NoPath) } outcome => Ok(outcome), } }, )?, )); } execute_sync_attempt_source_with_transfer_tracker::( state, parts, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await } pub(crate) async fn maybe_execute_stream_via_local_openai_responses_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some((attempt_source, _candidate_count)) = build_local_openai_responses_stream_attempt_source_for_kind( state, parts, trace_id, decision, body_json, plan_kind, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; execute_stream_attempt_source_with_transfer_tracker::( state, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await } pub(crate) async fn maybe_execute_sync_via_standard_family_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, resolve_sync_spec: fn(&str) -> Option, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some(spec) = resolve_sync_spec(plan_kind) else { return Ok(LocalExecutionRequestOutcome::NoPath); }; let Some((attempt_source, _candidate_count)) = build_standard_family_sync_attempt_source( state, parts, trace_id, decision, body_json, spec, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; if standard_text_sync_heartbeat_should_wrap( plan_kind, attempt_source.routing_execution_policy(), ) { let parts_for_task = parts.clone(); let body_json_for_task = body_json.clone(); let transfer_tracker_for_task = transfer_tracker.clone(); return Ok(LocalExecutionRequestOutcome::responded( build_standard_text_sync_heartbeat_shell_response( state.clone(), parts_for_task, trace_id.to_string(), decision.clone(), plan_kind.to_string(), move |state, parts, trace_id, decision, plan_kind, started_at| async move { let Some((attempt_source, _candidate_count)) = build_standard_family_sync_attempt_source( &state, &parts, trace_id.as_str(), &decision, &body_json_for_task, spec, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; let outcome = execute_sync_attempt_source_with_transfer_tracker::( &state, &parts, trace_id.as_str(), &decision, plan_kind.as_str(), attempt_source, &transfer_tracker_for_task, ) .await?; match outcome { LocalExecutionRequestOutcome::Exhausted(exhaustion) => { record_standard_text_sync_heartbeat_exhaustion( &state, exhaustion, &started_at, ) .await; Ok(LocalExecutionRequestOutcome::NoPath) } outcome => Ok(outcome), } }, )?, )); } execute_sync_attempt_source_with_transfer_tracker::( state, parts, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await } pub(crate) async fn maybe_execute_stream_via_standard_family_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, resolve_stream_spec: fn(&str) -> Option, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some(spec) = resolve_stream_spec(plan_kind) else { return Ok(LocalExecutionRequestOutcome::NoPath); }; let Some((attempt_source, _candidate_count)) = build_standard_family_stream_attempt_source( state, parts, trace_id, decision, body_json, spec, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; execute_stream_attempt_source_with_transfer_tracker::( state, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await } pub(crate) async fn maybe_execute_sync_via_local_standard_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let mut exhausted = None; match maybe_execute_sync_via_standard_family_decision( state, parts, trace_id, decision, body_json, plan_kind, resolve_claude_sync_spec, transfer_tracker, ) .await? { LocalExecutionRequestOutcome::Responded(response) => { return Ok(LocalExecutionRequestOutcome::Responded(response)); } LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome), LocalExecutionRequestOutcome::NoPath => {} } match maybe_execute_sync_via_standard_family_decision( state, parts, trace_id, decision, body_json, plan_kind, resolve_gemini_sync_spec, transfer_tracker, ) .await? { LocalExecutionRequestOutcome::Responded(response) => { Ok(LocalExecutionRequestOutcome::Responded(response)) } LocalExecutionRequestOutcome::Exhausted(outcome) => { Ok(LocalExecutionRequestOutcome::Exhausted(outcome)) } LocalExecutionRequestOutcome::NoPath => Ok(exhausted .map(LocalExecutionRequestOutcome::Exhausted) .unwrap_or(LocalExecutionRequestOutcome::NoPath)), } } pub(crate) async fn maybe_execute_stream_via_local_standard_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let mut exhausted = None; match maybe_execute_stream_via_standard_family_decision( state, parts, trace_id, decision, body_json, plan_kind, resolve_claude_stream_spec, transfer_tracker, ) .await? { LocalExecutionRequestOutcome::Responded(response) => { return Ok(LocalExecutionRequestOutcome::Responded(response)); } LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome), LocalExecutionRequestOutcome::NoPath => {} } match maybe_execute_stream_via_standard_family_decision( state, parts, trace_id, decision, body_json, plan_kind, resolve_gemini_stream_spec, transfer_tracker, ) .await? { LocalExecutionRequestOutcome::Responded(response) => { Ok(LocalExecutionRequestOutcome::Responded(response)) } LocalExecutionRequestOutcome::Exhausted(outcome) => { Ok(LocalExecutionRequestOutcome::Exhausted(outcome)) } LocalExecutionRequestOutcome::NoPath => Ok(exhausted .map(LocalExecutionRequestOutcome::Exhausted) .unwrap_or(LocalExecutionRequestOutcome::NoPath)), } } pub(crate) async fn maybe_execute_sync_via_local_same_format_provider_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some(spec) = resolve_local_same_format_sync_spec(plan_kind) else { return Ok(LocalExecutionRequestOutcome::NoPath); }; let Some((attempt_source, _candidate_count)) = build_local_same_format_sync_attempt_source( state, parts, trace_id, decision, body_json, spec, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; if standard_text_sync_heartbeat_should_wrap( plan_kind, attempt_source.routing_execution_policy(), ) { let parts_for_task = parts.clone(); let body_json_for_task = body_json.clone(); let transfer_tracker_for_task = transfer_tracker.clone(); return Ok(LocalExecutionRequestOutcome::responded( build_standard_text_sync_heartbeat_shell_response( state.clone(), parts_for_task, trace_id.to_string(), decision.clone(), plan_kind.to_string(), move |state, parts, trace_id, decision, plan_kind, started_at| async move { let Some((attempt_source, _candidate_count)) = build_local_same_format_sync_attempt_source( &state, &parts, trace_id.as_str(), &decision, &body_json_for_task, spec, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; let outcome = execute_sync_attempt_source_with_transfer_tracker::( &state, &parts, trace_id.as_str(), &decision, plan_kind.as_str(), attempt_source, &transfer_tracker_for_task, ) .await?; match outcome { LocalExecutionRequestOutcome::Exhausted(exhaustion) => { record_standard_text_sync_heartbeat_exhaustion( &state, exhaustion, &started_at, ) .await; Ok(LocalExecutionRequestOutcome::NoPath) } outcome => Ok(outcome), } }, )?, )); } execute_sync_attempt_source_with_transfer_tracker::( state, parts, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await } pub(crate) async fn maybe_execute_stream_via_local_same_format_provider_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, body_json: &serde_json::Value, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some(spec) = resolve_local_same_format_stream_spec(plan_kind) else { return Ok(LocalExecutionRequestOutcome::NoPath); }; let Some((attempt_source, _candidate_count)) = build_local_same_format_stream_attempt_source( state, parts, trace_id, decision, body_json, spec, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; execute_stream_attempt_source_with_transfer_tracker::( state, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await } pub(crate) async fn maybe_execute_sync_via_local_gemini_files_decision( state: &AppState, parts: &http::request::Parts, body_json: &serde_json::Value, body_base64: Option<&str>, body_is_empty: bool, trace_id: &str, decision: &GatewayControlDecision, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some((attempt_source, _candidate_count)) = build_local_gemini_files_sync_attempt_source_for_kind( state, parts, body_json, body_base64, body_is_empty, trace_id, decision, plan_kind, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; execute_sync_attempt_source_with_transfer_tracker::( state, parts, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await } fn standard_text_sync_heartbeat_applies_to_plan_kind(plan_kind: &str) -> bool { matches!( plan_kind, "openai_chat_sync" | "openai_responses_sync" | "openai_responses_compact_sync" | "claude_chat_sync" | "claude_cli_sync" | "gemini_chat_sync" | "gemini_cli_sync" ) } fn standard_text_sync_heartbeat_should_wrap( plan_kind: &str, execution_policy: Option, ) -> bool { standard_text_sync_heartbeat_applies_to_plan_kind(plan_kind) && execution_policy.is_some_and(|policy| policy.enable_cf_heartbeat) } fn standard_text_sync_heartbeat_client_api_format_for_plan_kind(plan_kind: &str) -> &'static str { match plan_kind { "openai_responses_sync" => "openai:responses", "openai_responses_compact_sync" => "openai:responses:compact", "claude_chat_sync" | "claude_cli_sync" => "claude:messages", "gemini_chat_sync" | "gemini_cli_sync" => "gemini:generate_content", _ => "openai:chat", } } fn build_standard_text_sync_heartbeat_shell_response( state: AppState, parts: http::request::Parts, trace_id: String, decision: GatewayControlDecision, plan_kind: String, execute: F, ) -> Result, GatewayError> where F: FnOnce( AppState, http::request::Parts, String, GatewayControlDecision, String, Instant, ) -> Fut + Send + 'static, Fut: std::future::Future> + Send + 'static, { let request_id = (!trace_id.trim().is_empty()).then(|| trace_id.clone()); let client_api_format = standard_text_sync_heartbeat_client_api_format_for_plan_kind(plan_kind.as_str()) .to_string(); let redaction_slot = parts .extensions .get::() .cloned(); let trace_id_for_response = trace_id.clone(); let decision_for_response = decision.clone(); let started_at = Instant::now(); let (tx, rx) = mpsc::channel::>(1); let request_diagnostics = current_request_diagnostics(); let cancel_on_disconnect = crate::request_lifecycle::cancel_on_client_disconnect(); tokio::spawn(async move { scope_request_diagnostics_with(request_diagnostics, async move { let completion = standard_text_sync_heartbeat_final_bytes( client_api_format.as_str(), redaction_slot.as_ref(), tokio::select! { biased; _ = tx.closed(), if cancel_on_disconnect => return, result = execute(state, parts, trace_id, decision, plan_kind, started_at) => result, }, ); let bytes = completion.await; let _ = tx.send(Ok(Bytes::from(bytes))).await; }) .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_sync_json_whitespace_heartbeat_stream(rx)), 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, request_id.as_deref(), None) } async fn record_standard_text_sync_heartbeat_exhaustion( state: &AppState, exhaustion: LocalExecutionExhaustion, started_at: &Instant, ) { record_failed_usage_for_exhausted_request( state, exhaustion, started_at, "Standard text sync heartbeat exhausted all local candidates", EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, None, ) .await; } async fn standard_text_sync_heartbeat_final_bytes( client_api_format: &str, redaction_slot: Option<&crate::privacy::RedactionSessionSlot>, result: Result, ) -> Vec { match result { Ok(LocalExecutionRequestOutcome::Responded(response)) => { standard_text_sync_heartbeat_response_body_bytes( client_api_format, redaction_slot, response, ) .await } Ok(LocalExecutionRequestOutcome::Exhausted(_)) | Ok(LocalExecutionRequestOutcome::NoPath) => standard_text_sync_heartbeat_error_body( client_api_format, STANDARD_TEXT_SYNC_HEARTBEAT_EXHAUSTED_STATUS, "standard text sync exhausted all local candidates", ), Err(_err) => standard_text_sync_heartbeat_error_body( client_api_format, STANDARD_TEXT_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, "internal gateway error while executing request", ), } } async fn standard_text_sync_heartbeat_response_body_bytes( client_api_format: &str, redaction_slot: Option<&crate::privacy::RedactionSessionSlot>, response: Response, ) -> Vec { let status_code = response.status().as_u16(); let (parts, body) = response.into_parts(); match to_bytes(body, crate::headers::max_internal_buffered_body_bytes()).await { Ok(bytes) => { let body = match standard_text_sync_heartbeat_restore_response_body( redaction_slot, &parts.headers, bytes.as_ref(), ) { Ok(body) => body, Err(_err) => { return standard_text_sync_heartbeat_error_body( client_api_format, STANDARD_TEXT_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, "internal gateway error while restoring response", ); } }; if (200..300).contains(&status_code) && !body.is_empty() { return body; } if !(200..300).contains(&status_code) { return standard_text_sync_heartbeat_error_body_from_response( client_api_format, status_code, body.as_ref(), ); } standard_text_sync_heartbeat_error_body( client_api_format, STANDARD_TEXT_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, "empty standard text sync response", ) } Err(_err) => standard_text_sync_heartbeat_error_body( client_api_format, STANDARD_TEXT_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, "internal gateway error while reading response", ), } } fn standard_text_sync_heartbeat_restore_response_body( redaction_slot: Option<&crate::privacy::RedactionSessionSlot>, headers: &http::HeaderMap, body: &[u8], ) -> Result, GatewayError> { let Some(redaction_slot) = redaction_slot else { return Ok(body.to_vec()); }; let candidate_id = headers .get(CONTROL_CANDIDATE_ID_HEADER) .and_then(|value| value.to_str().ok()) .map(str::trim) .filter(|value| !value.is_empty()); let Some(session) = redaction_slot.take_for_candidate(candidate_id) else { return Ok(body.to_vec()); }; let mut header_values = headers .iter() .map(|(name, value)| { ( name.as_str().to_string(), value.to_str().unwrap_or_default().to_string(), ) }) .collect::>(); crate::privacy::restore_sync_response_body(&mut header_values, body, &session) .map(|restored| restored.body) } fn standard_text_sync_heartbeat_error_body_from_response( client_api_format: &str, status_code: u16, body: &[u8], ) -> Vec { if let Ok(mut value) = serde_json::from_slice::(body) { if standard_text_sync_heartbeat_insert_upstream_status(&mut value, status_code) { return serde_json::to_vec(&value).unwrap_or_else(|_| { standard_text_sync_heartbeat_error_body( client_api_format, status_code, &format!("upstream returned status {status_code}"), ) }); } } let message = standard_text_sync_heartbeat_error_message_from_body(status_code, body); standard_text_sync_heartbeat_error_body(client_api_format, status_code, message.as_str()) } fn standard_text_sync_heartbeat_insert_upstream_status( value: &mut Value, status_code: u16, ) -> bool { let Some(error) = value.get_mut("error").and_then(Value::as_object_mut) else { return false; }; error.insert("upstream_status".to_string(), Value::from(status_code)); error .entry("message".to_string()) .or_insert_with(|| Value::String(format!("upstream returned status {status_code}"))); true } fn standard_text_sync_heartbeat_error_message_from_body(status_code: u16, body: &[u8]) -> String { let text = String::from_utf8_lossy(body).trim().to_string(); if text.is_empty() { return format!("upstream returned status {status_code}"); } text.chars() .take(STANDARD_TEXT_SYNC_HEARTBEAT_ERROR_MESSAGE_LIMIT) .collect() } fn standard_text_sync_heartbeat_error_body( client_api_format: &str, status_code: u16, message: &str, ) -> Vec { let mut body = build_core_error_body_for_client_format( client_api_format, message, Some("upstream_error"), standard_text_sync_heartbeat_error_kind(status_code), ) .unwrap_or_else(|| { json!({ "error": { "type": "upstream_error", "message": message, "code": status_code, } }) }); if !standard_text_sync_heartbeat_insert_upstream_status(&mut body, status_code) { body = json!({ "error": { "type": "upstream_error", "message": message, "code": status_code, "upstream_status": status_code, } }); } serde_json::to_vec(&body).unwrap_or_else(|_| { format!( "{{\"error\":{{\"type\":\"upstream_error\",\"code\":{status_code},\"upstream_status\":{status_code}}}}}" ) .into_bytes() }) } fn standard_text_sync_heartbeat_error_kind(status_code: u16) -> LocalCoreSyncErrorKind { match status_code { 400 => LocalCoreSyncErrorKind::InvalidRequest, 401 => LocalCoreSyncErrorKind::Authentication, 403 => LocalCoreSyncErrorKind::PermissionDenied, 404 => LocalCoreSyncErrorKind::NotFound, 413 => LocalCoreSyncErrorKind::RequestTooLarge, 429 => LocalCoreSyncErrorKind::RateLimit, 503 => LocalCoreSyncErrorKind::Overloaded, _ => LocalCoreSyncErrorKind::ServerError, } } fn build_openai_image_sync_heartbeat_shell_response( state: AppState, request_path: String, trace_id: String, decision: GatewayControlDecision, plan_kind: String, attempts: Vec, transfer_tracker: ProviderTransferTracker, ) -> Result, GatewayError> { let request_id = attempts .first() .map(|attempt| attempt.plan.request_id.clone()) .filter(|value| !value.trim().is_empty()); let trace_id_for_response = trace_id.clone(); let decision_for_response = decision.clone(); let started_at = Instant::now(); let (tx, rx) = mpsc::channel::>(1); let request_diagnostics = current_request_diagnostics(); let cancel_on_disconnect = crate::request_lifecycle::cancel_on_client_disconnect(); tokio::spawn(async move { scope_request_diagnostics_with(request_diagnostics, async move { let execution = execute_openai_image_sync_heartbeat_attempts( state, request_path, trace_id, decision, plan_kind, attempts, transfer_tracker, started_at, ); let outcome = tokio::select! { biased; _ = tx.closed(), if cancel_on_disconnect => return, result = execution => result, }; let bytes = openai_image_sync_heartbeat_final_bytes(outcome).await; let _ = tx.send(Ok(Bytes::from(bytes))).await; }) .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_openai_image_sync_json_whitespace_heartbeat_stream(rx)), 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, request_id.as_deref(), None) } async fn execute_openai_image_sync_heartbeat_attempts( state: AppState, request_path: String, trace_id: String, decision: GatewayControlDecision, plan_kind: String, attempts: Vec, transfer_tracker: ProviderTransferTracker, started_at: Instant, ) -> Result { let (parts, _) = http::Request::builder() .uri(request_path.as_str()) .body(()) .map_err(|err| GatewayError::Internal(err.to_string()))? .into_parts(); match execute_sync_plan_and_reports_with_transfer_tracker( &state, &parts, trace_id.as_str(), &decision, plan_kind.as_str(), attempts, &transfer_tracker, ) .await? { LocalExecutionRequestOutcome::Exhausted(exhaustion) => { record_failed_usage_for_exhausted_request( &state, exhaustion, &started_at, "OpenAI image sync heartbeat exhausted all local candidates", EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, None, ) .await; Ok(LocalExecutionRequestOutcome::NoPath) } outcome => Ok(outcome), } } async fn openai_image_sync_heartbeat_final_bytes( result: Result, ) -> Vec { match result { Ok(LocalExecutionRequestOutcome::Responded(response)) => { openai_image_sync_heartbeat_response_body_bytes(response).await } Ok(LocalExecutionRequestOutcome::Exhausted(_)) | Ok(LocalExecutionRequestOutcome::NoPath) => openai_image_sync_heartbeat_error_body( OPENAI_IMAGE_SYNC_HEARTBEAT_EXHAUSTED_STATUS, "OpenAI image sync exhausted all local candidates", ), Err(_err) => openai_image_sync_heartbeat_error_body( OPENAI_IMAGE_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, "internal gateway error while executing image request", ), } } async fn openai_image_sync_heartbeat_response_body_bytes(response: Response) -> Vec { let status_code = response.status().as_u16(); match to_bytes( response.into_body(), crate::headers::max_internal_buffered_body_bytes(), ) .await { Ok(bytes) if status_code < 400 && !bytes.is_empty() => bytes.to_vec(), Ok(bytes) if status_code >= 400 => { openai_image_sync_heartbeat_error_body_from_response(status_code, bytes.as_ref()) } Ok(_) => openai_image_sync_heartbeat_error_body( OPENAI_IMAGE_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, "empty sync image response", ), Err(_err) => openai_image_sync_heartbeat_error_body( OPENAI_IMAGE_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS, "internal gateway error while reading image response", ), } } fn openai_image_sync_heartbeat_error_body_from_response(status_code: u16, body: &[u8]) -> Vec { if let Ok(mut value) = serde_json::from_slice::(body) { if let Some(error) = value.get_mut("error").and_then(Value::as_object_mut) { error.insert("upstream_status".to_string(), Value::from(status_code)); error .entry("type".to_string()) .or_insert_with(|| Value::String("upstream_error".to_string())); error.entry("message".to_string()).or_insert_with(|| { Value::String(format!("upstream returned status {status_code}")) }); return serde_json::to_vec(&value).unwrap_or_else(|_| { openai_image_sync_heartbeat_error_body( status_code, &format!("upstream returned status {status_code}"), ) }); } } let message = openai_image_sync_heartbeat_error_message_from_body(status_code, body); openai_image_sync_heartbeat_error_body(status_code, message.as_str()) } fn openai_image_sync_heartbeat_error_message_from_body(status_code: u16, body: &[u8]) -> String { let text = String::from_utf8_lossy(body).trim().to_string(); if text.is_empty() { return format!("upstream returned status {status_code}"); } text.chars() .take(OPENAI_IMAGE_SYNC_HEARTBEAT_ERROR_MESSAGE_LIMIT) .collect() } fn openai_image_sync_heartbeat_error_body(status_code: u16, message: &str) -> Vec { serde_json::to_vec(&json!({ "error": { "type": "upstream_error", "message": message, "code": status_code, "upstream_status": status_code, } })) .unwrap_or_else(|_| { format!( "{{\"error\":{{\"type\":\"upstream_error\",\"code\":{status_code},\"upstream_status\":{status_code}}}}}" ) .into_bytes() }) } pub(crate) async fn maybe_execute_sync_via_local_image_decision( state: &AppState, parts: &http::request::Parts, body_json: &serde_json::Value, body_base64: Option<&str>, trace_id: &str, decision: &GatewayControlDecision, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some((mut attempt_source, candidate_count)) = build_local_image_sync_attempt_source_for_kind( state, parts, body_json, body_base64, trace_id, decision, plan_kind, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; if attempt_source .routing_execution_policy() .is_some_and(|policy| policy.enable_cf_heartbeat) { let mut attempts = Vec::new(); while let Some(attempt) = attempt_source.next_execution_attempt().await? { attempts.push(attempt); } return Ok(LocalExecutionRequestOutcome::responded( build_openai_image_sync_heartbeat_shell_response( state.clone(), parts.uri.path().to_string(), trace_id.to_string(), decision.clone(), plan_kind.to_string(), attempts, transfer_tracker.clone(), )?, )); } let outcome = execute_sync_attempt_source_with_transfer_tracker::( state, parts, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await?; if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome { set_local_openai_image_execution_exhausted_diagnostic( state, trace_id, decision, plan_kind, body_json, candidate_count, ); } Ok(outcome) } pub(crate) async fn maybe_execute_stream_via_local_gemini_files_decision( state: &AppState, parts: &http::request::Parts, trace_id: &str, decision: &GatewayControlDecision, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some((attempt_source, _candidate_count)) = build_local_gemini_files_stream_attempt_source_for_kind( state, parts, trace_id, decision, plan_kind, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; execute_stream_attempt_source_with_transfer_tracker::( state, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await } pub(crate) async fn maybe_execute_stream_via_local_image_decision( state: &AppState, parts: &http::request::Parts, body_json: &serde_json::Value, body_base64: Option<&str>, trace_id: &str, decision: &GatewayControlDecision, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some((attempt_source, candidate_count)) = build_local_image_stream_attempt_source_for_kind( state, parts, body_json, body_base64, trace_id, decision, plan_kind, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; let outcome = execute_stream_attempt_source_with_transfer_tracker::( state, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await?; if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome { set_local_openai_image_execution_exhausted_diagnostic( state, trace_id, decision, plan_kind, body_json, candidate_count, ); } Ok(outcome) } pub(crate) async fn maybe_execute_sync_via_local_video_decision( state: &AppState, parts: &http::request::Parts, body_json: &serde_json::Value, trace_id: &str, decision: &GatewayControlDecision, plan_kind: &str, transfer_tracker: &ProviderTransferTracker, ) -> Result { let Some((attempt_source, _candidate_count)) = build_local_video_sync_attempt_source_for_kind( state, parts, body_json, trace_id, decision, plan_kind, ) .await? else { return Ok(LocalExecutionRequestOutcome::NoPath); }; execute_sync_attempt_source_with_transfer_tracker::( state, parts, trace_id, decision, plan_kind, attempt_source, transfer_tracker, ) .await } fn supports_local_video_get( parts: &http::request::Parts, decision: &GatewayControlDecision, ) -> bool { parts.method == http::Method::GET && decision.route_kind.as_deref() == Some("video") && (crate::video_tasks::resolve_video_task_read_lookup_key( decision.route_family.as_deref(), parts.uri.path(), ) .is_some() || (decision.route_family.as_deref() == Some("openai") && crate::video_tasks::extract_openai_task_id_from_content_path(parts.uri.path()) .is_some())) } pub(crate) fn maybe_execute_sync_request<'a>( state: &'a AppState, parts: &'a http::request::Parts, body_bytes: &'a axum::body::Bytes, trace_id: &'a str, decision: Option<&'a GatewayControlDecision>, ) -> Pin> + Send + 'a>> { Box::pin(async move { let Some(decision) = decision else { return Ok(LocalExecutionRequestOutcome::NoPath); }; #[cfg(not(test))] { if parts.method != http::Method::POST && !supports_local_video_get(parts, decision) { return Ok(LocalExecutionRequestOutcome::NoPath); } return maybe_execute_sync_local_path(state, parts, body_bytes, trace_id, decision) .await; } #[cfg(test)] { if state .execution_runtime_override_base_url() .unwrap_or_default() .is_empty() && parts.method != http::Method::POST && !supports_local_video_get(parts, decision) { return Ok(LocalExecutionRequestOutcome::NoPath); } maybe_execute_sync_local_path(state, parts, body_bytes, trace_id, decision).await } }) } pub(crate) fn maybe_execute_stream_request<'a>( state: &'a AppState, parts: &'a http::request::Parts, body_bytes: &'a axum::body::Bytes, trace_id: &'a str, decision: Option<&'a GatewayControlDecision>, ) -> Pin> + Send + 'a>> { Box::pin(async move { let Some(decision) = decision else { return Ok(LocalExecutionRequestOutcome::NoPath); }; #[cfg(not(test))] { if parts.method != http::Method::POST && !supports_local_video_get(parts, decision) { return Ok(LocalExecutionRequestOutcome::NoPath); } return maybe_execute_stream_local_path(state, parts, body_bytes, trace_id, decision) .await; } #[cfg(test)] { if state .execution_runtime_override_base_url() .unwrap_or_default() .is_empty() && parts.method != http::Method::POST && !supports_local_video_get(parts, decision) { return Ok(LocalExecutionRequestOutcome::NoPath); } maybe_execute_stream_local_path(state, parts, body_bytes, trace_id, decision).await } }) } pub(crate) fn planner_decision_action(action: &str) -> bool { matches!( action, EXECUTION_RUNTIME_SYNC_DECISION_ACTION | EXECUTION_RUNTIME_STREAM_DECISION_ACTION ) } pub(crate) fn parse_local_request_body( parts: &http::request::Parts, body_bytes: &axum::body::Bytes, ) -> Option<(serde_json::Value, Option)> { parse_direct_request_body(parts, body_bytes) } pub(crate) fn decision_payload_is_direct_execution(payload: &AiExecutionDecision) -> bool { planner_decision_action(payload.action.as_str()) } #[cfg(test)] mod tests { use super::*; use aether_data::repository::candidates::InMemoryRequestCandidateRepository; use aether_data::repository::usage::InMemoryUsageReadRepository; use aether_data_contracts::repository::usage::UsageReadRepository; use aether_usage_runtime::UsageRuntimeConfig; use futures_util::StreamExt; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::Duration; const TEST_OPENAI_IMAGE_SYNC_PLAN_KIND: &str = "openai_image_sync"; const TEST_STANDARD_TEXT_SYNC_PLAN_KIND: &str = "openai_responses_compact_sync"; const HEARTBEAT_USAGE_POLL_INTERVAL: Duration = Duration::from_millis(10); const HEARTBEAT_USAGE_SETTLE_TIMEOUT: Duration = Duration::from_secs(30); struct TestSyncAttemptSource { attempts: VecDeque, } impl TestSyncAttemptSource { fn new(attempts: Vec) -> Self { Self { attempts: VecDeque::from(attempts), } } } #[async_trait::async_trait] impl LocalExecutionAttemptSource for TestSyncAttemptSource { async fn next_execution_attempt(&mut self) -> Result, GatewayError> { Ok(self.attempts.pop_front()) } async fn drain_execution_attempts(&mut self) -> Result, GatewayError> { Ok(self.attempts.drain(..).collect()) } async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { self.attempts .retain(|attempt| attempt.plan.key_id != key_id); Ok(()) } async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { self.attempts .retain(|attempt| attempt.plan.endpoint_id != endpoint_id); Ok(()) } async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.attempts .retain(|attempt| attempt.plan.provider_id != provider_id); Ok(()) } } fn test_openai_image_heartbeat_decision() -> GatewayControlDecision { GatewayControlDecision::synthetic( "/v1/images/generations", Some("ai_public".to_string()), Some("openai".to_string()), Some("image".to_string()), Some("openai:image".to_string()), ) .with_execution_runtime_candidate(true) } fn test_openai_image_heartbeat_plan( endpoint_id: &str, candidate_id: &str, ) -> aether_contracts::ExecutionPlan { aether_contracts::ExecutionPlan { request_id: "trace-image-heartbeat-retry".to_string(), candidate_id: Some(candidate_id.to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-openai".to_string(), endpoint_id: endpoint_id.to_string(), key_id: "key-openai".to_string(), method: "POST".to_string(), url: "https://api.openai.com/v1/images/generations".to_string(), headers: BTreeMap::new(), content_type: Some("application/json".to_string()), content_encoding: None, body: aether_contracts::RequestBody::from_json(json!({"prompt": "test"})), stream: false, client_api_format: "openai:image".to_string(), provider_api_format: "openai:image".to_string(), model_name: Some("gpt-image-1".to_string()), proxy: None, transport_profile: None, timeouts: None, } } fn test_openai_image_heartbeat_attempt( candidate_index: u32, endpoint_id: &str, candidate_id: &str, ) -> AiSyncAttempt { test_openai_image_heartbeat_attempt_with_sticky_key_attempts( candidate_index, endpoint_id, candidate_id, 1, ) } /// `sticky_key_attempts` is pinned so these tests exercise candidate /// failover; the default same-key retry is covered separately. fn test_openai_image_heartbeat_attempt_with_sticky_key_attempts( candidate_index: u32, endpoint_id: &str, candidate_id: &str, sticky_key_attempts: u32, ) -> AiSyncAttempt { AiSyncAttempt { plan: test_openai_image_heartbeat_plan(endpoint_id, candidate_id), report_kind: None, report_context: Some(json!({ "candidate_index": candidate_index, "retry_index": 0, "sticky_key_attempts": sticky_key_attempts, })), } } fn test_openai_image_execution_result( plan: &aether_contracts::ExecutionPlan, status_code: u16, body_json: Value, ) -> aether_contracts::ExecutionResult { aether_contracts::ExecutionResult { request_id: plan.request_id.clone(), candidate_id: plan.candidate_id.clone(), status_code, headers: BTreeMap::from([( CONTENT_TYPE.as_str().to_string(), "application/json".to_string(), )]), response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(body_json), body_bytes_b64: None, }), telemetry: Some(aether_contracts::ExecutionTelemetry { ttfb_ms: None, elapsed_ms: Some(10), upstream_bytes: None, }), error: None, } } fn heartbeat_usage_test_state( response_body: Value, ) -> (AppState, Arc) { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let state = AppState::new() .expect("state should build") .with_data_state_for_tests( crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests( request_candidate_repository, Arc::clone(&usage_repository), ), ) .with_usage_runtime_for_tests(UsageRuntimeConfig { enabled: true, ..UsageRuntimeConfig::default() }) .with_execution_runtime_sync_override_for_tests(move |plan| { let mut result = test_openai_image_execution_result( plan, StatusCode::OK.as_u16(), response_body.clone(), ); if let Some(telemetry) = result.telemetry.as_mut() { telemetry.ttfb_ms = Some(5); } Ok(result) }); (state, usage_repository) } async fn assert_usage_has_end_to_end_timings( usage_repository: &InMemoryUsageReadRepository, request_id: &str, ) { let deadline = Instant::now() + HEARTBEAT_USAGE_SETTLE_TIMEOUT; let usage = loop { let usage = usage_repository .find_by_request_id(request_id) .await .expect("usage should read"); if usage.as_ref().is_some_and(|usage| { matches!(usage.status.as_str(), "completed" | "failed" | "cancelled") }) { break usage.expect("terminal usage should be recorded"); } let now = Instant::now(); let last_status = usage.as_ref().map(|usage| usage.status.as_str()); assert!( now < deadline, "terminal usage should be recorded within {HEARTBEAT_USAGE_SETTLE_TIMEOUT:?}; \ last status: {}", last_status.unwrap_or("") ); tokio::time::sleep(HEARTBEAT_USAGE_POLL_INTERVAL.min(deadline - now)).await; }; assert_eq!( usage.status, "completed", "heartbeat usage should complete successfully" ); let request_metadata = usage .request_metadata .as_ref() .expect("terminal usage should retain request diagnostics"); let end_to_end_time_ms = request_metadata .get("end_to_end_time_ms") .and_then(Value::as_u64) .expect("end-to-end time should be recorded"); let end_to_end_first_byte_time_ms = request_metadata .get("end_to_end_first_byte_time_ms") .and_then(Value::as_u64) .expect("end-to-end first-byte time should be recorded"); assert!(end_to_end_first_byte_time_ms <= end_to_end_time_ms); } fn test_standard_text_heartbeat_decision() -> GatewayControlDecision { GatewayControlDecision::synthetic( "/v1/responses", Some("ai_public".to_string()), Some("openai".to_string()), Some("responses".to_string()), Some("openai:responses:compact".to_string()), ) .with_execution_runtime_candidate(true) } fn test_standard_text_heartbeat_plan( endpoint_id: &str, candidate_id: &str, client_api_format: &str, ) -> aether_contracts::ExecutionPlan { aether_contracts::ExecutionPlan { request_id: "trace-standard-text-heartbeat-retry".to_string(), candidate_id: Some(candidate_id.to_string()), provider_name: Some("OpenAI".to_string()), provider_id: "provider-openai".to_string(), endpoint_id: endpoint_id.to_string(), key_id: "key-openai".to_string(), method: "POST".to_string(), url: "https://api.openai.com/v1/responses".to_string(), headers: BTreeMap::new(), content_type: Some("application/json".to_string()), content_encoding: None, body: aether_contracts::RequestBody::from_json(json!({"model": "gpt-5"})), stream: false, client_api_format: client_api_format.to_string(), provider_api_format: client_api_format.to_string(), model_name: Some("gpt-5".to_string()), proxy: None, transport_profile: None, timeouts: None, } } fn test_standard_text_heartbeat_attempt( candidate_index: u32, endpoint_id: &str, candidate_id: &str, client_api_format: &str, ) -> AiSyncAttempt { AiSyncAttempt { plan: test_standard_text_heartbeat_plan(endpoint_id, candidate_id, client_api_format), report_kind: None, report_context: Some(json!({ "candidate_index": candidate_index, "retry_index": 0, // Pin to a single attempt so this helper exercises candidate // failover rather than the default same-key retry. "sticky_key_attempts": 1, "client_api_format": client_api_format, "provider_api_format": client_api_format, })), } } #[tokio::test] async fn openai_image_sync_heartbeat_success_body_is_unchanged() { let response = Response::builder() .status(StatusCode::OK) .body(Body::from(r#"{"data":[{"b64_json":"x"}]}"#)) .expect("response should build"); let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await; let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(body, json!({"data": [{"b64_json": "x"}]})); } #[test] fn openai_image_sync_heartbeat_missing_routing_policy_defaults_disabled() { assert!(!Option::::None .is_some_and(|policy| policy.enable_cf_heartbeat)); } #[tokio::test] async fn openai_image_sync_heartbeat_error_body_includes_upstream_status() { let response = Response::builder() .status(StatusCode::TOO_MANY_REQUESTS) .body(Body::from( r#"{"error":{"type":"rate_limit","message":"slow down"}}"#, )) .expect("response should build"); let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await; let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(body["error"]["type"], json!("rate_limit")); assert_eq!(body["error"]["message"], json!("slow down")); assert_eq!(body["error"]["upstream_status"], json!(429)); } #[test] fn openai_image_sync_heartbeat_non_json_error_body_is_wrapped() { let bytes = openai_image_sync_heartbeat_error_body_from_response(502, b"bad gateway from upstream"); let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(body["error"]["type"], json!("upstream_error")); assert_eq!(body["error"]["message"], json!("bad gateway from upstream")); assert_eq!(body["error"]["upstream_status"], json!(502)); } #[tokio::test] async fn openai_image_sync_heartbeat_no_path_returns_json_error_body() { let bytes = openai_image_sync_heartbeat_final_bytes(Ok(LocalExecutionRequestOutcome::NoPath)).await; let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(body["error"]["type"], json!("upstream_error")); assert_eq!(body["error"]["upstream_status"], json!(503)); } #[tokio::test] async fn openai_image_sync_heartbeat_propagates_request_diagnostics_to_terminal_usage() { let (state, usage_repository) = heartbeat_usage_test_state(json!({ "data": [{"b64_json": "heartbeat-image"}] })); let response = crate::request_diagnostics::scope_request_diagnostics(async move { crate::request_diagnostics::record_request_accepted_at( Instant::now() - Duration::from_millis(25), ); build_openai_image_sync_heartbeat_shell_response( state, "/v1/images/generations".to_string(), "trace-image-heartbeat-retry".to_string(), test_openai_image_heartbeat_decision(), TEST_OPENAI_IMAGE_SYNC_PLAN_KIND.to_string(), vec![test_openai_image_heartbeat_attempt( 0, "endpoint-success", "candidate-success", )], ProviderTransferTracker::default(), ) }) .await .expect("heartbeat shell should build"); let body = to_bytes( response.into_body(), crate::headers::max_internal_buffered_body_bytes(), ) .await .expect("heartbeat response body should complete"); assert!(!body.is_empty()); assert_usage_has_end_to_end_timings( usage_repository.as_ref(), "trace-image-heartbeat-retry", ) .await; } #[tokio::test] async fn openai_image_sync_heartbeat_attempts_retry_first_candidate_then_return_second() { let call_count = Arc::new(AtomicUsize::new(0)); let call_count_for_override = Arc::clone(&call_count); let state = AppState::new() .expect("state should build") .with_execution_runtime_sync_override_for_tests(move |plan| { call_count_for_override.fetch_add(1, Ordering::SeqCst); if plan.endpoint_id == "endpoint-retry" { Ok(test_openai_image_execution_result( plan, StatusCode::TOO_MANY_REQUESTS.as_u16(), json!({"error": {"message": "retry this candidate"}}), )) } else { Ok(test_openai_image_execution_result( plan, StatusCode::OK.as_u16(), json!({"data": [{"b64_json": "second-candidate"}]}), )) } }); let attempts = vec![ test_openai_image_heartbeat_attempt(0, "endpoint-retry", "candidate-retry"), test_openai_image_heartbeat_attempt(1, "endpoint-success", "candidate-success"), ]; let outcome = execute_openai_image_sync_heartbeat_attempts( state, "/v1/images/generations".to_string(), "trace-image-heartbeat-retry".to_string(), test_openai_image_heartbeat_decision(), TEST_OPENAI_IMAGE_SYNC_PLAN_KIND.to_string(), attempts, ProviderTransferTracker::default(), Instant::now(), ) .await .expect("heartbeat attempts should execute"); let LocalExecutionRequestOutcome::Responded(response) = outcome else { panic!("second candidate should return a response"); }; let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await; let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(call_count.load(Ordering::SeqCst), 2); assert_eq!(body, json!({"data": [{"b64_json": "second-candidate"}]})); } #[tokio::test] async fn openai_image_sync_heartbeat_retries_sticky_key_lazily_before_failover() { let seen_plans = Arc::new(std::sync::Mutex::new(Vec::<(String, Option)>::new())); let seen_plans_for_override = Arc::clone(&seen_plans); let state = AppState::new() .expect("state should build") .with_execution_runtime_sync_override_for_tests(move |plan| { seen_plans_for_override .lock() .expect("mutex should lock") .push((plan.endpoint_id.clone(), plan.candidate_id.clone())); if plan.endpoint_id == "endpoint-retry" { Ok(test_openai_image_execution_result( plan, StatusCode::TOO_MANY_REQUESTS.as_u16(), json!({"error": {"message": "retry this candidate"}}), )) } else { Ok(test_openai_image_execution_result( plan, StatusCode::OK.as_u16(), json!({"data": [{"b64_json": "second-candidate"}]}), )) } }); // Three total attempts on the sticky key; only one attempt is // materialized up front, the other two are derived after each failure. let attempts = vec![ test_openai_image_heartbeat_attempt_with_sticky_key_attempts( 0, "endpoint-retry", "candidate-retry", 3, ), test_openai_image_heartbeat_attempt_with_sticky_key_attempts( 1, "endpoint-success", "candidate-success", 3, ), ]; let outcome = execute_openai_image_sync_heartbeat_attempts( state, "/v1/images/generations".to_string(), "trace-image-heartbeat-sticky-retry".to_string(), test_openai_image_heartbeat_decision(), TEST_OPENAI_IMAGE_SYNC_PLAN_KIND.to_string(), attempts, ProviderTransferTracker::default(), Instant::now(), ) .await .expect("heartbeat attempts should execute"); let LocalExecutionRequestOutcome::Responded(response) = outcome else { panic!("second candidate should return a response"); }; let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await; let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); let seen_plans = seen_plans.lock().expect("mutex should lock").clone(); assert_eq!( seen_plans .iter() .map(|(endpoint_id, _)| endpoint_id.as_str()) .collect::>(), [ "endpoint-retry", "endpoint-retry", "endpoint-retry", "endpoint-success" ] ); let sticky_candidate_ids = seen_plans[..3] .iter() .map(|(_, candidate_id)| candidate_id.clone()) .collect::>(); assert_eq!( sticky_candidate_ids.len(), 3, "each derived same-key retry must carry a fresh candidate id" ); assert_eq!(body, json!({"data": [{"b64_json": "second-candidate"}]})); } #[tokio::test] async fn openai_image_sync_heartbeat_honors_provider_transfer_limit() { let call_count = Arc::new(AtomicUsize::new(0)); let call_count_for_override = Arc::clone(&call_count); let state = AppState::new() .expect("state should build") .with_execution_runtime_sync_override_for_tests(move |plan| { call_count_for_override.fetch_add(1, Ordering::SeqCst); if plan.provider_id == "provider-fallback" { Ok(test_openai_image_execution_result( plan, StatusCode::OK.as_u16(), json!({"data": [{"b64_json": "fallback-provider"}]}), )) } else { Ok(test_openai_image_execution_result( plan, StatusCode::TOO_MANY_REQUESTS.as_u16(), json!({"error": {"message": "retry another key"}}), )) } }); let mut attempts = vec![ test_openai_image_heartbeat_attempt(0, "endpoint-key-1", "candidate-key-1"), test_openai_image_heartbeat_attempt(1, "endpoint-key-2", "candidate-key-2"), test_openai_image_heartbeat_attempt(2, "endpoint-key-3", "candidate-key-3"), test_openai_image_heartbeat_attempt(3, "endpoint-fallback", "candidate-fallback"), ]; for (index, attempt) in attempts.iter_mut().take(3).enumerate() { attempt.plan.key_id = format!("key-{}", index + 1); attempt.report_context = Some(json!({ "candidate_index": index, "retry_index": 0, "sticky_key_attempts": 1, "local_failover_policy": { "max_transfer_count": 1, "max_transfer_timeout_seconds": 0 } })); } attempts[3].plan.provider_id = "provider-fallback".to_string(); attempts[3].plan.key_id = "key-fallback".to_string(); let outcome = execute_openai_image_sync_heartbeat_attempts( state, "/v1/images/generations".to_string(), "trace-image-heartbeat-transfer-limit".to_string(), test_openai_image_heartbeat_decision(), TEST_OPENAI_IMAGE_SYNC_PLAN_KIND.to_string(), attempts, ProviderTransferTracker::default(), Instant::now(), ) .await .expect("heartbeat attempts should execute"); let LocalExecutionRequestOutcome::Responded(response) = outcome else { panic!("fallback provider should return a response"); }; let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await; let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(call_count.load(Ordering::SeqCst), 3); assert_eq!(body, json!({"data": [{"b64_json": "fallback-provider"}]})); } #[tokio::test] async fn standard_text_sync_heartbeat_no_local_candidates_preserves_no_path() { let state = AppState::new().expect("state should build"); let (parts, _) = http::Request::builder() .method(http::Method::POST) .uri("/v1/responses") .body(()) .expect("request should build") .into_parts(); let outcome = maybe_execute_sync_via_local_openai_responses_decision( &state, &parts, "trace-standard-text-heartbeat-no-path", &test_standard_text_heartbeat_decision(), &json!({"model": "missing-local-candidate"}), TEST_STANDARD_TEXT_SYNC_PLAN_KIND, &ProviderTransferTracker::default(), ) .await .expect("heartbeat no-path check should execute"); assert!(matches!(outcome, LocalExecutionRequestOutcome::NoPath)); } #[tokio::test] async fn standard_text_sync_heartbeat_success_body_is_unchanged() { let response = Response::builder() .status(StatusCode::OK) .body(Body::from(r#"{"id":"resp_123","output":[]}"#)) .expect("response should build"); let bytes = standard_text_sync_heartbeat_response_body_bytes( "openai:responses:compact", None, response, ) .await; let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(body, json!({"id": "resp_123", "output": []})); } #[tokio::test] async fn standard_text_sync_heartbeat_claude_error_body_includes_upstream_status() { let response = Response::builder() .status(StatusCode::TOO_MANY_REQUESTS) .body(Body::from( r#"{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}"#, )) .expect("response should build"); let bytes = standard_text_sync_heartbeat_response_body_bytes("claude:messages", None, response) .await; let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(body["type"], json!("error")); assert_eq!(body["error"]["type"], json!("rate_limit_error")); assert_eq!(body["error"]["message"], json!("slow down")); assert_eq!(body["error"]["upstream_status"], json!(429)); } #[test] fn standard_text_sync_heartbeat_applies_to_chat_and_cli_plan_kinds() { assert!(standard_text_sync_heartbeat_applies_to_plan_kind( "claude_chat_sync" )); assert!(standard_text_sync_heartbeat_applies_to_plan_kind( "claude_cli_sync" )); assert!(standard_text_sync_heartbeat_applies_to_plan_kind( "gemini_chat_sync" )); assert!(standard_text_sync_heartbeat_applies_to_plan_kind( "gemini_cli_sync" )); assert!(!standard_text_sync_heartbeat_applies_to_plan_kind( "openai_embedding_sync" )); } #[tokio::test] async fn standard_text_sync_heartbeat_redirect_status_is_wrapped_as_error() { let response = Response::builder() .status(StatusCode::TEMPORARY_REDIRECT) .body(Body::from(r#"{"location":"https://upstream.example"}"#)) .expect("response should build"); let bytes = standard_text_sync_heartbeat_response_body_bytes( "openai:responses:compact", None, response, ) .await; let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(body["error"]["type"], json!("server_error")); assert_eq!(body["error"]["upstream_status"], json!(307)); } #[tokio::test] async fn standard_text_sync_heartbeat_shell_sends_whitespace_before_background_finishes() { let state = AppState::new().expect("state should build"); let (parts, _) = http::Request::builder() .method(http::Method::POST) .uri("/v1/responses") .body(()) .expect("request should build") .into_parts(); let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); let response = build_standard_text_sync_heartbeat_shell_response( state, parts, "trace-standard-text-heartbeat-shell".to_string(), test_standard_text_heartbeat_decision(), TEST_STANDARD_TEXT_SYNC_PLAN_KIND.to_string(), move |_state, _parts, _trace_id, _decision, _plan_kind, _started_at| async move { let _ = release_rx.await; Ok(LocalExecutionRequestOutcome::responded( Response::builder() .status(StatusCode::OK) .body(Body::from(r#"{"id":"resp_done","output":[]}"#)) .expect("response should build"), )) }, ) .expect("heartbeat shell should build"); let mut body_stream = response.into_body().into_data_stream(); let first = body_stream .next() .await .expect("heartbeat stream should yield") .expect("heartbeat chunk should be ok"); assert_eq!(first.as_ref(), b"\n"); let _ = release_tx.send(()); } #[tokio::test] async fn standard_text_sync_heartbeat_background_holds_request_admission_after_disconnect() { let state = AppState::new().expect("state should build"); let gate = aether_runtime::ConcurrencyGate::new("heartbeat_request", 1); let admission = aether_runtime::AdmissionPermit::from( gate.try_acquire().expect("request admission permit"), ); let (mut parts, _) = http::Request::builder() .method(http::Method::POST) .uri("/v1/responses") .body(()) .expect("request should build") .into_parts(); parts.extensions.insert( crate::executor::candidate_loop::BackgroundAdmissionPermit::new(admission.clone()), ); drop(admission); let (started_tx, started_rx) = tokio::sync::oneshot::channel::<()>(); let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); let response = build_standard_text_sync_heartbeat_shell_response( state, parts, "trace-standard-text-heartbeat-admission".to_string(), test_standard_text_heartbeat_decision(), TEST_STANDARD_TEXT_SYNC_PLAN_KIND.to_string(), move |_state, parts, _trace_id, _decision, _plan_kind, _started_at| async move { assert!( parts .extensions .get::() .is_some(), "background request parts should carry admission" ); let _ = started_tx.send(()); let _ = release_rx.await; Ok(LocalExecutionRequestOutcome::responded( Response::builder() .status(StatusCode::OK) .body(Body::from(r#"{"id":"resp_done","output":[]}"#)) .expect("response should build"), )) }, ) .expect("heartbeat shell should build"); started_rx.await.expect("background execution should start"); drop(response); assert_eq!(gate.snapshot().in_flight, 1); assert!( gate.try_acquire().is_err(), "disconnect must not release background admission" ); let _ = release_tx.send(()); tokio::time::timeout(Duration::from_secs(1), async { while gate.snapshot().in_flight != 0 { tokio::task::yield_now().await; } }) .await .expect("background completion should release admission"); } #[tokio::test] async fn standard_text_sync_heartbeat_cancels_when_routing_policy_enables_it() { let (started_tx, started_rx) = tokio::sync::oneshot::channel(); let (mut release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); let response = crate::request_lifecycle::run_request(async move { crate::request_lifecycle::configure_client_disconnect( aether_routing_core::RoutingExecutionPolicy { cancel_on_client_disconnect: true, ..Default::default() }, ); let (parts, _) = http::Request::builder() .method("POST") .uri("/v1/responses") .body(()) .unwrap() .into_parts(); build_standard_text_sync_heartbeat_shell_response( AppState::new().unwrap(), parts, "trace-heartbeat-disconnect".to_string(), test_standard_text_heartbeat_decision(), TEST_STANDARD_TEXT_SYNC_PLAN_KIND.to_string(), move |_, _, _, _, _, _| async move { started_tx.send(()).unwrap(); release_rx.await.unwrap(); Ok(LocalExecutionRequestOutcome::NoPath) }, ) }) .await .unwrap(); started_rx.await.unwrap(); drop(response); tokio::time::timeout(Duration::from_secs(1), release_tx.closed()) .await .expect("heartbeat must drop upstream execution immediately"); } #[tokio::test] async fn standard_text_sync_heartbeat_propagates_request_diagnostics_to_terminal_usage() { let (state, usage_repository) = heartbeat_usage_test_state(json!({ "id": "resp_heartbeat", "output": [] })); let (parts, _) = http::Request::builder() .method(http::Method::POST) .uri("/v1/responses") .body(()) .expect("request should build") .into_parts(); let response = crate::request_diagnostics::scope_request_diagnostics(async move { crate::request_diagnostics::record_request_accepted_at( Instant::now() - Duration::from_millis(25), ); build_standard_text_sync_heartbeat_shell_response( state, parts, "trace-standard-text-heartbeat-retry".to_string(), test_standard_text_heartbeat_decision(), TEST_STANDARD_TEXT_SYNC_PLAN_KIND.to_string(), move |state, parts, trace_id, decision, plan_kind, _started_at| async move { execute_sync_attempt_source::( &state, &parts, trace_id.as_str(), &decision, plan_kind.as_str(), TestSyncAttemptSource::new(vec![test_standard_text_heartbeat_attempt( 0, "endpoint-success", "candidate-success", "openai:responses:compact", )]), ) .await }, ) }) .await .expect("heartbeat shell should build"); let body = to_bytes( response.into_body(), crate::headers::max_internal_buffered_body_bytes(), ) .await .expect("heartbeat response body should complete"); assert!(!body.is_empty()); assert_usage_has_end_to_end_timings( usage_repository.as_ref(), "trace-standard-text-heartbeat-retry", ) .await; } #[test] fn standard_text_sync_heartbeat_compact_non_json_error_body_is_wrapped_in_client_format() { let bytes = standard_text_sync_heartbeat_error_body_from_response( "openai:responses:compact", 502, b"bad gateway from upstream", ); let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(body["error"]["type"], json!("server_error")); assert_eq!(body["error"]["message"], json!("bad gateway from upstream")); assert_eq!(body["error"]["upstream_status"], json!(502)); } #[tokio::test] async fn standard_text_sync_heartbeat_attempts_retry_first_candidate_then_return_second() { let call_count = Arc::new(AtomicUsize::new(0)); let call_count_for_override = Arc::clone(&call_count); let state = AppState::new() .expect("state should build") .with_execution_runtime_sync_override_for_tests(move |plan| { call_count_for_override.fetch_add(1, Ordering::SeqCst); if plan.endpoint_id == "endpoint-retry" { Ok(test_openai_image_execution_result( plan, StatusCode::TOO_MANY_REQUESTS.as_u16(), json!({"error": {"message": "retry this candidate"}}), )) } else { Ok(test_openai_image_execution_result( plan, StatusCode::OK.as_u16(), json!({"id": "resp_second_candidate", "output": []}), )) } }); let attempts = vec![ test_standard_text_heartbeat_attempt( 0, "endpoint-retry", "candidate-retry", "openai:responses:compact", ), test_standard_text_heartbeat_attempt( 1, "endpoint-success", "candidate-success", "openai:responses:compact", ), ]; let (parts, _) = http::Request::builder() .method(http::Method::POST) .uri("/v1/responses") .body(()) .expect("request should build") .into_parts(); let outcome = execute_sync_attempt_source::( &state, &parts, "trace-standard-text-heartbeat-retry", &test_standard_text_heartbeat_decision(), TEST_STANDARD_TEXT_SYNC_PLAN_KIND, TestSyncAttemptSource::new(attempts), ) .await .expect("heartbeat attempts should execute"); let LocalExecutionRequestOutcome::Responded(response) = outcome else { panic!("second candidate should return a response"); }; let bytes = standard_text_sync_heartbeat_response_body_bytes( "openai:responses:compact", None, response, ) .await; let body: Value = serde_json::from_slice(&bytes).expect("body should decode"); assert_eq!(call_count.load(Ordering::SeqCst), 2); assert_eq!(body, json!({"id": "resp_second_candidate", "output": []})); } }