i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-06 01:17:46 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
e83399db2fd1d16f07d5a6e8c563c9a199c190c1
Aether/apps/aether-gateway/src/executor/orchestration.rs
T

2519 lines
91 KiB
Rust
Raw Blame History

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<LocalExecutionRequestOutcome, GatewayError> {
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<LocalExecutionRequestOutcome, GatewayError> {
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiSyncAttempt, _>(
&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::<AiSyncAttempt, _>(
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiStreamAttempt, _>(
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiSyncAttempt, _>(
&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::<AiSyncAttempt, _>(
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiStreamAttempt, _>(
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<LocalStandardSpec>,
transfer_tracker: &ProviderTransferTracker,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
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::<AiSyncAttempt, _>(
&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::<AiSyncAttempt, _>(
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<LocalStandardSpec>,
transfer_tracker: &ProviderTransferTracker,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
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::<AiStreamAttempt, _>(
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<LocalExecutionRequestOutcome, GatewayError> {
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<LocalExecutionRequestOutcome, GatewayError> {
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiSyncAttempt, _>(
&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::<AiSyncAttempt, _>(
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiStreamAttempt, _>(
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiSyncAttempt, _>(
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<aether_routing_core::RoutingExecutionPolicy>,
) -> 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<F, Fut>(
state: AppState,
parts: http::request::Parts,
trace_id: String,
decision: GatewayControlDecision,
plan_kind: String,
execute: F,
) -> Result<Response<Body>, GatewayError>
where
F: FnOnce(
AppState,
http::request::Parts,
String,
GatewayControlDecision,
String,
Instant,
) -> Fut
+ Send
+ 'static,
Fut: std::future::Future<Output = Result<LocalExecutionRequestOutcome, GatewayError>>
+ 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::<crate::privacy::RedactionSessionSlot>()
.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::<Result<Bytes, IoError>>(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<LocalExecutionRequestOutcome, GatewayError>,
) -> Vec<u8> {
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<Body>,
) -> Vec<u8> {
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<Vec<u8>, 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::<BTreeMap<_, _>>();
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<u8> {
if let Ok(mut value) = serde_json::from_slice::<Value>(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<u8> {
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<AiSyncAttempt>,
transfer_tracker: ProviderTransferTracker,
) -> Result<Response<Body>, 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::<Result<Bytes, IoError>>(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<AiSyncAttempt>,
transfer_tracker: ProviderTransferTracker,
started_at: Instant,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
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<LocalExecutionRequestOutcome, GatewayError>,
) -> Vec<u8> {
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<Body>) -> Vec<u8> {
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<u8> {
if let Ok(mut value) = serde_json::from_slice::<Value>(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<u8> {
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiSyncAttempt, _>(
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiStreamAttempt, _>(
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiStreamAttempt, _>(
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<LocalExecutionRequestOutcome, GatewayError> {
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::<AiSyncAttempt, _>(
state,
parts,
trace_id,
decision,
plan_kind,
attempt_source,
transfer_tracker,
)
.await
}
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<Box<dyn Future<Output = Result<LocalExecutionRequestOutcome, GatewayError>> + Send + 'a>> {
Box::pin(async move {
let Some(decision) = decision else {
return Ok(LocalExecutionRequestOutcome::NoPath);
};
#[cfg(not(test))]
{
if parts.method != http::Method::POST {
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
{
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<Box<dyn Future<Output = Result<LocalExecutionRequestOutcome, GatewayError>> + Send + 'a>> {
Box::pin(async move {
let Some(decision) = decision else {
return Ok(LocalExecutionRequestOutcome::NoPath);
};
#[cfg(not(test))]
{
if parts.method != http::Method::POST {
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
{
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<String>)> {
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<AiSyncAttempt>,
}
impl TestSyncAttemptSource {
fn new(attempts: Vec<AiSyncAttempt>) -> Self {
Self {
attempts: VecDeque::from(attempts),
}
}
}
#[async_trait::async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for TestSyncAttemptSource {
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
Ok(self.attempts.pop_front())
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<AiSyncAttempt>, 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<InMemoryUsageReadRepository>) {
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("<missing>")
);
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::<aether_routing_core::RoutingExecutionPolicy>::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<String>)>::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::<Vec<_>>(),
[
"endpoint-retry",
"endpoint-retry",
"endpoint-retry",
"endpoint-success"
]
);
let sticky_candidate_ids = seen_plans[..3]
.iter()
.map(|(_, candidate_id)| candidate_id.clone())
.collect::<std::collections::BTreeSet<_>>();
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::<crate::executor::candidate_loop::BackgroundAdmissionPermit>()
.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::<AiSyncAttempt, _>(
&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::<AiSyncAttempt, _>(
&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": []}));
}
}
Reference in New Issue View Git Blame Copy Permalink