mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 09:27:46 +08:00
stabeyandClaude Opus 5
04c4a97766
Expose the xAI Imagine image and video surfaces on top of the `xai` provider, and make the shared OpenAI video-task layer survive the production configuration they need. Native video requests live under /v1 (generations, edits, extensions, with /v1/videos as a creation alias that only selects xAI candidates); the OpenAI-compatible adapter stays under /openai/v1/videos and maps `seconds` / `size` onto numeric duration, aspect ratio and resolution. Clients receive an opaque Aether task ID scoped to the owning user; polling uses the upstream task ID and the original credential, and completed downloads fetch the returned media URL without forwarding provider authorization to the media host. Three fixes to the shared video layer are required for this to work outside tests: - OpenAI/xAI task persistence now supplies a stable 16-character short_id, which the PostgreSQL schema requires. Existing rows keep their original value across reconstruction, so no schema change or historical rewrite is needed. - Task retrieval and content downloads are admitted by the production GET execution gate, and reconstructed tasks resolve proxy nodes, system proxy defaults, tunnel affinity and transport profiles through the same deployment resolver used for creation. A configured proxy route no longer silently becomes a direct request after restart. - When the gateway also serves the frontend, /openai/v1/videos and its subpaths bypass the static SPA handler. Otherwise a video query returns HTTP 200 with text/html instead of the task JSON. Co-Authored-By: Claude Opus 5 <[email protected]>
2537 lines
92 KiB
Rust
2537 lines
92 KiB
Rust
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
|
|
}
|
|
|
|
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<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 && !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<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 && !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<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": []}));
|
|
}
|
|
}
|