mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 09:57:47 +08:00
elky
066ea87d72
Add dashboard and overview analytics, health monitoring, provider expense tracking, and announcement updates across the gateway and frontend. Keep schema migrations free of historical backfills while preserving automatic backfill execution. Bound migration deadlines, run schema preparation before Compose replacement, and anonymize deleted dashboard users. Include the current documentation cleanup and regression coverage.
4639 lines
173 KiB
Rust
4639 lines
173 KiB
Rust
use std::collections::BTreeMap;
|
|
use std::io::Error as IoError;
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
|
|
use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope, UPSTREAM_IS_STREAM_KEY};
|
|
use aether_contracts::{
|
|
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan,
|
|
ExecutionResponseObservation, ExecutionResult, ExecutionTelemetry,
|
|
};
|
|
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
|
|
use aether_scheduler_core::{
|
|
execution_error_details, parse_request_candidate_report_context,
|
|
SchedulerRequestCandidateStatusUpdate,
|
|
};
|
|
use aether_usage_runtime::{
|
|
build_lifecycle_usage_seed, build_sync_terminal_usage_payload_seed,
|
|
build_terminal_usage_context_seed, build_usage_event_data_seed, UsageEvent, UsageEventType,
|
|
};
|
|
use async_stream::stream;
|
|
use axum::body::{to_bytes, Body, Bytes};
|
|
use axum::http::header::{CACHE_CONTROL, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE};
|
|
use axum::http::{HeaderName, HeaderValue, Response, StatusCode};
|
|
use futures_util::StreamExt;
|
|
use http_body_util::BodyExt;
|
|
use serde_json::{json, Value};
|
|
use tokio::sync::mpsc;
|
|
use tokio::sync::Mutex;
|
|
use tokio::time::MissedTickBehavior;
|
|
use tracing::{debug, warn};
|
|
|
|
use crate::ai_serving::api::{
|
|
build_core_error_body_for_client_format, extract_provider_private_stream_error_body,
|
|
implicit_sync_finalize_report_kind, maybe_build_sync_finalize_outcome, LocalCoreSyncErrorKind,
|
|
LocalCoreSyncFinalizeOutcome,
|
|
};
|
|
use crate::ai_serving::record_local_runtime_candidate_skip_reason;
|
|
use crate::api::response::{
|
|
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
|
|
build_client_response_from_parts_with_mutator,
|
|
};
|
|
use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
|
|
use crate::control::GatewayControlDecision;
|
|
use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync;
|
|
use crate::execution_runtime::grok::maybe_execute_grok_sync;
|
|
use crate::execution_runtime::kiro_cache::{
|
|
build_kiro_prompt_cache_profile, compute_kiro_prompt_cache_usage,
|
|
estimate_kiro_prompt_input_tokens, kiro_simulated_cache_enabled_from_provider_config,
|
|
kiro_simulated_cache_enabled_from_report_context, KiroPromptCacheUsage,
|
|
KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD,
|
|
};
|
|
use crate::execution_runtime::oauth_retry::refresh_oauth_plan_auth_for_retry;
|
|
#[cfg(test)]
|
|
use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_runtime;
|
|
use crate::execution_runtime::submission::{
|
|
resolve_local_sync_error_status_code, submit_local_core_error_or_sync_finalize,
|
|
};
|
|
use crate::execution_runtime::transport::{
|
|
append_upstream_response_body_chunk_with_limit, build_execution_response_body,
|
|
build_request_body, collect_response_headers, decode_response_body_bytes_with_limit,
|
|
execution_plan_response_body_limit_bytes, execution_response_body_mode,
|
|
format_hyper_error_chain, format_upstream_request_error, format_wreq_upstream_request_error,
|
|
response_body_is_json, safe_transport_error_message, send_request, DirectHttpResponse,
|
|
DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
|
|
};
|
|
use crate::execution_runtime::windsurf::maybe_execute_windsurf_sync;
|
|
use crate::execution_runtime::{
|
|
ai_attempt_retry_scope_from_failure_disposition, analyze_local_candidate_failover_sync,
|
|
apply_endpoint_response_header_rules, attach_provider_response_headers_to_report_context,
|
|
local_failover_response_text, resolve_core_sync_error_finalize_report_kind,
|
|
should_fallback_to_control_sync, should_finalize_sync_response, LocalFailoverDecision,
|
|
};
|
|
use crate::log_ids::short_request_id;
|
|
use crate::orchestration::{
|
|
apply_local_execution_effect, build_local_error_flow_metadata,
|
|
spawn_local_oauth_success_effect, trace_upstream_response_body, with_error_flow_report_context,
|
|
with_upstream_response_report_context, LocalAdaptiveRateLimitEffect,
|
|
LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
|
|
LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect,
|
|
LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, LocalPoolErrorEffect,
|
|
};
|
|
use crate::provider_pool_demand::{
|
|
acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission,
|
|
};
|
|
use crate::request_candidate_runtime::{
|
|
ensure_execution_request_candidate_slot, record_local_request_candidate_extra_data,
|
|
record_local_request_candidate_status, record_local_request_candidate_status_snapshot,
|
|
snapshot_local_request_candidate_status,
|
|
};
|
|
use crate::request_diagnostics::{
|
|
attach_current_request_diagnostics_and_candidate_start_timing_to_report_context,
|
|
attach_request_diagnostics_to_report_context, calibrate_candidate_first_byte_elapsed_ms,
|
|
current_request_diagnostics, RequestDiagnostics,
|
|
};
|
|
use crate::usage::{spawn_sync_report, submit_sync_report};
|
|
use crate::video_tasks::VideoTaskSyncReportMode;
|
|
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
|
|
|
|
#[path = "execution/policy.rs"]
|
|
mod policy;
|
|
#[path = "execution/response.rs"]
|
|
mod response;
|
|
|
|
use policy::{decode_execution_result_body, decode_execution_result_body_with_limit};
|
|
pub(crate) use response::{
|
|
maybe_build_local_sync_finalize_response, maybe_build_local_video_error_response,
|
|
maybe_build_local_video_success_outcome, resolve_local_sync_error_background_report_kind,
|
|
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessBuild,
|
|
LocalVideoSyncSuccessOutcome,
|
|
};
|
|
|
|
const OPENAI_IMAGE_SYNC_PLAN_KIND: &str = "openai_image_sync";
|
|
const OPENAI_IMAGE_SYNC_DEFAULT_TOTAL_TIMEOUT_MS: u64 = 900_000;
|
|
const SYNC_EXECUTION_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60);
|
|
const OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15);
|
|
const OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_BYTES: &[u8] = b"\n";
|
|
const OPENAI_IMAGE_SYNC_PROGRESS_WRITE_INTERVAL: Duration = Duration::from_secs(5);
|
|
// Progress parsing must retain only the incomplete SSE record. This is not a
|
|
// response-body limit: the full upstream body is still handled by the normal
|
|
// execution body policy, while malformed streams cannot grow telemetry state
|
|
// forever by withholding a record separator.
|
|
const OPENAI_IMAGE_SYNC_PROGRESS_MAX_BUFFER_BYTES: usize = 16 * 1024 * 1024;
|
|
const INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE: &str = "Provider returned HTTP 200 but the Gemini response did not contain visible model output; refusing to finalize it as a successful response.";
|
|
|
|
fn elapsed_ms_since(started_at: Instant) -> u64 {
|
|
started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64
|
|
}
|
|
|
|
fn calibrated_sync_candidate_first_byte_elapsed_ms(
|
|
candidate_started_at: Instant,
|
|
result: &ExecutionResult,
|
|
) -> Option<u64> {
|
|
let telemetry = result.telemetry.as_ref()?;
|
|
calibrate_candidate_first_byte_elapsed_ms(
|
|
elapsed_ms_since(candidate_started_at),
|
|
telemetry.elapsed_ms,
|
|
telemetry.ttfb_ms,
|
|
)
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct SyncExecutionFailure {
|
|
error_type: &'static str,
|
|
message: String,
|
|
status_code: Option<u16>,
|
|
latency_ms: Option<u64>,
|
|
fallback_kind: Option<SyncExecutionFailureFallbackKind>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
enum SyncExecutionFailureFallbackKind {
|
|
UpstreamResponseTooLarge,
|
|
UpstreamResponseDecode,
|
|
}
|
|
|
|
impl SyncExecutionFailureFallbackKind {
|
|
fn error_type(self) -> &'static str {
|
|
match self {
|
|
Self::UpstreamResponseTooLarge => "upstream_response_too_large",
|
|
Self::UpstreamResponseDecode => "upstream_response_decode_failed",
|
|
}
|
|
}
|
|
|
|
fn client_message(self) -> &'static str {
|
|
match self {
|
|
Self::UpstreamResponseTooLarge => "Upstream response too large",
|
|
Self::UpstreamResponseDecode => "Failed to decode upstream response",
|
|
}
|
|
}
|
|
}
|
|
|
|
struct SyncAttemptTerminalGuard {
|
|
state: AppState,
|
|
plan: ExecutionPlan,
|
|
report_context: Option<Value>,
|
|
request_diagnostics: Option<Arc<RequestDiagnostics>>,
|
|
candidate_started_unix_ms: u64,
|
|
candidate_started_at: Instant,
|
|
armed: bool,
|
|
}
|
|
|
|
/// Keep forced-terminal records useful for operations without copying an
|
|
/// arbitrary `GatewayError` into the candidate/usage stores. Gateway errors
|
|
/// can wrap provider URLs, credentials, query strings, or database details;
|
|
/// those values belong in the internal logging path only.
|
|
fn persisted_sync_abort_message(error: &GatewayError) -> &'static str {
|
|
match error {
|
|
GatewayError::UpstreamUnavailable { .. } => {
|
|
"local sync attempt aborted before terminal finalization: upstream unavailable"
|
|
}
|
|
GatewayError::ControlUnavailable { .. } => {
|
|
"local sync attempt aborted before terminal finalization: control unavailable"
|
|
}
|
|
GatewayError::LocalExecutionPlanningTimeout { .. } => {
|
|
"local sync attempt aborted before terminal finalization: planning timeout"
|
|
}
|
|
GatewayError::AdmissionTimeout { .. } => {
|
|
"local sync attempt aborted before terminal finalization: admission timeout"
|
|
}
|
|
GatewayError::Client { .. } => {
|
|
"local sync attempt aborted before terminal finalization: client error"
|
|
}
|
|
GatewayError::PlanUsageLimited(_) => {
|
|
"local sync attempt aborted before terminal finalization: usage limit"
|
|
}
|
|
GatewayError::LastActiveAdminUpdateDenied | GatewayError::LastActiveAdminDeleteDenied => {
|
|
"local sync attempt aborted before terminal finalization: policy denied"
|
|
}
|
|
GatewayError::Internal(_) => {
|
|
"local sync attempt aborted before terminal finalization: internal error"
|
|
}
|
|
}
|
|
}
|
|
|
|
impl SyncAttemptTerminalGuard {
|
|
fn new(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<Value>,
|
|
candidate_started_unix_ms: u64,
|
|
candidate_started_at: Instant,
|
|
) -> Self {
|
|
Self {
|
|
state: state.clone(),
|
|
plan: plan.clone(),
|
|
report_context,
|
|
request_diagnostics: current_request_diagnostics(),
|
|
candidate_started_unix_ms,
|
|
candidate_started_at,
|
|
armed: true,
|
|
}
|
|
}
|
|
|
|
fn disarm(&mut self) {
|
|
self.armed = false;
|
|
}
|
|
|
|
async fn fail_and_disarm(&mut self, error: &GatewayError) {
|
|
if !self.armed {
|
|
return;
|
|
}
|
|
self.armed = false;
|
|
record_sync_attempt_forced_terminal_state(
|
|
self.state.clone(),
|
|
self.plan.clone(),
|
|
crate::usage::reporting::failure::gateway_error_analytics_context(
|
|
self.report_context.as_ref(),
|
|
error,
|
|
),
|
|
self.request_diagnostics.clone(),
|
|
self.candidate_started_unix_ms,
|
|
self.candidate_started_at,
|
|
UsageEventType::Failed,
|
|
RequestCandidateStatus::Failed,
|
|
StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
|
|
"local_sync_attempt_aborted",
|
|
persisted_sync_abort_message(error),
|
|
)
|
|
.await;
|
|
}
|
|
}
|
|
|
|
impl Drop for SyncAttemptTerminalGuard {
|
|
fn drop(&mut self) {
|
|
if !self.armed {
|
|
return;
|
|
}
|
|
self.armed = false;
|
|
let state = self.state.clone();
|
|
let plan = self.plan.clone();
|
|
let report_context = self.report_context.clone();
|
|
let request_diagnostics = self.request_diagnostics.clone();
|
|
let candidate_started_unix_ms = self.candidate_started_unix_ms;
|
|
let candidate_started_at = self.candidate_started_at;
|
|
if let Ok(handle) = tokio::runtime::Handle::try_current() {
|
|
let usage_producer = state.usage_runtime.track_producer();
|
|
handle.spawn(async move {
|
|
let _usage_producer = usage_producer;
|
|
record_sync_attempt_forced_terminal_state(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
request_diagnostics,
|
|
candidate_started_unix_ms,
|
|
candidate_started_at,
|
|
UsageEventType::Cancelled,
|
|
RequestCandidateStatus::Cancelled,
|
|
499,
|
|
"local_sync_attempt_cancelled",
|
|
"Local sync attempt was dropped before terminal finalization, usually because the client disconnected or the request task was cancelled.",
|
|
)
|
|
.await;
|
|
});
|
|
} else {
|
|
warn!(
|
|
event_name = "local_sync_attempt_terminal_guard_no_runtime",
|
|
log_type = "ops",
|
|
request_id = %short_request_id(self.plan.request_id.as_str()),
|
|
candidate_id = ?self.plan.candidate_id,
|
|
"gateway could not finalize dropped local sync attempt because no Tokio runtime is available"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
async fn record_sync_attempt_forced_terminal_state(
|
|
state: AppState,
|
|
plan: ExecutionPlan,
|
|
report_context: Option<Value>,
|
|
request_diagnostics: Option<Arc<RequestDiagnostics>>,
|
|
candidate_started_unix_ms: u64,
|
|
candidate_started_at: Instant,
|
|
usage_event_type: UsageEventType,
|
|
candidate_status: RequestCandidateStatus,
|
|
status_code: u16,
|
|
error_type: &'static str,
|
|
error_message: impl Into<String>,
|
|
) {
|
|
let error_message = error_message.into();
|
|
let report_context =
|
|
attach_request_diagnostics_to_report_context(report_context, request_diagnostics.as_ref());
|
|
let report_context = if matches!(usage_event_type, UsageEventType::Cancelled) {
|
|
crate::usage::reporting::failure::with_analytics_failure(
|
|
report_context.as_ref(),
|
|
"unknown",
|
|
"finalize",
|
|
"request_task_cancelled",
|
|
)
|
|
} else {
|
|
report_context
|
|
};
|
|
let terminal_unix_ms = current_request_candidate_unix_ms();
|
|
let latency_ms = elapsed_ms_since(candidate_started_at);
|
|
record_local_request_candidate_status(
|
|
&state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: candidate_status,
|
|
status_code: Some(status_code),
|
|
error_type: Some(error_type.to_string()),
|
|
error_message: Some(error_message.clone()),
|
|
latency_ms: Some(latency_ms),
|
|
started_at_unix_ms: Some(candidate_started_unix_ms),
|
|
finished_at_unix_ms: Some(terminal_unix_ms),
|
|
},
|
|
)
|
|
.await;
|
|
|
|
if !state.usage_runtime.is_enabled() {
|
|
return;
|
|
}
|
|
|
|
let mut usage_data = build_usage_event_data_seed(&plan, report_context.as_ref());
|
|
usage_data.status_code = Some(status_code);
|
|
usage_data.error_message = Some(error_message.clone());
|
|
usage_data.error_category = Some(
|
|
match usage_event_type {
|
|
UsageEventType::Cancelled => "cancelled",
|
|
_ => "server_error",
|
|
}
|
|
.to_string(),
|
|
);
|
|
usage_data.response_time_ms = Some(latency_ms);
|
|
let error_body = json!({
|
|
"error": {
|
|
"type": error_type,
|
|
"message": error_message,
|
|
"code": status_code
|
|
}
|
|
});
|
|
usage_data.response_headers = Some(json!({"content-type": "application/json"}));
|
|
usage_data.response_body = Some(error_body.clone());
|
|
usage_data.client_response_headers = Some(json!({"content-type": "application/json"}));
|
|
usage_data.client_response_body = Some(error_body);
|
|
|
|
state
|
|
.usage_runtime
|
|
.record_terminal_event_direct(
|
|
state.usage_lifecycle_data_state().as_ref(),
|
|
UsageEvent::new(usage_event_type, plan.request_id.clone(), usage_data),
|
|
)
|
|
.await;
|
|
}
|
|
|
|
impl SyncExecutionFailure {
|
|
fn from_transport(err: ExecutionRuntimeTransportError) -> Self {
|
|
let fallback_kind = match &err {
|
|
ExecutionRuntimeTransportError::UpstreamResponseTooLarge { .. } => {
|
|
Some(SyncExecutionFailureFallbackKind::UpstreamResponseTooLarge)
|
|
}
|
|
ExecutionRuntimeTransportError::UpstreamResponseDecode { .. } => {
|
|
Some(SyncExecutionFailureFallbackKind::UpstreamResponseDecode)
|
|
}
|
|
_ => None,
|
|
};
|
|
Self {
|
|
error_type: fallback_kind
|
|
.map(SyncExecutionFailureFallbackKind::error_type)
|
|
.unwrap_or("execution_runtime_unavailable"),
|
|
message: safe_transport_error_message(&err),
|
|
status_code: fallback_kind.map(|_| StatusCode::BAD_GATEWAY.as_u16()),
|
|
latency_ms: None,
|
|
fallback_kind,
|
|
}
|
|
}
|
|
|
|
fn image_sync_total_timeout(timeout_ms: u64, elapsed_ms: u64) -> Self {
|
|
Self {
|
|
error_type: "image_sync_total_timeout",
|
|
message: format!(
|
|
"OpenAI image sync execution exceeded total timeout of {timeout_ms}ms"
|
|
),
|
|
status_code: Some(StatusCode::GATEWAY_TIMEOUT.as_u16()),
|
|
latency_ms: Some(elapsed_ms),
|
|
fallback_kind: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
fn build_sync_execution_failure_fallback_body(
|
|
client_api_format: &str,
|
|
kind: SyncExecutionFailureFallbackKind,
|
|
) -> Value {
|
|
let message = kind.client_message();
|
|
let error_type = kind.error_type();
|
|
match crate::ai_serving::normalize_api_format_alias(client_api_format).as_str() {
|
|
"claude:messages" => json!({
|
|
"type": "error",
|
|
"error": {
|
|
"type": "upstream_error",
|
|
"message": message,
|
|
}
|
|
}),
|
|
"gemini:generate_content" => json!({
|
|
"error": {
|
|
"code": StatusCode::BAD_GATEWAY.as_u16(),
|
|
"message": message,
|
|
"status": "BAD_GATEWAY",
|
|
}
|
|
}),
|
|
_ => json!({
|
|
"error": {
|
|
"type": "upstream_error",
|
|
"message": message,
|
|
"code": error_type,
|
|
}
|
|
}),
|
|
}
|
|
}
|
|
|
|
fn build_sync_execution_failure_fallback_response(
|
|
failure: &SyncExecutionFailure,
|
|
plan: &ExecutionPlan,
|
|
trace_id: &str,
|
|
decision: &GatewayControlDecision,
|
|
) -> Result<Option<Response<Body>>, GatewayError> {
|
|
let Some(kind) = failure.fallback_kind else {
|
|
return Ok(None);
|
|
};
|
|
let body_json = build_sync_execution_failure_fallback_body(&plan.client_api_format, kind);
|
|
let body_bytes = serde_json::to_vec(&body_json)
|
|
.map_err(|error| GatewayError::Internal(error.to_string()))?;
|
|
let headers = BTreeMap::from([
|
|
("content-type".to_string(), "application/json".to_string()),
|
|
("content-length".to_string(), body_bytes.len().to_string()),
|
|
]);
|
|
let response = build_client_response_from_parts(
|
|
StatusCode::BAD_GATEWAY.as_u16(),
|
|
&headers,
|
|
Body::from(body_bytes),
|
|
trace_id,
|
|
Some(decision),
|
|
)?;
|
|
attach_control_metadata_headers(
|
|
response,
|
|
Some(plan.request_id.as_str()),
|
|
plan.candidate_id.as_deref(),
|
|
)
|
|
.map(Some)
|
|
}
|
|
|
|
fn maybe_store_sync_execution_failure_fallback(
|
|
failure: &SyncExecutionFailure,
|
|
plan: &ExecutionPlan,
|
|
trace_id: &str,
|
|
decision: &GatewayControlDecision,
|
|
retry_scope_out: &mut Option<&mut AiAttemptRetryScope>,
|
|
retry_fallback_out: &mut Option<&mut Option<Response<Body>>>,
|
|
) -> Result<(), GatewayError> {
|
|
if failure.fallback_kind.is_none() {
|
|
return Ok(());
|
|
}
|
|
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
|
|
*retry_scope = AiAttemptRetryScope::Candidate;
|
|
}
|
|
if let Some(retry_fallback) = retry_fallback_out.as_deref_mut() {
|
|
*retry_fallback =
|
|
build_sync_execution_failure_fallback_response(failure, plan, trace_id, decision)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn maybe_build_sync_transport_error_stop_response(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&Value>,
|
|
trace_id: &str,
|
|
decision: &GatewayControlDecision,
|
|
error_type: &str,
|
|
error_message: &str,
|
|
elapsed_ms: u64,
|
|
) -> Result<Option<Response<Body>>, GatewayError> {
|
|
let analysis = crate::orchestration::resolve_local_transport_failover_analysis_for_attempt(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
)
|
|
.await;
|
|
if !matches!(analysis.decision, LocalFailoverDecision::StopLocalFailover) {
|
|
return Ok(None);
|
|
}
|
|
|
|
crate::execution_runtime::build_transport_error_stop_response(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
trace_id,
|
|
decision,
|
|
StatusCode::BAD_GATEWAY.as_u16(),
|
|
error_type,
|
|
error_message,
|
|
elapsed_ms,
|
|
)
|
|
.await
|
|
.map(Some)
|
|
}
|
|
|
|
struct ImplicitSyncFinalizeOutcome {
|
|
payload: GatewaySyncReportRequest,
|
|
outcome: LocalCoreSyncFinalizeOutcome,
|
|
}
|
|
|
|
fn spawn_sync_candidate_status_update(
|
|
state: AppState,
|
|
snapshot: crate::request_candidate_runtime::LocalRequestCandidateStatusSnapshot,
|
|
status_update: SchedulerRequestCandidateStatusUpdate,
|
|
) {
|
|
tokio::spawn(async move {
|
|
record_local_request_candidate_status_snapshot(&state, &snapshot, status_update).await;
|
|
});
|
|
}
|
|
|
|
fn record_sync_response_started(
|
|
state: &AppState,
|
|
lifecycle_seed: aether_usage_runtime::LifecycleUsageSeed,
|
|
request_candidate_status_snapshot: Option<
|
|
crate::request_candidate_runtime::LocalRequestCandidateStatusSnapshot,
|
|
>,
|
|
candidate_started_unix_ms: u64,
|
|
status_code: u16,
|
|
ttfb_ms: u64,
|
|
) {
|
|
state.usage_runtime.record_stream_started_immediate_async(
|
|
state.usage_lifecycle_data_state().as_ref(),
|
|
lifecycle_seed,
|
|
status_code,
|
|
Some(ExecutionTelemetry {
|
|
ttfb_ms: Some(ttfb_ms),
|
|
elapsed_ms: Some(ttfb_ms),
|
|
upstream_bytes: None,
|
|
}),
|
|
);
|
|
|
|
if let Some(snapshot) = request_candidate_status_snapshot {
|
|
spawn_sync_candidate_status_update(
|
|
state.clone(),
|
|
snapshot,
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Streaming,
|
|
status_code: Some(status_code),
|
|
error_type: None,
|
|
error_message: None,
|
|
latency_ms: Some(ttfb_ms),
|
|
started_at_unix_ms: Some(candidate_started_unix_ms),
|
|
finished_at_unix_ms: None,
|
|
},
|
|
);
|
|
}
|
|
}
|
|
|
|
fn record_sync_execution_active(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&Value>,
|
|
candidate_started_unix_ms: u64,
|
|
) {
|
|
let lifecycle_seed = build_lifecycle_usage_seed(plan, report_context);
|
|
state.usage_runtime.record_sync_active_immediate_async(
|
|
state.usage_lifecycle_data_state().as_ref(),
|
|
lifecycle_seed,
|
|
);
|
|
|
|
if let Some(snapshot) = snapshot_local_request_candidate_status(plan, report_context) {
|
|
spawn_sync_candidate_status_update(
|
|
state.clone(),
|
|
snapshot,
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Streaming,
|
|
status_code: None,
|
|
error_type: None,
|
|
error_message: None,
|
|
latency_ms: None,
|
|
started_at_unix_ms: Some(candidate_started_unix_ms),
|
|
finished_at_unix_ms: None,
|
|
},
|
|
);
|
|
}
|
|
}
|
|
|
|
async fn record_sync_terminal_usage(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&serde_json::Value>,
|
|
payload: &GatewaySyncReportRequest,
|
|
candidate_started_at: Instant,
|
|
candidate_first_byte_elapsed_ms: Option<u64>,
|
|
) {
|
|
let analytics_context =
|
|
crate::usage::reporting::failure::sync_analytics_context(report_context, payload);
|
|
let report_context_with_diagnostics =
|
|
attach_current_request_diagnostics_and_candidate_start_timing_to_report_context(
|
|
analytics_context.as_ref(),
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
);
|
|
let context_seed = build_terminal_usage_context_seed(
|
|
plan,
|
|
report_context_with_diagnostics
|
|
.as_ref()
|
|
.or(analytics_context.as_ref()),
|
|
);
|
|
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
|
|
state
|
|
.usage_runtime
|
|
.record_sync_terminal(
|
|
state.usage_lifecycle_data_state().as_ref(),
|
|
context_seed,
|
|
payload_seed,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
async fn record_sync_terminal_usage_and_disarm_guard(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&serde_json::Value>,
|
|
payload: &GatewaySyncReportRequest,
|
|
candidate_started_at: Instant,
|
|
candidate_first_byte_elapsed_ms: Option<u64>,
|
|
terminal_guard: &mut SyncAttemptTerminalGuard,
|
|
) {
|
|
record_sync_terminal_usage(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
)
|
|
.await;
|
|
terminal_guard.disarm();
|
|
}
|
|
|
|
fn with_sync_error_trace_context(
|
|
report_context: Option<&serde_json::Value>,
|
|
status_code: u16,
|
|
headers: &BTreeMap<String, String>,
|
|
body_json: Option<&serde_json::Value>,
|
|
body_bytes: &[u8],
|
|
response_text: Option<&str>,
|
|
local_failover_analysis: crate::orchestration::LocalFailoverAnalysis,
|
|
) -> Option<serde_json::Value> {
|
|
let body = trace_upstream_response_body(body_json, body_bytes);
|
|
let upstream_context = with_upstream_response_report_context(
|
|
report_context,
|
|
status_code,
|
|
Some(headers),
|
|
body.as_ref(),
|
|
None,
|
|
None,
|
|
);
|
|
with_error_flow_report_context(
|
|
upstream_context.as_ref().or(report_context),
|
|
build_local_error_flow_metadata(status_code, response_text, local_failover_analysis),
|
|
)
|
|
}
|
|
|
|
fn build_sync_report_payload(
|
|
trace_id: &str,
|
|
report_kind: String,
|
|
report_context: Option<serde_json::Value>,
|
|
status_code: u16,
|
|
headers: BTreeMap<String, String>,
|
|
body_json: Option<serde_json::Value>,
|
|
body_base64: Option<String>,
|
|
telemetry: Option<ExecutionTelemetry>,
|
|
) -> GatewaySyncReportRequest {
|
|
GatewaySyncReportRequest {
|
|
trace_id: trace_id.to_string(),
|
|
report_kind,
|
|
report_context,
|
|
status_code,
|
|
headers,
|
|
body_json,
|
|
client_body_json: None,
|
|
body_base64,
|
|
telemetry,
|
|
}
|
|
}
|
|
|
|
fn seed_kiro_sync_report_context_input_tokens(
|
|
plan: &ExecutionPlan,
|
|
report_context: &mut Option<Value>,
|
|
) {
|
|
if !plan
|
|
.provider_name
|
|
.as_deref()
|
|
.is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro"))
|
|
{
|
|
return;
|
|
}
|
|
|
|
let Some(context) = report_context.as_mut().and_then(Value::as_object_mut) else {
|
|
return;
|
|
};
|
|
if context
|
|
.get("input_tokens")
|
|
.and_then(Value::as_u64)
|
|
.is_some_and(|input_tokens| input_tokens > 0)
|
|
{
|
|
return;
|
|
}
|
|
|
|
let Some(original_request_body) = context.get("original_request_body").cloned() else {
|
|
return;
|
|
};
|
|
let estimated_input_tokens = estimate_kiro_prompt_input_tokens(&original_request_body);
|
|
context.insert(
|
|
"input_tokens".to_string(),
|
|
Value::from(estimated_input_tokens),
|
|
);
|
|
}
|
|
|
|
async fn seed_kiro_sync_simulated_cache_enabled(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: &mut Option<Value>,
|
|
) {
|
|
if !plan
|
|
.provider_name
|
|
.as_deref()
|
|
.is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro"))
|
|
{
|
|
return;
|
|
}
|
|
|
|
let enabled = match state
|
|
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&plan.provider_id))
|
|
.await
|
|
{
|
|
Ok(providers) => providers
|
|
.iter()
|
|
.find(|provider| provider.id == plan.provider_id)
|
|
.filter(|provider| provider.provider_type.eq_ignore_ascii_case("kiro"))
|
|
.is_some_and(|provider| {
|
|
kiro_simulated_cache_enabled_from_provider_config(provider.config.as_ref())
|
|
}),
|
|
Err(err) => {
|
|
warn!(
|
|
event_name = "kiro_simulated_cache_config_read_failed",
|
|
log_type = "event",
|
|
request_id = %plan.request_id,
|
|
provider_id = %plan.provider_id,
|
|
error = ?err,
|
|
"failed to read Kiro simulated cache provider config; defaulting disabled"
|
|
);
|
|
false
|
|
}
|
|
};
|
|
|
|
let Some(context) = report_context.as_mut().and_then(Value::as_object_mut) else {
|
|
return;
|
|
};
|
|
if enabled {
|
|
context.insert(
|
|
KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD.to_string(),
|
|
Value::Bool(true),
|
|
);
|
|
} else {
|
|
context.remove(KIRO_SIMULATED_CACHE_ENABLED_CONTEXT_FIELD);
|
|
}
|
|
}
|
|
|
|
async fn seed_kiro_sync_report_context_prompt_cache_usage(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: &mut Option<Value>,
|
|
) {
|
|
if !plan
|
|
.provider_name
|
|
.as_deref()
|
|
.is_some_and(|provider_name| provider_name.eq_ignore_ascii_case("Kiro"))
|
|
{
|
|
return;
|
|
}
|
|
|
|
let simulated_cache_enabled =
|
|
kiro_simulated_cache_enabled_from_report_context(report_context.as_ref());
|
|
let Some(context) = report_context.as_mut().and_then(Value::as_object_mut) else {
|
|
return;
|
|
};
|
|
if context
|
|
.get("kiro_web_search_mcp")
|
|
.and_then(Value::as_bool)
|
|
.unwrap_or(false)
|
|
{
|
|
return;
|
|
}
|
|
if !simulated_cache_enabled {
|
|
return;
|
|
}
|
|
if kiro_cache_usage_from_context_object(context).is_some() {
|
|
return;
|
|
}
|
|
|
|
let Some(original_request_body) = context.get("original_request_body").cloned() else {
|
|
return;
|
|
};
|
|
let input_tokens = context
|
|
.get("input_tokens")
|
|
.and_then(Value::as_u64)
|
|
.filter(|value| *value > 0)
|
|
.unwrap_or_else(|| {
|
|
let estimated = estimate_kiro_prompt_input_tokens(&original_request_body);
|
|
context.insert("input_tokens".to_string(), Value::from(estimated));
|
|
estimated
|
|
});
|
|
let Some(profile) = build_kiro_prompt_cache_profile(&original_request_body, input_tokens)
|
|
else {
|
|
return;
|
|
};
|
|
|
|
let cache_usage = compute_kiro_prompt_cache_usage(
|
|
state.runtime_state(),
|
|
kiro_sync_cache_credential_id(plan),
|
|
&profile,
|
|
)
|
|
.await;
|
|
if cache_usage.cache_creation_input_tokens == 0 && cache_usage.cache_read_input_tokens == 0 {
|
|
return;
|
|
}
|
|
context.insert(
|
|
"cache_creation_input_tokens".to_string(),
|
|
Value::from(cache_usage.cache_creation_input_tokens),
|
|
);
|
|
context.insert(
|
|
"cache_read_input_tokens".to_string(),
|
|
Value::from(cache_usage.cache_read_input_tokens),
|
|
);
|
|
}
|
|
|
|
fn kiro_sync_cache_credential_id(plan: &ExecutionPlan) -> String {
|
|
format!("{}:{}:{}", plan.provider_id, plan.endpoint_id, plan.key_id)
|
|
}
|
|
|
|
fn kiro_cache_usage_from_context_object(
|
|
context: &serde_json::Map<String, Value>,
|
|
) -> Option<KiroPromptCacheUsage> {
|
|
let cache_creation_input_tokens = context
|
|
.get("cache_creation_input_tokens")
|
|
.and_then(Value::as_u64)
|
|
.unwrap_or(0);
|
|
let cache_read_input_tokens = context
|
|
.get("cache_read_input_tokens")
|
|
.and_then(Value::as_u64)
|
|
.unwrap_or(0);
|
|
(cache_creation_input_tokens > 0 || cache_read_input_tokens > 0).then_some(
|
|
KiroPromptCacheUsage {
|
|
cache_creation_input_tokens,
|
|
cache_read_input_tokens,
|
|
},
|
|
)
|
|
}
|
|
|
|
fn invalid_gemini_provider_success_message(
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&Value>,
|
|
status_code: u16,
|
|
body_json: Option<&Value>,
|
|
) -> Option<&'static str> {
|
|
if status_code >= 400 {
|
|
return None;
|
|
}
|
|
if !provider_api_format_is_gemini_generate_content(plan, report_context) {
|
|
return None;
|
|
}
|
|
let body_json = body_json?;
|
|
if body_json
|
|
.as_object()
|
|
.is_some_and(|object| object.get("error").is_some_and(|error| !error.is_null()))
|
|
{
|
|
return None;
|
|
}
|
|
let normalized_body_json = report_context
|
|
.filter(|context| {
|
|
context
|
|
.get("has_envelope")
|
|
.and_then(Value::as_bool)
|
|
.unwrap_or(false)
|
|
})
|
|
.and_then(|context| {
|
|
crate::ai_serving::normalize_provider_private_response_value(body_json.clone(), context)
|
|
});
|
|
let body_json = normalized_body_json.as_ref().unwrap_or(body_json);
|
|
if crate::ai_serving::gemini_generate_content_response_has_visible_output(body_json) {
|
|
return None;
|
|
}
|
|
Some(INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE)
|
|
}
|
|
|
|
fn invalid_gemini_provider_stream_success_message(
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&Value>,
|
|
status_code: u16,
|
|
body_json: Option<&Value>,
|
|
body_bytes: &[u8],
|
|
has_body_bytes: bool,
|
|
) -> Option<&'static str> {
|
|
if status_code >= 400 || body_json.is_some() || !has_body_bytes {
|
|
return None;
|
|
}
|
|
if !provider_api_format_is_gemini_generate_content(plan, report_context) {
|
|
return None;
|
|
}
|
|
let Some(body_json) = crate::ai_serving::aggregate_gemini_stream_sync_response(body_bytes)
|
|
else {
|
|
return Some(INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE);
|
|
};
|
|
if crate::ai_serving::gemini_generate_content_response_has_visible_output(&body_json) {
|
|
return None;
|
|
}
|
|
Some(INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE)
|
|
}
|
|
|
|
fn provider_api_format_is_gemini_generate_content(
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&Value>,
|
|
) -> bool {
|
|
let provider_api_format = report_context
|
|
.and_then(|value| value.get("provider_api_format"))
|
|
.and_then(Value::as_str)
|
|
.unwrap_or(plan.provider_api_format.as_str());
|
|
crate::ai_serving::normalize_api_format_alias(provider_api_format) == "gemini:generate_content"
|
|
}
|
|
|
|
fn invalid_gemini_provider_success_execution_error(message: &str) -> ExecutionError {
|
|
ExecutionError {
|
|
kind: ExecutionErrorKind::Upstream5xx,
|
|
phase: ExecutionPhase::Finalize,
|
|
message: message.to_string(),
|
|
upstream_status: Some(StatusCode::OK.as_u16()),
|
|
retryable: true,
|
|
failover_recommended: true,
|
|
}
|
|
}
|
|
|
|
fn build_invalid_provider_success_body(
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&Value>,
|
|
message: &str,
|
|
) -> Option<Value> {
|
|
let client_api_format = report_context
|
|
.and_then(|value| value.get("client_api_format"))
|
|
.and_then(Value::as_str)
|
|
.unwrap_or(plan.client_api_format.as_str());
|
|
build_core_error_body_for_client_format(
|
|
client_api_format,
|
|
message,
|
|
Some("invalid_provider_success_response"),
|
|
LocalCoreSyncErrorKind::ServerError,
|
|
)
|
|
}
|
|
|
|
fn provider_private_error_details(body_json: &Value) -> (Option<String>, Option<String>) {
|
|
let body_object = body_json.as_object();
|
|
let error_object = body_object
|
|
.and_then(|object| object.get("error"))
|
|
.and_then(Value::as_object);
|
|
let error_type =
|
|
first_non_empty_error_text(error_object, body_object, &["type", "code", "status"]);
|
|
let error_message = first_non_empty_error_text(
|
|
error_object,
|
|
body_object,
|
|
&["message", "detail", "reason", "status", "type", "code"],
|
|
);
|
|
(error_type, error_message)
|
|
}
|
|
|
|
fn first_non_empty_error_text(
|
|
error_object: Option<&serde_json::Map<String, Value>>,
|
|
body_object: Option<&serde_json::Map<String, Value>>,
|
|
keys: &[&str],
|
|
) -> Option<String> {
|
|
for object in [error_object, body_object].into_iter().flatten() {
|
|
for key in keys {
|
|
let Some(value) = object.get(*key) else {
|
|
continue;
|
|
};
|
|
match value {
|
|
Value::String(text) if !text.trim().is_empty() => {
|
|
return Some(text.trim().to_string());
|
|
}
|
|
Value::Number(number) => return Some(number.to_string()),
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct OpenAiImageSyncProgressSnapshot {
|
|
phase: &'static str,
|
|
upstream_ttfb_ms: Option<u64>,
|
|
upstream_sse_frame_count: u64,
|
|
last_upstream_event: Option<String>,
|
|
last_upstream_frame_at_unix_ms: Option<u64>,
|
|
partial_image_count: u64,
|
|
last_client_visible_event: Option<String>,
|
|
downstream_heartbeat_count: u64,
|
|
last_downstream_heartbeat_at_unix_ms: Option<u64>,
|
|
downstream_heartbeat_interval_ms: Option<u64>,
|
|
}
|
|
|
|
struct OpenAiImageSyncProgressRecorder<'a> {
|
|
state: &'a AppState,
|
|
plan: &'a ExecutionPlan,
|
|
report_context: Option<&'a Value>,
|
|
snapshot: Arc<Mutex<OpenAiImageSyncProgressSnapshot>>,
|
|
buffer: Vec<u8>,
|
|
last_persist_at: Option<Instant>,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct OpenAiImageSyncJsonHeartbeatContext {
|
|
state: AppState,
|
|
plan: ExecutionPlan,
|
|
report_context: Option<Value>,
|
|
snapshot: Arc<Mutex<OpenAiImageSyncProgressSnapshot>>,
|
|
started_at: Instant,
|
|
trace_id: String,
|
|
request_id_for_log: String,
|
|
candidate_id: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct OpenAiImageSyncSseFrame {
|
|
event_name: String,
|
|
is_partial_image: bool,
|
|
is_completed: bool,
|
|
is_failed: bool,
|
|
client_visible_event: Option<&'static str>,
|
|
}
|
|
|
|
impl OpenAiImageSyncProgressSnapshot {
|
|
fn new() -> Self {
|
|
Self {
|
|
phase: "upstream_connecting",
|
|
upstream_ttfb_ms: None,
|
|
upstream_sse_frame_count: 0,
|
|
last_upstream_event: None,
|
|
last_upstream_frame_at_unix_ms: None,
|
|
partial_image_count: 0,
|
|
last_client_visible_event: None,
|
|
downstream_heartbeat_count: 0,
|
|
last_downstream_heartbeat_at_unix_ms: None,
|
|
downstream_heartbeat_interval_ms: None,
|
|
}
|
|
}
|
|
|
|
fn to_json(&self) -> Value {
|
|
json!({
|
|
"phase": self.phase,
|
|
"upstream_ttfb_ms": self.upstream_ttfb_ms,
|
|
"upstream_sse_frame_count": self.upstream_sse_frame_count,
|
|
"last_upstream_event": self.last_upstream_event,
|
|
"last_upstream_frame_at_unix_ms": self.last_upstream_frame_at_unix_ms,
|
|
"partial_image_count": self.partial_image_count,
|
|
"last_client_visible_event": self.last_client_visible_event,
|
|
"downstream_heartbeat_count": self.downstream_heartbeat_count,
|
|
"last_downstream_heartbeat_at_unix_ms": self.last_downstream_heartbeat_at_unix_ms,
|
|
"downstream_heartbeat_interval_ms": self.downstream_heartbeat_interval_ms,
|
|
})
|
|
}
|
|
}
|
|
|
|
impl<'a> OpenAiImageSyncProgressRecorder<'a> {
|
|
fn new(
|
|
state: &'a AppState,
|
|
plan: &'a ExecutionPlan,
|
|
report_context: Option<&'a Value>,
|
|
snapshot: Option<Arc<Mutex<OpenAiImageSyncProgressSnapshot>>>,
|
|
) -> Self {
|
|
Self {
|
|
state,
|
|
plan,
|
|
report_context,
|
|
snapshot: snapshot
|
|
.unwrap_or_else(|| Arc::new(Mutex::new(OpenAiImageSyncProgressSnapshot::new()))),
|
|
buffer: Vec::new(),
|
|
last_persist_at: None,
|
|
}
|
|
}
|
|
|
|
async fn persist(
|
|
&mut self,
|
|
status: RequestCandidateStatus,
|
|
status_code: Option<u16>,
|
|
latency_ms: Option<u64>,
|
|
force: bool,
|
|
) {
|
|
let now = Instant::now();
|
|
if !force
|
|
&& self.last_persist_at.is_some_and(|last| {
|
|
now.duration_since(last) < OPENAI_IMAGE_SYNC_PROGRESS_WRITE_INTERVAL
|
|
})
|
|
{
|
|
return;
|
|
}
|
|
let snapshot = self.snapshot.lock().await.clone();
|
|
let extra_data = json!({
|
|
"image_progress": snapshot.to_json(),
|
|
});
|
|
record_local_request_candidate_extra_data(
|
|
self.state,
|
|
self.plan,
|
|
self.report_context,
|
|
status,
|
|
status_code,
|
|
latency_ms,
|
|
extra_data,
|
|
)
|
|
.await;
|
|
self.last_persist_at = Some(now);
|
|
}
|
|
|
|
async fn record_connecting(&mut self) {
|
|
self.snapshot.lock().await.phase = "upstream_connecting";
|
|
self.persist(RequestCandidateStatus::Pending, None, None, true)
|
|
.await;
|
|
}
|
|
|
|
async fn record_response_started(&mut self, status_code: u16, ttfb_ms: u64) {
|
|
{
|
|
let mut snapshot = self.snapshot.lock().await;
|
|
snapshot.phase = if status_code >= 400 {
|
|
"failed"
|
|
} else {
|
|
"upstream_streaming"
|
|
};
|
|
snapshot.upstream_ttfb_ms = Some(ttfb_ms);
|
|
}
|
|
self.persist(
|
|
if status_code >= 400 {
|
|
RequestCandidateStatus::Failed
|
|
} else {
|
|
RequestCandidateStatus::Streaming
|
|
},
|
|
Some(status_code),
|
|
Some(ttfb_ms),
|
|
true,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
async fn observe_chunk(&mut self, chunk: &[u8], status_code: u16, elapsed_ms: u64) {
|
|
if chunk.is_empty() {
|
|
return;
|
|
}
|
|
let mut blocks = Vec::new();
|
|
let mut remaining = chunk;
|
|
let mut parser_overflowed = false;
|
|
loop {
|
|
while let Some(block_end) = find_sse_block_end(&self.buffer) {
|
|
blocks.push(self.buffer.drain(..block_end).collect::<Vec<_>>());
|
|
}
|
|
if remaining.is_empty() {
|
|
break;
|
|
}
|
|
|
|
let capacity =
|
|
OPENAI_IMAGE_SYNC_PROGRESS_MAX_BUFFER_BYTES.saturating_sub(self.buffer.len());
|
|
if capacity == 0 {
|
|
// The progress recorder is observational. Drop an incomplete
|
|
// oversized record and keep the client-facing response alive.
|
|
self.buffer.clear();
|
|
parser_overflowed = true;
|
|
break;
|
|
}
|
|
let take = find_sse_block_end(remaining)
|
|
.map_or(remaining.len(), |block_end| block_end)
|
|
.min(capacity);
|
|
if take == 0 {
|
|
self.buffer.clear();
|
|
parser_overflowed = true;
|
|
break;
|
|
}
|
|
self.buffer.extend_from_slice(&remaining[..take]);
|
|
remaining = &remaining[take..];
|
|
}
|
|
if parser_overflowed {
|
|
debug!("openai image sync progress parser dropped an oversized incomplete SSE record");
|
|
}
|
|
let mut force_persist = false;
|
|
for block in blocks {
|
|
let Some(frame) = parse_openai_image_sync_sse_frame(&block) else {
|
|
continue;
|
|
};
|
|
{
|
|
let mut snapshot = self.snapshot.lock().await;
|
|
snapshot.upstream_sse_frame_count =
|
|
snapshot.upstream_sse_frame_count.saturating_add(1);
|
|
snapshot.last_upstream_event = Some(frame.event_name);
|
|
snapshot.last_upstream_frame_at_unix_ms = Some(current_request_candidate_unix_ms());
|
|
if frame.is_partial_image {
|
|
snapshot.partial_image_count = snapshot.partial_image_count.saturating_add(1);
|
|
}
|
|
if let Some(client_visible_event) = frame.client_visible_event {
|
|
snapshot.last_client_visible_event = Some(client_visible_event.to_string());
|
|
force_persist = true;
|
|
}
|
|
if frame.is_failed || status_code >= 400 {
|
|
snapshot.phase = "failed";
|
|
force_persist = true;
|
|
} else if frame.is_completed {
|
|
snapshot.phase = "upstream_completed";
|
|
force_persist = true;
|
|
} else {
|
|
snapshot.phase = "upstream_streaming";
|
|
}
|
|
}
|
|
}
|
|
let phase = self.snapshot.lock().await.phase;
|
|
self.persist(
|
|
if phase == "failed" {
|
|
RequestCandidateStatus::Failed
|
|
} else {
|
|
RequestCandidateStatus::Streaming
|
|
},
|
|
Some(status_code),
|
|
Some(elapsed_ms),
|
|
force_persist,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
async fn finish(&mut self, status_code: u16, elapsed_ms: u64) {
|
|
{
|
|
let mut snapshot = self.snapshot.lock().await;
|
|
if status_code >= 400 || snapshot.phase == "failed" {
|
|
snapshot.phase = "failed";
|
|
} else {
|
|
snapshot.phase = "upstream_completed";
|
|
}
|
|
}
|
|
self.persist(
|
|
if status_code >= 400 {
|
|
RequestCandidateStatus::Failed
|
|
} else {
|
|
RequestCandidateStatus::Streaming
|
|
},
|
|
Some(status_code),
|
|
Some(elapsed_ms),
|
|
true,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
async fn fail(&mut self, status_code: Option<u16>, elapsed_ms: u64) {
|
|
self.snapshot.lock().await.phase = "failed";
|
|
self.persist(
|
|
RequestCandidateStatus::Failed,
|
|
status_code,
|
|
Some(elapsed_ms),
|
|
true,
|
|
)
|
|
.await;
|
|
}
|
|
}
|
|
|
|
impl OpenAiImageSyncJsonHeartbeatContext {
|
|
async fn record_heartbeat(&self, heartbeat_kind: &'static str, heartbeat_interval: Duration) {
|
|
let now_unix_ms = current_request_candidate_unix_ms();
|
|
let elapsed_ms = self.started_at.elapsed().as_millis() as u64;
|
|
let interval_ms = heartbeat_interval
|
|
.as_millis()
|
|
.try_into()
|
|
.unwrap_or(u64::MAX);
|
|
let (count, phase, progress_json) = {
|
|
let mut snapshot = self.snapshot.lock().await;
|
|
snapshot.downstream_heartbeat_count =
|
|
snapshot.downstream_heartbeat_count.saturating_add(1);
|
|
snapshot.last_downstream_heartbeat_at_unix_ms = Some(now_unix_ms);
|
|
snapshot.downstream_heartbeat_interval_ms = Some(interval_ms);
|
|
(
|
|
snapshot.downstream_heartbeat_count,
|
|
snapshot.phase,
|
|
snapshot.to_json(),
|
|
)
|
|
};
|
|
let status = match phase {
|
|
"failed" => RequestCandidateStatus::Failed,
|
|
"upstream_connecting" => RequestCandidateStatus::Pending,
|
|
_ => RequestCandidateStatus::Streaming,
|
|
};
|
|
record_local_request_candidate_extra_data(
|
|
&self.state,
|
|
&self.plan,
|
|
self.report_context.as_ref(),
|
|
status,
|
|
None,
|
|
Some(elapsed_ms),
|
|
json!({ "image_progress": progress_json }),
|
|
)
|
|
.await;
|
|
debug!(
|
|
event_name = "openai_image_sync_json_heartbeat_sent",
|
|
log_type = "event",
|
|
trace_id = %self.trace_id,
|
|
request_id = %self.request_id_for_log,
|
|
candidate_id = self.candidate_id.as_deref().unwrap_or("-"),
|
|
heartbeat_kind,
|
|
heartbeat_count = count,
|
|
heartbeat_interval_ms = interval_ms,
|
|
elapsed_ms,
|
|
phase,
|
|
"gateway emitted OpenAI image sync JSON whitespace heartbeat"
|
|
);
|
|
}
|
|
}
|
|
|
|
fn find_sse_block_end(buffer: &[u8]) -> Option<usize> {
|
|
buffer
|
|
.windows(2)
|
|
.position(|window| window == b"\n\n")
|
|
.map(|index| index + 2)
|
|
.or_else(|| {
|
|
buffer
|
|
.windows(4)
|
|
.position(|window| window == b"\r\n\r\n")
|
|
.map(|index| index + 4)
|
|
})
|
|
}
|
|
|
|
fn parse_openai_image_sync_sse_frame(block: &[u8]) -> Option<OpenAiImageSyncSseFrame> {
|
|
let text = std::str::from_utf8(block).ok()?.trim();
|
|
if text.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
let mut event_name = None;
|
|
let mut data_lines = Vec::new();
|
|
for line in text.lines() {
|
|
let line = line.trim_end_matches('\r').trim();
|
|
if line.is_empty() || line.starts_with(':') {
|
|
continue;
|
|
}
|
|
if let Some(value) = line.strip_prefix("event:") {
|
|
event_name = Some(value.trim().to_string());
|
|
} else if let Some(value) = line.strip_prefix("data:") {
|
|
data_lines.push(value.trim().to_string());
|
|
}
|
|
}
|
|
|
|
let data_text = data_lines.join("\n");
|
|
if data_text.trim().eq("[DONE]") {
|
|
let event_name = event_name.unwrap_or_else(|| "done".to_string());
|
|
return Some(OpenAiImageSyncSseFrame {
|
|
event_name,
|
|
is_partial_image: false,
|
|
is_completed: true,
|
|
is_failed: false,
|
|
client_visible_event: None,
|
|
});
|
|
}
|
|
|
|
let data_event_name = serde_json::from_str::<Value>(&data_text)
|
|
.ok()
|
|
.and_then(|value| {
|
|
value
|
|
.get("type")
|
|
.and_then(Value::as_str)
|
|
.map(str::to_string)
|
|
.or_else(|| {
|
|
value
|
|
.get("error")
|
|
.and_then(Value::as_object)
|
|
.map(|_| "error".to_string())
|
|
})
|
|
});
|
|
let event_name = event_name.or(data_event_name)?;
|
|
let is_partial_image = event_name == "response.image_generation_call.partial_image";
|
|
let is_completed = event_name == "response.completed";
|
|
let is_failed = event_name == "response.failed"
|
|
|| event_name == "response.error"
|
|
|| event_name == "error"
|
|
|| event_name.ends_with(".failed");
|
|
let client_visible_event = if is_partial_image {
|
|
Some("image_generation.partial_image")
|
|
} else if is_completed {
|
|
Some("image_generation.completed")
|
|
} else if is_failed {
|
|
Some("image_generation.failed")
|
|
} else {
|
|
None
|
|
};
|
|
|
|
Some(OpenAiImageSyncSseFrame {
|
|
event_name,
|
|
is_partial_image,
|
|
is_completed,
|
|
is_failed,
|
|
client_visible_event,
|
|
})
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
async fn execute_direct_sync_runtime_candidate(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&Value>,
|
|
trace_id: &str,
|
|
plan_kind: &str,
|
|
candidate_started_unix_ms: u64,
|
|
request_id_for_log: &str,
|
|
candidate_id: Option<&str>,
|
|
provider_name: &str,
|
|
endpoint_id: &str,
|
|
key_id: &str,
|
|
model_name: &str,
|
|
candidate_index: &str,
|
|
progress_snapshot: Option<Arc<Mutex<OpenAiImageSyncProgressSnapshot>>>,
|
|
) -> Result<ExecutionResult, SyncExecutionFailure> {
|
|
if let Some(result) = maybe_execute_windsurf_sync(state, plan, report_context)
|
|
.await
|
|
.map_err(SyncExecutionFailure::from_transport)?
|
|
{
|
|
return Ok(result);
|
|
}
|
|
if !should_track_openai_image_sync_upstream_sse(plan_kind, plan, report_context) {
|
|
let state_for_response_started = state.clone();
|
|
let response_started_lifecycle_seed = build_lifecycle_usage_seed(plan, report_context);
|
|
let response_started_candidate_snapshot =
|
|
snapshot_local_request_candidate_status(plan, report_context);
|
|
return DirectSyncExecutionRuntime::new()
|
|
.execute_sync_with_response_started(plan, move |event| {
|
|
record_sync_response_started(
|
|
&state_for_response_started,
|
|
response_started_lifecycle_seed,
|
|
response_started_candidate_snapshot,
|
|
candidate_started_unix_ms,
|
|
event.status_code,
|
|
event.ttfb_ms,
|
|
);
|
|
spawn_local_oauth_success_effect(
|
|
state_for_response_started.clone(),
|
|
plan,
|
|
report_context,
|
|
LocalOAuthSuccessEffect {
|
|
status_code: event.status_code,
|
|
request_started_at_unix_ms: Some(
|
|
event.response_observation.request_started_at_unix_ms,
|
|
),
|
|
request_order_id: Some(&event.response_observation.request_order_id),
|
|
},
|
|
);
|
|
})
|
|
.await
|
|
.map_err(SyncExecutionFailure::from_transport);
|
|
}
|
|
|
|
let started_at = Instant::now();
|
|
let timeout_ms = resolve_openai_image_sync_total_timeout_ms(plan);
|
|
let mut execution = Box::pin(execute_openai_image_sync_upstream_sse_candidate(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
progress_snapshot.clone(),
|
|
));
|
|
let mut idle_interval = tokio::time::interval(SYNC_EXECUTION_IDLE_LOG_INTERVAL);
|
|
idle_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
|
idle_interval.tick().await;
|
|
let mut total_timeout = Box::pin(tokio::time::sleep(Duration::from_millis(timeout_ms)));
|
|
|
|
loop {
|
|
tokio::select! {
|
|
result = execution.as_mut() => {
|
|
match result {
|
|
Ok(result) => return Ok(result),
|
|
Err(err) => {
|
|
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
|
let status_code = err.status_code;
|
|
record_openai_image_sync_failed_progress(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
status_code,
|
|
elapsed_ms,
|
|
progress_snapshot.clone(),
|
|
)
|
|
.await;
|
|
return Err(err);
|
|
}
|
|
}
|
|
}
|
|
_ = idle_interval.tick() => {
|
|
warn!(
|
|
event_name = "openai_image_sync_execution_idle",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %request_id_for_log,
|
|
candidate_id = candidate_id.unwrap_or("-"),
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index,
|
|
elapsed_ms = started_at.elapsed().as_millis() as u64,
|
|
timeout_ms,
|
|
"gateway OpenAI image sync execution still waiting for upstream response"
|
|
);
|
|
}
|
|
_ = total_timeout.as_mut() => {
|
|
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
|
record_openai_image_sync_failed_progress(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
Some(StatusCode::GATEWAY_TIMEOUT.as_u16()),
|
|
elapsed_ms,
|
|
progress_snapshot.clone(),
|
|
)
|
|
.await;
|
|
warn!(
|
|
event_name = "openai_image_sync_total_timeout",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %request_id_for_log,
|
|
candidate_id = candidate_id.unwrap_or("-"),
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index,
|
|
elapsed_ms,
|
|
timeout_ms,
|
|
"gateway OpenAI image sync execution exceeded total timeout"
|
|
);
|
|
return Err(SyncExecutionFailure::image_sync_total_timeout(
|
|
timeout_ms,
|
|
elapsed_ms,
|
|
));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn execute_openai_image_sync_upstream_sse_candidate(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&Value>,
|
|
progress_snapshot: Option<Arc<Mutex<OpenAiImageSyncProgressSnapshot>>>,
|
|
) -> Result<ExecutionResult, SyncExecutionFailure> {
|
|
let request_body = build_request_body(plan).map_err(SyncExecutionFailure::from_transport)?;
|
|
let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan);
|
|
let started_at = Instant::now();
|
|
let mut progress =
|
|
OpenAiImageSyncProgressRecorder::new(state, plan, report_context, progress_snapshot);
|
|
progress.record_connecting().await;
|
|
|
|
let request_started_at_unix_ms = current_request_candidate_unix_ms();
|
|
let request_order_id = uuid::Uuid::now_v7().to_string();
|
|
let response = send_request(plan, request_body)
|
|
.await
|
|
.map_err(SyncExecutionFailure::from_transport)?;
|
|
let ttfb_ms = started_at.elapsed().as_millis() as u64;
|
|
let response_headers_observed_at_unix_ms = current_request_candidate_unix_ms();
|
|
let status_code = response.status_code();
|
|
let headers = response.headers();
|
|
spawn_local_oauth_success_effect(
|
|
state.clone(),
|
|
plan,
|
|
report_context,
|
|
LocalOAuthSuccessEffect {
|
|
status_code,
|
|
request_started_at_unix_ms: Some(request_started_at_unix_ms),
|
|
request_order_id: Some(&request_order_id),
|
|
},
|
|
);
|
|
progress.record_response_started(status_code, ttfb_ms).await;
|
|
|
|
let mut body_bytes = Vec::new();
|
|
match response {
|
|
DirectHttpResponse::Reqwest(response) => {
|
|
let mut upstream_stream = response.bytes_stream();
|
|
while let Some(chunk) = upstream_stream.next().await {
|
|
let chunk = chunk.map_err(|err| {
|
|
SyncExecutionFailure::from_transport(
|
|
ExecutionRuntimeTransportError::UpstreamRequest(
|
|
format_upstream_request_error(&err),
|
|
),
|
|
)
|
|
})?;
|
|
append_upstream_response_body_chunk_with_limit(
|
|
&mut body_bytes,
|
|
&chunk,
|
|
response_body_limit_bytes,
|
|
)
|
|
.map_err(SyncExecutionFailure::from_transport)?;
|
|
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
|
progress
|
|
.observe_chunk(&chunk, status_code, elapsed_ms)
|
|
.await;
|
|
}
|
|
}
|
|
DirectHttpResponse::HyperH2c(response) => {
|
|
let mut upstream_stream = response.into_body().into_data_stream();
|
|
while let Some(chunk) = upstream_stream.next().await {
|
|
let chunk = chunk.map_err(|err| {
|
|
SyncExecutionFailure::from_transport(
|
|
ExecutionRuntimeTransportError::UpstreamRequest(format_hyper_error_chain(
|
|
&err,
|
|
)),
|
|
)
|
|
})?;
|
|
append_upstream_response_body_chunk_with_limit(
|
|
&mut body_bytes,
|
|
&chunk,
|
|
response_body_limit_bytes,
|
|
)
|
|
.map_err(SyncExecutionFailure::from_transport)?;
|
|
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
|
progress
|
|
.observe_chunk(&chunk, status_code, elapsed_ms)
|
|
.await;
|
|
}
|
|
}
|
|
DirectHttpResponse::BrowserWreq(response) => {
|
|
let mut upstream_stream = response.bytes_stream();
|
|
while let Some(chunk) = upstream_stream.next().await {
|
|
let chunk = chunk.map_err(|err| {
|
|
SyncExecutionFailure::from_transport(
|
|
ExecutionRuntimeTransportError::UpstreamRequest(
|
|
format_wreq_upstream_request_error(&err),
|
|
),
|
|
)
|
|
})?;
|
|
append_upstream_response_body_chunk_with_limit(
|
|
&mut body_bytes,
|
|
&chunk,
|
|
response_body_limit_bytes,
|
|
)
|
|
.map_err(SyncExecutionFailure::from_transport)?;
|
|
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
|
progress
|
|
.observe_chunk(&chunk, status_code, elapsed_ms)
|
|
.await;
|
|
}
|
|
}
|
|
}
|
|
|
|
let decoded_body_bytes =
|
|
decode_response_body_bytes_with_limit(&headers, &body_bytes, response_body_limit_bytes)
|
|
.map_err(SyncExecutionFailure::from_transport)?;
|
|
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
|
let upstream_bytes = body_bytes.len() as u64;
|
|
progress.finish(status_code, elapsed_ms).await;
|
|
|
|
let body = build_execution_response_body(
|
|
&headers,
|
|
&body_bytes,
|
|
decoded_body_bytes.as_ref(),
|
|
plan.stream,
|
|
execution_response_body_mode(plan),
|
|
)
|
|
.map_err(SyncExecutionFailure::from_transport)?;
|
|
|
|
Ok(ExecutionResult {
|
|
request_id: plan.request_id.clone(),
|
|
candidate_id: plan.candidate_id.clone(),
|
|
status_code,
|
|
headers,
|
|
response_observation: Some(ExecutionResponseObservation {
|
|
request_started_at_unix_ms,
|
|
response_headers_observed_at_unix_ms,
|
|
request_order_id,
|
|
}),
|
|
body,
|
|
telemetry: Some(ExecutionTelemetry {
|
|
ttfb_ms: Some(ttfb_ms),
|
|
elapsed_ms: Some(elapsed_ms),
|
|
upstream_bytes: Some(upstream_bytes),
|
|
}),
|
|
error: None,
|
|
})
|
|
}
|
|
|
|
async fn record_openai_image_sync_failed_progress(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&Value>,
|
|
status_code: Option<u16>,
|
|
elapsed_ms: u64,
|
|
progress_snapshot: Option<Arc<Mutex<OpenAiImageSyncProgressSnapshot>>>,
|
|
) {
|
|
let mut progress =
|
|
OpenAiImageSyncProgressRecorder::new(state, plan, report_context, progress_snapshot);
|
|
progress.fail(status_code, elapsed_ms).await;
|
|
}
|
|
|
|
fn resolve_openai_image_sync_total_timeout_ms(plan: &ExecutionPlan) -> u64 {
|
|
plan.timeouts
|
|
.as_ref()
|
|
.and_then(|timeouts| timeouts.total_ms)
|
|
.unwrap_or(OPENAI_IMAGE_SYNC_DEFAULT_TOTAL_TIMEOUT_MS)
|
|
.max(1)
|
|
}
|
|
|
|
fn report_context_upstream_is_stream(report_context: Option<&Value>) -> bool {
|
|
report_context
|
|
.and_then(|value| value.get(UPSTREAM_IS_STREAM_KEY))
|
|
.and_then(Value::as_bool)
|
|
.unwrap_or(false)
|
|
}
|
|
|
|
fn should_track_openai_image_sync_upstream_sse(
|
|
plan_kind: &str,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&Value>,
|
|
) -> bool {
|
|
plan_kind == OPENAI_IMAGE_SYNC_PLAN_KIND
|
|
&& (plan.stream || report_context_upstream_is_stream(report_context))
|
|
}
|
|
|
|
fn should_enable_openai_image_sync_json_heartbeat(
|
|
_plan_kind: &str,
|
|
_plan: &ExecutionPlan,
|
|
_report_context: Option<&Value>,
|
|
) -> bool {
|
|
false
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn build_openai_image_sync_json_heartbeat_response(
|
|
state: AppState,
|
|
request_path: String,
|
|
plan: ExecutionPlan,
|
|
trace_id: String,
|
|
decision: GatewayControlDecision,
|
|
plan_kind: String,
|
|
report_kind: Option<String>,
|
|
report_context: Option<Value>,
|
|
) -> Result<Response<Body>, GatewayError> {
|
|
let request_id = plan.request_id.clone();
|
|
let candidate_id = plan.candidate_id.clone();
|
|
let trace_id_for_response = trace_id.clone();
|
|
let decision_for_response = decision.clone();
|
|
let progress_snapshot = Arc::new(Mutex::new(OpenAiImageSyncProgressSnapshot::new()));
|
|
let heartbeat_context = OpenAiImageSyncJsonHeartbeatContext {
|
|
state: state.clone(),
|
|
plan: plan.clone(),
|
|
report_context: report_context.clone(),
|
|
snapshot: progress_snapshot.clone(),
|
|
started_at: Instant::now(),
|
|
trace_id: trace_id.clone(),
|
|
request_id_for_log: short_request_id(request_id.as_str()),
|
|
candidate_id: candidate_id.clone(),
|
|
};
|
|
let (tx, rx) = mpsc::channel::<Result<Bytes, IoError>>(1);
|
|
|
|
tokio::spawn(async move {
|
|
let bytes = openai_image_sync_json_heartbeat_final_bytes(
|
|
execute_execution_runtime_sync_impl(
|
|
&state,
|
|
request_path.as_str(),
|
|
plan,
|
|
trace_id.as_str(),
|
|
&decision,
|
|
plan_kind.as_str(),
|
|
report_kind,
|
|
report_context,
|
|
false,
|
|
Some(progress_snapshot),
|
|
None,
|
|
None,
|
|
)
|
|
.await,
|
|
)
|
|
.await;
|
|
let _ = tx.send(Ok(Bytes::from(bytes))).await;
|
|
});
|
|
|
|
let headers = BTreeMap::from([(
|
|
CONTENT_TYPE.as_str().to_string(),
|
|
"application/json".to_string(),
|
|
)]);
|
|
let response = build_client_response_from_parts_with_mutator(
|
|
StatusCode::OK.as_u16(),
|
|
&headers,
|
|
Body::from_stream(build_json_whitespace_heartbeat_stream(
|
|
rx,
|
|
OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_INTERVAL,
|
|
Some(heartbeat_context),
|
|
)),
|
|
trace_id_for_response.as_str(),
|
|
Some(&decision_for_response),
|
|
|headers| {
|
|
headers.remove(CONTENT_LENGTH);
|
|
headers.remove(CONTENT_ENCODING);
|
|
headers.insert(
|
|
CACHE_CONTROL,
|
|
HeaderValue::from_static("no-cache, no-transform"),
|
|
);
|
|
headers.insert(
|
|
HeaderName::from_static("x-accel-buffering"),
|
|
HeaderValue::from_static("no"),
|
|
);
|
|
Ok(())
|
|
},
|
|
)?;
|
|
attach_control_metadata_headers(response, Some(request_id.as_str()), candidate_id.as_deref())
|
|
}
|
|
|
|
fn build_json_whitespace_heartbeat_stream(
|
|
mut rx: mpsc::Receiver<Result<Bytes, IoError>>,
|
|
heartbeat_interval: Duration,
|
|
heartbeat_context: Option<OpenAiImageSyncJsonHeartbeatContext>,
|
|
) -> impl futures_util::Stream<Item = Result<Bytes, IoError>> + Send + 'static {
|
|
stream! {
|
|
if let Some(context) = heartbeat_context.as_ref() {
|
|
context.record_heartbeat("initial", heartbeat_interval).await;
|
|
}
|
|
yield Ok(Bytes::from_static(OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_BYTES));
|
|
|
|
let mut heartbeat = tokio::time::interval(heartbeat_interval);
|
|
heartbeat.set_missed_tick_behavior(MissedTickBehavior::Delay);
|
|
heartbeat.tick().await;
|
|
loop {
|
|
tokio::select! {
|
|
biased;
|
|
item = rx.recv() => {
|
|
let Some(item) = item else {
|
|
break;
|
|
};
|
|
yield item;
|
|
break;
|
|
}
|
|
_ = heartbeat.tick() => {
|
|
if let Some(context) = heartbeat_context.as_ref() {
|
|
context.record_heartbeat("interval", heartbeat_interval).await;
|
|
}
|
|
yield Ok(Bytes::from_static(OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_BYTES));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) fn build_sync_json_whitespace_heartbeat_stream(
|
|
rx: mpsc::Receiver<Result<Bytes, IoError>>,
|
|
) -> impl futures_util::Stream<Item = Result<Bytes, IoError>> + Send + 'static {
|
|
build_json_whitespace_heartbeat_stream(rx, OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_INTERVAL, None)
|
|
}
|
|
|
|
pub(crate) fn build_openai_image_sync_json_whitespace_heartbeat_stream(
|
|
rx: mpsc::Receiver<Result<Bytes, IoError>>,
|
|
) -> impl futures_util::Stream<Item = Result<Bytes, IoError>> + Send + 'static {
|
|
build_sync_json_whitespace_heartbeat_stream(rx)
|
|
}
|
|
|
|
async fn openai_image_sync_json_heartbeat_final_bytes(
|
|
result: Result<Option<Response<Body>>, GatewayError>,
|
|
) -> Vec<u8> {
|
|
match result {
|
|
Ok(Some(response)) => match to_bytes(
|
|
response.into_body(),
|
|
crate::headers::max_internal_buffered_body_bytes(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(bytes) if !bytes.is_empty() => bytes.to_vec(),
|
|
Ok(_) => openai_image_sync_json_heartbeat_error_body("empty sync image response"),
|
|
Err(_err) => {
|
|
openai_image_sync_json_heartbeat_error_body("sync image response read failed")
|
|
}
|
|
},
|
|
Ok(None) => openai_image_sync_json_heartbeat_error_body(
|
|
"sync image execution ended without a local response",
|
|
),
|
|
// Do not serialize the internal error: Debug output can contain
|
|
// upstream URLs, credentials, or other request-specific details.
|
|
Err(_err) => openai_image_sync_json_heartbeat_error_body("sync image execution failed"),
|
|
}
|
|
}
|
|
|
|
fn openai_image_sync_json_heartbeat_error_body(message: &str) -> Vec<u8> {
|
|
serde_json::to_vec(&json!({
|
|
"error": {
|
|
"type": "aether_gateway_error",
|
|
"message": message,
|
|
}
|
|
}))
|
|
.unwrap_or_else(|_| b"{\"error\":{\"type\":\"aether_gateway_error\"}}".to_vec())
|
|
}
|
|
|
|
async fn apply_sync_success_effects(
|
|
state: &AppState,
|
|
plan: &ExecutionPlan,
|
|
report_context: Option<&serde_json::Value>,
|
|
payload: &GatewaySyncReportRequest,
|
|
) {
|
|
if let Some(report_context) = report_context {
|
|
crate::ai_serving::persist_converted_response_history(
|
|
state,
|
|
report_context,
|
|
payload
|
|
.client_body_json
|
|
.as_ref()
|
|
.or(payload.body_json.as_ref()),
|
|
)
|
|
.await;
|
|
}
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan,
|
|
report_context,
|
|
},
|
|
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan,
|
|
report_context,
|
|
},
|
|
LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan,
|
|
report_context,
|
|
},
|
|
LocalExecutionEffect::PoolSuccessSync { payload },
|
|
)
|
|
.await;
|
|
}
|
|
|
|
#[cfg(test)]
|
|
enum RemoteSyncFallbackOutcome {
|
|
Executed(ExecutionResult),
|
|
ClientResponse(Response<Body>),
|
|
Unavailable,
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
|
|
pub(crate) async fn execute_execution_runtime_sync(
|
|
state: &AppState,
|
|
request_path: &str,
|
|
mut plan: ExecutionPlan,
|
|
trace_id: &str,
|
|
decision: &GatewayControlDecision,
|
|
plan_kind: &str,
|
|
report_kind: Option<String>,
|
|
mut report_context: Option<serde_json::Value>,
|
|
) -> Result<Option<Response<Body>>, GatewayError> {
|
|
execute_execution_runtime_sync_impl(
|
|
state,
|
|
request_path,
|
|
plan,
|
|
trace_id,
|
|
decision,
|
|
plan_kind,
|
|
report_kind,
|
|
report_context,
|
|
true,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.await
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub(crate) async fn execute_execution_runtime_sync_with_retry_scope(
|
|
state: &AppState,
|
|
request_path: &str,
|
|
plan: ExecutionPlan,
|
|
trace_id: &str,
|
|
decision: &GatewayControlDecision,
|
|
plan_kind: &str,
|
|
report_kind: Option<String>,
|
|
report_context: Option<serde_json::Value>,
|
|
) -> Result<AiAttemptExecutionOutcome<Response<Body>>, GatewayError> {
|
|
let mut retry_scope = AiAttemptRetryScope::Candidate;
|
|
let mut fallback_response = None;
|
|
let response = execute_execution_runtime_sync_impl(
|
|
state,
|
|
request_path,
|
|
plan,
|
|
trace_id,
|
|
decision,
|
|
plan_kind,
|
|
report_kind,
|
|
report_context,
|
|
true,
|
|
None,
|
|
Some(&mut retry_scope),
|
|
Some(&mut fallback_response),
|
|
)
|
|
.await?;
|
|
Ok(match response {
|
|
Some(response) => AiAttemptExecutionOutcome::Responded(response),
|
|
None => AiAttemptExecutionOutcome::Retry {
|
|
scope: retry_scope,
|
|
fallback_response,
|
|
},
|
|
})
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
|
|
async fn execute_execution_runtime_sync_impl(
|
|
state: &AppState,
|
|
request_path: &str,
|
|
mut plan: ExecutionPlan,
|
|
trace_id: &str,
|
|
decision: &GatewayControlDecision,
|
|
plan_kind: &str,
|
|
report_kind: Option<String>,
|
|
mut report_context: Option<serde_json::Value>,
|
|
allow_json_heartbeat: bool,
|
|
progress_snapshot: Option<Arc<Mutex<OpenAiImageSyncProgressSnapshot>>>,
|
|
mut retry_scope_out: Option<&mut AiAttemptRetryScope>,
|
|
mut retry_fallback_out: Option<&mut Option<Response<Body>>>,
|
|
) -> Result<Option<Response<Body>>, GatewayError> {
|
|
if allow_json_heartbeat
|
|
&& should_enable_openai_image_sync_json_heartbeat(plan_kind, &plan, report_context.as_ref())
|
|
{
|
|
return build_openai_image_sync_json_heartbeat_response(
|
|
state.clone(),
|
|
request_path.to_string(),
|
|
plan,
|
|
trace_id.to_string(),
|
|
decision.clone(),
|
|
plan_kind.to_string(),
|
|
report_kind,
|
|
report_context,
|
|
)
|
|
.map(Some);
|
|
}
|
|
|
|
ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await;
|
|
let plan_request_id = plan.request_id.clone();
|
|
let plan_request_id_for_log = short_request_id(plan_request_id.as_str());
|
|
let plan_candidate_id = plan.candidate_id.clone();
|
|
let provider_name = plan
|
|
.provider_name
|
|
.clone()
|
|
.unwrap_or_else(|| "-".to_string());
|
|
let endpoint_id = plan.endpoint_id.clone();
|
|
let key_id = plan.key_id.clone();
|
|
let model_name = plan.model_name.clone().unwrap_or_else(|| "-".to_string());
|
|
let candidate_index = parse_request_candidate_report_context(report_context.as_ref())
|
|
.and_then(|context| context.candidate_index)
|
|
.map(|value| value.to_string())
|
|
.unwrap_or_else(|| "-".to_string());
|
|
let candidate_started_at = Instant::now();
|
|
let candidate_started_unix_secs = current_request_candidate_unix_ms();
|
|
let _provider_pool_in_flight_guard =
|
|
match acquire_provider_pool_execution_guard(state, &plan, report_context.as_ref()).await? {
|
|
ProviderPoolInFlightAdmission::Acquired(guard) => guard,
|
|
ProviderPoolInFlightAdmission::Saturated { limit } => {
|
|
record_local_runtime_candidate_skip_reason(
|
|
state,
|
|
trace_id,
|
|
"provider_key_concurrency_limit_reached",
|
|
);
|
|
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
|
|
*retry_scope = AiAttemptRetryScope::Candidate;
|
|
}
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Skipped,
|
|
status_code: Some(StatusCode::TOO_MANY_REQUESTS.as_u16()),
|
|
error_type: Some("provider_key_concurrency_limit_reached".to_string()),
|
|
error_message: Some(format!(
|
|
"provider key concurrency limit reached: {limit}"
|
|
)),
|
|
latency_ms: Some(0),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(candidate_started_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
return Ok(None);
|
|
}
|
|
};
|
|
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
|
|
let usage_data = state.usage_lifecycle_data_state().as_ref().clone();
|
|
state
|
|
.usage_runtime
|
|
.record_pending_direct(&usage_data, lifecycle_seed)
|
|
.await;
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Pending,
|
|
status_code: None,
|
|
error_type: None,
|
|
error_message: None,
|
|
latency_ms: None,
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: None,
|
|
},
|
|
)
|
|
.await;
|
|
let mut terminal_guard = SyncAttemptTerminalGuard::new(
|
|
state,
|
|
&plan,
|
|
report_context.clone(),
|
|
candidate_started_unix_secs,
|
|
candidate_started_at,
|
|
);
|
|
let result = (async {
|
|
record_sync_execution_active(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
candidate_started_unix_secs,
|
|
);
|
|
#[cfg(not(test))]
|
|
let mut result = {
|
|
match maybe_execute_grok_sync(&plan, report_context.as_ref()).await {
|
|
Ok(Some(result)) => result,
|
|
Ok(None) => {
|
|
match maybe_execute_chatgpt_web_image_sync(state, &plan, report_context.as_ref())
|
|
.await
|
|
{
|
|
Ok(Some(result)) => result,
|
|
Ok(None) => match execute_direct_sync_runtime_candidate(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
plan_kind,
|
|
candidate_started_unix_secs,
|
|
plan_request_id_for_log.as_str(),
|
|
plan_candidate_id.as_deref(),
|
|
provider_name.as_str(),
|
|
endpoint_id.as_str(),
|
|
key_id.as_str(),
|
|
model_name.as_str(),
|
|
candidate_index.as_str(),
|
|
progress_snapshot.clone(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => result,
|
|
Err(err) => {
|
|
let failure_error_type = err.error_type;
|
|
let failure_message = err.message.clone();
|
|
let failure_latency_ms = err
|
|
.latency_ms
|
|
.unwrap_or_else(|| elapsed_ms_since(candidate_started_at));
|
|
maybe_store_sync_execution_failure_fallback(
|
|
&err,
|
|
&plan,
|
|
trace_id,
|
|
decision,
|
|
&mut retry_scope_out,
|
|
&mut retry_fallback_out,
|
|
)?;
|
|
warn!(
|
|
event_name = "sync_execution_runtime_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error_type = err.error_type,
|
|
error = %err.message,
|
|
"gateway in-process sync execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some(failure_error_type.to_string()),
|
|
error_message: Some(err.message),
|
|
latency_ms: Some(failure_latency_ms),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
failure_error_type,
|
|
failure_message.as_str(),
|
|
failure_latency_ms,
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
},
|
|
Err(err) => {
|
|
let transport_error_message = safe_transport_error_message(&err);
|
|
warn!(
|
|
event_name = "chatgpt_web_image_execution_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = %transport_error_message,
|
|
"gateway ChatGPT-Web image execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some(
|
|
"chatgpt_web_image_execution_unavailable".to_string(),
|
|
),
|
|
error_message: Some(transport_error_message.clone()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
"chatgpt_web_image_execution_unavailable",
|
|
transport_error_message.as_str(),
|
|
elapsed_ms_since(candidate_started_at),
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
}
|
|
}
|
|
Err(err) => {
|
|
let transport_error_message = safe_transport_error_message(&err);
|
|
warn!(
|
|
event_name = "grok_execution_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = %transport_error_message,
|
|
"gateway Grok execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some("grok_execution_unavailable".to_string()),
|
|
error_message: Some(transport_error_message.clone()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
"grok_execution_unavailable",
|
|
transport_error_message.as_str(),
|
|
elapsed_ms_since(candidate_started_at),
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
}
|
|
};
|
|
#[cfg(test)]
|
|
let mut result = {
|
|
if let Some(override_fn) = state.execution_runtime_sync_override.as_ref() {
|
|
match (override_fn.0)(&plan) {
|
|
Ok(result) => result,
|
|
Err(err) => {
|
|
let transport_error_message = persisted_sync_abort_message(&err).to_string();
|
|
warn!(
|
|
event_name = "sync_execution_runtime_test_override_failed",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = ?err,
|
|
"gateway test sync execution override failed"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some("execution_runtime_unavailable".to_string()),
|
|
error_message: Some(transport_error_message.clone()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
"execution_runtime_unavailable",
|
|
transport_error_message.as_str(),
|
|
elapsed_ms_since(candidate_started_at),
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
}
|
|
} else if state
|
|
.execution_runtime_override_base_url()
|
|
.unwrap_or_default()
|
|
.trim()
|
|
.is_empty()
|
|
{
|
|
match maybe_execute_grok_sync(&plan, report_context.as_ref()).await {
|
|
Ok(Some(result)) => result,
|
|
Ok(None) => match maybe_execute_chatgpt_web_image_sync(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(Some(result)) => result,
|
|
Ok(None) => match execute_direct_sync_runtime_candidate(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
plan_kind,
|
|
candidate_started_unix_secs,
|
|
plan_request_id_for_log.as_str(),
|
|
plan_candidate_id.as_deref(),
|
|
provider_name.as_str(),
|
|
endpoint_id.as_str(),
|
|
key_id.as_str(),
|
|
model_name.as_str(),
|
|
candidate_index.as_str(),
|
|
progress_snapshot.clone(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => result,
|
|
Err(err) => {
|
|
let failure_error_type = err.error_type;
|
|
let failure_message = err.message.clone();
|
|
let failure_latency_ms = err
|
|
.latency_ms
|
|
.unwrap_or_else(|| elapsed_ms_since(candidate_started_at));
|
|
maybe_store_sync_execution_failure_fallback(
|
|
&err,
|
|
&plan,
|
|
trace_id,
|
|
decision,
|
|
&mut retry_scope_out,
|
|
&mut retry_fallback_out,
|
|
)?;
|
|
warn!(
|
|
event_name = "sync_execution_runtime_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error_type = err.error_type,
|
|
error = %err.message,
|
|
"gateway in-process sync execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some(failure_error_type.to_string()),
|
|
error_message: Some(err.message),
|
|
latency_ms: Some(failure_latency_ms),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
failure_error_type,
|
|
failure_message.as_str(),
|
|
failure_latency_ms,
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
},
|
|
Err(err) => {
|
|
let transport_error_message = safe_transport_error_message(&err);
|
|
warn!(
|
|
event_name = "chatgpt_web_image_execution_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = %transport_error_message,
|
|
"gateway ChatGPT-Web image execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some(
|
|
"chatgpt_web_image_execution_unavailable".to_string(),
|
|
),
|
|
error_message: Some(transport_error_message.clone()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
"chatgpt_web_image_execution_unavailable",
|
|
transport_error_message.as_str(),
|
|
elapsed_ms_since(candidate_started_at),
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
},
|
|
Err(err) => {
|
|
let transport_error_message = safe_transport_error_message(&err);
|
|
warn!(
|
|
event_name = "grok_execution_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = %transport_error_message,
|
|
"gateway Grok execution unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some("grok_execution_unavailable".to_string()),
|
|
error_message: Some(transport_error_message.clone()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
if let Some(response) = maybe_build_sync_transport_error_stop_response(
|
|
state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
trace_id,
|
|
decision,
|
|
"grok_execution_unavailable",
|
|
transport_error_message.as_str(),
|
|
elapsed_ms_since(candidate_started_at),
|
|
)
|
|
.await?
|
|
{
|
|
return Ok(Some(response));
|
|
}
|
|
return Ok(None);
|
|
}
|
|
}
|
|
} else {
|
|
let remote_execution_runtime_base_url = state
|
|
.execution_runtime_override_base_url()
|
|
.unwrap_or_default();
|
|
let remote_outcome = execute_sync_via_remote_execution_runtime(
|
|
state,
|
|
remote_execution_runtime_base_url,
|
|
trace_id,
|
|
decision,
|
|
&plan,
|
|
plan_request_id.as_str(),
|
|
plan_candidate_id.as_deref(),
|
|
report_context.as_ref(),
|
|
candidate_started_unix_secs,
|
|
candidate_started_at,
|
|
)
|
|
.await?;
|
|
match remote_outcome {
|
|
RemoteSyncFallbackOutcome::Executed(result) => result,
|
|
RemoteSyncFallbackOutcome::ClientResponse(response) => return Ok(Some(response)),
|
|
RemoteSyncFallbackOutcome::Unavailable => return Ok(None),
|
|
}
|
|
}
|
|
};
|
|
let mut candidate_first_byte_elapsed_ms =
|
|
calibrated_sync_candidate_first_byte_elapsed_ms(candidate_started_at, &result);
|
|
let initial_response_observed_at_unix_ms = current_request_candidate_unix_ms();
|
|
let mut provider_response_observation =
|
|
result
|
|
.response_observation
|
|
.clone()
|
|
.unwrap_or(ExecutionResponseObservation {
|
|
request_started_at_unix_ms: candidate_started_unix_secs,
|
|
response_headers_observed_at_unix_ms: initial_response_observed_at_unix_ms,
|
|
request_order_id: uuid::Uuid::now_v7().to_string(),
|
|
});
|
|
let mut oauth_retry_attempted = false;
|
|
let (
|
|
result_error_type,
|
|
result_error_message,
|
|
result_latency_ms,
|
|
headers,
|
|
body_bytes,
|
|
body_json,
|
|
body_base64,
|
|
local_failover_response_text,
|
|
local_failover_analysis,
|
|
) = loop {
|
|
spawn_local_oauth_success_effect(
|
|
state.clone(),
|
|
&plan,
|
|
report_context.as_ref(),
|
|
LocalOAuthSuccessEffect {
|
|
status_code: result.status_code,
|
|
request_started_at_unix_ms: Some(
|
|
provider_response_observation.request_started_at_unix_ms,
|
|
),
|
|
request_order_id: Some(&provider_response_observation.request_order_id),
|
|
},
|
|
);
|
|
let result_latency_ms = result
|
|
.telemetry
|
|
.as_ref()
|
|
.and_then(|telemetry| telemetry.elapsed_ms);
|
|
let mut headers = std::mem::take(&mut result.headers);
|
|
let result_body = result.body.take();
|
|
let chatgpt_web_image_result = plan
|
|
.provider_api_format
|
|
.eq_ignore_ascii_case("openai:image")
|
|
&& (plan.headers.iter().any(|(name, value)| {
|
|
name.eq_ignore_ascii_case("x-aether-chatgpt-web-image") && value == "1"
|
|
}) || report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("chatgpt_web_image"))
|
|
.and_then(Value::as_bool)
|
|
.unwrap_or(false));
|
|
let (body_bytes, mut body_json, body_base64) = if chatgpt_web_image_result {
|
|
decode_execution_result_body_with_limit(
|
|
result_body,
|
|
&mut headers,
|
|
crate::execution_runtime::chatgpt_web_image::chatgpt_web_image_sse_envelope_limit_bytes(),
|
|
)?
|
|
} else {
|
|
decode_execution_result_body(result_body, &mut headers)?
|
|
};
|
|
if let Some(message) = invalid_gemini_provider_success_message(
|
|
&plan,
|
|
report_context.as_ref(),
|
|
result.status_code,
|
|
body_json.as_ref(),
|
|
)
|
|
.or_else(|| {
|
|
invalid_gemini_provider_stream_success_message(
|
|
&plan,
|
|
report_context.as_ref(),
|
|
result.status_code,
|
|
body_json.as_ref(),
|
|
&body_bytes,
|
|
body_base64.is_some(),
|
|
)
|
|
}) {
|
|
result.status_code = StatusCode::BAD_GATEWAY.as_u16();
|
|
result.error = Some(invalid_gemini_provider_success_execution_error(message));
|
|
if let Some(error_body) =
|
|
build_invalid_provider_success_body(&plan, report_context.as_ref(), message)
|
|
{
|
|
body_json = Some(error_body);
|
|
headers.insert("content-type".to_string(), "application/json".to_string());
|
|
}
|
|
}
|
|
let (mut result_error_type, mut result_error_message) =
|
|
execution_error_details(result.error.as_ref(), body_json.as_ref());
|
|
if result.status_code < 400 && body_json.is_none() {
|
|
if let Some(error_body_json) =
|
|
extract_provider_private_stream_error_body(report_context.as_ref(), &body_bytes)
|
|
{
|
|
result.status_code =
|
|
resolve_local_sync_error_status_code(result.status_code, &error_body_json);
|
|
let (private_error_type, private_error_message) =
|
|
provider_private_error_details(&error_body_json);
|
|
result_error_type = private_error_type.or(result_error_type);
|
|
result_error_message = private_error_message.or(result_error_message);
|
|
body_json = Some(error_body_json);
|
|
}
|
|
}
|
|
let local_failover_response_text = local_failover_response_text(
|
|
body_json.as_ref(),
|
|
&body_bytes,
|
|
result.error.as_ref().map(|error| error.message.as_str()),
|
|
);
|
|
|
|
if result.status_code >= 400
|
|
&& !oauth_retry_attempted
|
|
&& refresh_oauth_plan_auth_for_retry(
|
|
state,
|
|
&mut plan,
|
|
result.status_code,
|
|
local_failover_response_text.as_deref(),
|
|
trace_id,
|
|
report_context.as_ref(),
|
|
Some(provider_response_observation.request_started_at_unix_ms),
|
|
Some(&provider_response_observation.request_order_id),
|
|
)
|
|
.await
|
|
{
|
|
oauth_retry_attempted = true;
|
|
let retry_started_at_unix_ms = current_request_candidate_unix_ms();
|
|
let retry_request_order_id = uuid::Uuid::now_v7().to_string();
|
|
match crate::execution_runtime::execute_execution_runtime_sync_plan(
|
|
state,
|
|
Some(trace_id),
|
|
&plan,
|
|
)
|
|
.await
|
|
{
|
|
Ok(retry_result) => {
|
|
let retry_response_observed_at_unix_ms = current_request_candidate_unix_ms();
|
|
provider_response_observation = retry_result
|
|
.response_observation
|
|
.clone()
|
|
.unwrap_or(ExecutionResponseObservation {
|
|
request_started_at_unix_ms: retry_started_at_unix_ms,
|
|
response_headers_observed_at_unix_ms:
|
|
retry_response_observed_at_unix_ms,
|
|
request_order_id: retry_request_order_id,
|
|
});
|
|
candidate_first_byte_elapsed_ms =
|
|
calibrated_sync_candidate_first_byte_elapsed_ms(
|
|
candidate_started_at,
|
|
&retry_result,
|
|
);
|
|
result = retry_result;
|
|
continue;
|
|
}
|
|
Err(err) => {
|
|
warn!(
|
|
event_name = "local_sync_oauth_retry_execution_failed",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
candidate_id = ?plan_candidate_id,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
error = ?err,
|
|
"gateway oauth retry sync execution failed"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
let local_failover_analysis = analyze_local_candidate_failover_sync(
|
|
state,
|
|
&plan,
|
|
plan_kind,
|
|
report_context.as_ref(),
|
|
&result,
|
|
local_failover_response_text.as_deref(),
|
|
)
|
|
.await;
|
|
break (
|
|
result_error_type,
|
|
result_error_message,
|
|
result_latency_ms,
|
|
headers,
|
|
body_bytes,
|
|
body_json,
|
|
body_base64,
|
|
local_failover_response_text,
|
|
local_failover_analysis,
|
|
);
|
|
};
|
|
let mut report_context = attach_provider_response_headers_to_report_context(
|
|
report_context,
|
|
&headers,
|
|
provider_response_observation.request_started_at_unix_ms,
|
|
provider_response_observation.response_headers_observed_at_unix_ms,
|
|
&provider_response_observation.request_order_id,
|
|
);
|
|
if let Some(error) = result.error.as_ref() {
|
|
report_context = crate::usage::reporting::failure::execution_error_analytics_context(report_context.as_ref(), error);
|
|
}
|
|
if result.status_code >= 400 {
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan: &plan,
|
|
report_context: report_context.as_ref(),
|
|
},
|
|
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
|
|
status_code: result.status_code,
|
|
classification: local_failover_analysis.classification,
|
|
}),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan: &plan,
|
|
report_context: report_context.as_ref(),
|
|
},
|
|
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
|
|
status_code: result.status_code,
|
|
classification: local_failover_analysis.classification,
|
|
headers: Some(&headers),
|
|
}),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan: &plan,
|
|
report_context: report_context.as_ref(),
|
|
},
|
|
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
|
|
status_code: result.status_code,
|
|
classification: local_failover_analysis.classification,
|
|
}),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan: &plan,
|
|
report_context: report_context.as_ref(),
|
|
},
|
|
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
|
|
status_code: result.status_code,
|
|
response_text: local_failover_response_text.as_deref(),
|
|
}),
|
|
)
|
|
.await;
|
|
apply_local_execution_effect(
|
|
state,
|
|
LocalExecutionEffectContext {
|
|
plan: &plan,
|
|
report_context: report_context.as_ref(),
|
|
},
|
|
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
|
|
status_code: result.status_code,
|
|
classification: local_failover_analysis.classification,
|
|
headers: &headers,
|
|
error_body: local_failover_response_text.as_deref(),
|
|
}),
|
|
)
|
|
.await;
|
|
}
|
|
if matches!(
|
|
local_failover_analysis.decision,
|
|
LocalFailoverDecision::RetryNextCandidate
|
|
) {
|
|
let failure_disposition = crate::orchestration::classify_failure_disposition(
|
|
&plan.provider_api_format,
|
|
local_failover_analysis.classification,
|
|
result.status_code,
|
|
);
|
|
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
|
|
*retry_scope =
|
|
ai_attempt_retry_scope_from_failure_disposition(failure_disposition);
|
|
}
|
|
if failure_disposition.preserve_upstream_error {
|
|
if let Some(retry_fallback) = retry_fallback_out.as_deref_mut() {
|
|
let mut fallback_headers = headers.clone();
|
|
apply_endpoint_response_header_rules(
|
|
state,
|
|
&plan,
|
|
&mut fallback_headers,
|
|
body_json.as_ref(),
|
|
)
|
|
.await?;
|
|
*retry_fallback = Some(attach_control_metadata_headers(
|
|
build_client_response_from_parts(
|
|
result.status_code,
|
|
&fallback_headers,
|
|
Body::from(body_bytes.clone()),
|
|
trace_id,
|
|
Some(decision),
|
|
)?,
|
|
Some(plan.request_id.as_str()),
|
|
plan.candidate_id.as_deref(),
|
|
)?);
|
|
}
|
|
}
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
let error_trace_report_context = with_sync_error_trace_context(
|
|
report_context.as_ref(),
|
|
result.status_code,
|
|
&headers,
|
|
body_json.as_ref(),
|
|
&body_bytes,
|
|
local_failover_response_text.as_deref(),
|
|
local_failover_analysis,
|
|
);
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
error_trace_report_context
|
|
.as_ref()
|
|
.or(report_context.as_ref()),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: Some(result.status_code),
|
|
error_type: result_error_type.clone(),
|
|
error_message: result_error_message.clone(),
|
|
latency_ms: result_latency_ms,
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
warn!(
|
|
event_name = "local_sync_candidate_retry_scheduled",
|
|
log_type = "event",
|
|
trace_id = %trace_id,
|
|
request_id = %plan_request_id_for_log,
|
|
status_code = result.status_code,
|
|
provider_name,
|
|
endpoint_id,
|
|
key_id,
|
|
model_name,
|
|
candidate_index = candidate_index.as_str(),
|
|
"gateway local sync decision retrying next candidate after retryable execution runtime result"
|
|
);
|
|
return Ok(None);
|
|
}
|
|
let status_code = result.status_code;
|
|
let has_body_bytes = body_base64.is_some();
|
|
if (200..300).contains(&status_code) {
|
|
seed_kiro_sync_simulated_cache_enabled(state, &plan, &mut report_context).await;
|
|
if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) {
|
|
seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context);
|
|
}
|
|
seed_kiro_sync_report_context_prompt_cache_usage(state, &plan, &mut report_context).await;
|
|
}
|
|
let mut client_headers = headers.clone();
|
|
apply_endpoint_response_header_rules(state, &plan, &mut client_headers, body_json.as_ref())
|
|
.await?;
|
|
let explicit_finalize = should_finalize_sync_response(report_kind.as_deref());
|
|
let mapped_error_finalize_kind =
|
|
resolve_core_sync_error_finalize_report_kind(plan_kind, &result, body_json.as_ref());
|
|
let implicit_finalize = if !explicit_finalize && mapped_error_finalize_kind.is_none() {
|
|
maybe_build_implicit_sync_finalize_outcome(
|
|
trace_id,
|
|
decision,
|
|
plan_kind,
|
|
&report_context,
|
|
status_code,
|
|
&client_headers,
|
|
&body_json,
|
|
&body_base64,
|
|
&result.telemetry,
|
|
)?
|
|
} else {
|
|
None
|
|
};
|
|
if !matches!(
|
|
local_failover_analysis.decision,
|
|
LocalFailoverDecision::StopLocalFailover
|
|
) && should_fallback_to_control_sync(
|
|
plan_kind,
|
|
&result,
|
|
body_json.as_ref(),
|
|
has_body_bytes,
|
|
explicit_finalize || implicit_finalize.is_some(),
|
|
mapped_error_finalize_kind.is_some(),
|
|
) {
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
let error_trace_report_context = with_sync_error_trace_context(
|
|
report_context.as_ref(),
|
|
result.status_code,
|
|
&headers,
|
|
body_json.as_ref(),
|
|
&body_bytes,
|
|
local_failover_response_text.as_deref(),
|
|
local_failover_analysis,
|
|
);
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
error_trace_report_context
|
|
.as_ref()
|
|
.or(report_context.as_ref()),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: Some(result.status_code),
|
|
error_type: result_error_type.clone(),
|
|
error_message: result_error_message.clone(),
|
|
latency_ms: result_latency_ms,
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
return Ok(None);
|
|
}
|
|
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
let error_flow_report_context = (result.status_code >= 400)
|
|
.then(|| {
|
|
with_sync_error_trace_context(
|
|
report_context.as_ref(),
|
|
result.status_code,
|
|
&headers,
|
|
body_json.as_ref(),
|
|
&body_bytes,
|
|
local_failover_response_text.as_deref(),
|
|
local_failover_analysis,
|
|
)
|
|
})
|
|
.flatten();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
&plan,
|
|
error_flow_report_context
|
|
.as_ref()
|
|
.or(report_context.as_ref()),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: if result.status_code >= 400 {
|
|
RequestCandidateStatus::Failed
|
|
} else {
|
|
RequestCandidateStatus::Success
|
|
},
|
|
status_code: Some(result.status_code),
|
|
error_type: result_error_type.clone(),
|
|
error_message: result_error_message.clone(),
|
|
latency_ms: result_latency_ms,
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
|
|
let request_id_owned = result.request_id;
|
|
let candidate_id_owned = result.candidate_id;
|
|
let request_id = (!request_id_owned.trim().is_empty())
|
|
.then_some(request_id_owned.as_str())
|
|
.or(Some(plan_request_id.as_str()));
|
|
let request_id_for_log = short_request_id(request_id.unwrap_or("-"));
|
|
let candidate_id = candidate_id_owned
|
|
.as_deref()
|
|
.or(plan_candidate_id.as_deref());
|
|
let report_context = report_context;
|
|
let body_json = body_json;
|
|
let telemetry = result.telemetry;
|
|
|
|
if let Some(implicit_finalize) = implicit_finalize {
|
|
let usage_payload = implicit_finalize
|
|
.outcome
|
|
.background_report
|
|
.as_ref()
|
|
.unwrap_or(&implicit_finalize.payload);
|
|
apply_sync_success_effects(
|
|
state,
|
|
&plan,
|
|
implicit_finalize.payload.report_context.as_ref(),
|
|
usage_payload,
|
|
)
|
|
.await;
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
implicit_finalize.payload.report_context.as_ref(),
|
|
usage_payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
if let Some(report_payload) = implicit_finalize.outcome.background_report {
|
|
spawn_sync_report(state.clone(), report_payload);
|
|
} else {
|
|
warn!(
|
|
event_name = "local_core_finalize_missing_success_report_mapping",
|
|
log_type = "event",
|
|
trace_id = %trace_id,
|
|
report_kind = %implicit_finalize.payload.report_kind,
|
|
"gateway implicit local core finalize produced response without background success report mapping"
|
|
);
|
|
}
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
implicit_finalize.outcome.response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
|
|
let finalize_report_kind = if explicit_finalize {
|
|
report_kind.clone()
|
|
} else {
|
|
mapped_error_finalize_kind
|
|
};
|
|
|
|
if let Some(finalize_report_kind) = finalize_report_kind {
|
|
let mut payload = build_sync_report_payload(
|
|
trace_id,
|
|
finalize_report_kind,
|
|
report_context,
|
|
status_code,
|
|
client_headers,
|
|
body_json,
|
|
body_base64,
|
|
telemetry,
|
|
);
|
|
if let Some(outcome) = maybe_build_sync_finalize_outcome(trace_id, decision, &payload)? {
|
|
let usage_payload = outcome.background_report.as_ref().unwrap_or(&payload);
|
|
if status_code < 400 {
|
|
apply_sync_success_effects(
|
|
state,
|
|
&plan,
|
|
payload.report_context.as_ref(),
|
|
usage_payload,
|
|
)
|
|
.await;
|
|
}
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
payload.report_context.as_ref(),
|
|
usage_payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
if let Some(report_payload) = outcome.background_report {
|
|
spawn_sync_report(state.clone(), report_payload);
|
|
} else {
|
|
warn!(
|
|
event_name = "local_core_finalize_missing_success_report_mapping",
|
|
log_type = "event",
|
|
trace_id = %trace_id,
|
|
report_kind = %payload.report_kind,
|
|
"gateway local core finalize produced response without background success report mapping"
|
|
);
|
|
}
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
outcome.response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
let mut payload = match maybe_build_local_video_success_outcome(
|
|
trace_id,
|
|
decision,
|
|
payload,
|
|
&state.video_tasks,
|
|
&plan,
|
|
)? {
|
|
LocalVideoSyncSuccessBuild::Handled(outcome) => {
|
|
let LocalVideoSyncSuccessOutcome {
|
|
response,
|
|
report_payload,
|
|
original_report_context,
|
|
report_mode,
|
|
local_task_snapshot,
|
|
} = outcome;
|
|
apply_sync_success_effects(
|
|
state,
|
|
&plan,
|
|
original_report_context.as_ref(),
|
|
&report_payload,
|
|
)
|
|
.await;
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
original_report_context.as_ref(),
|
|
&report_payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
if let Some(snapshot) = local_task_snapshot {
|
|
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
|
|
state.video_tasks.record_snapshot(snapshot);
|
|
}
|
|
match report_mode {
|
|
VideoTaskSyncReportMode::InlineSync => {
|
|
submit_sync_report(state, report_payload).await?;
|
|
}
|
|
VideoTaskSyncReportMode::Background => {
|
|
spawn_sync_report(state.clone(), report_payload);
|
|
}
|
|
}
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
LocalVideoSyncSuccessBuild::NotHandled(payload) => payload,
|
|
};
|
|
if let Some(response) =
|
|
maybe_build_local_sync_finalize_response(trace_id, decision, &payload)?
|
|
{
|
|
let background_success_report_kind =
|
|
resolve_local_sync_success_background_report_kind(payload.report_kind.as_str());
|
|
apply_sync_success_effects(state, &plan, payload.report_context.as_ref(), &payload)
|
|
.await;
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
payload.report_context.as_ref(),
|
|
&payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
state
|
|
.video_tasks
|
|
.apply_finalize_mutation(request_path, payload.report_kind.as_str());
|
|
if let Some(snapshot) = state
|
|
.video_tasks
|
|
.snapshot_for_route(decision.route_family.as_deref(), request_path)
|
|
{
|
|
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
|
|
}
|
|
if let Some(success_report_kind) = background_success_report_kind {
|
|
payload.report_kind = success_report_kind.to_string();
|
|
}
|
|
if background_success_report_kind.is_some() {
|
|
spawn_sync_report(state.clone(), payload);
|
|
} else {
|
|
warn!(
|
|
event_name = "local_video_finalize_missing_success_report_mapping",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %request_id_for_log,
|
|
candidate_id = ?candidate_id,
|
|
report_kind = %payload.report_kind,
|
|
"gateway local video finalize produced response without background success report mapping"
|
|
);
|
|
}
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
if let Some(response) =
|
|
maybe_build_local_video_error_response(trace_id, decision, &payload)?
|
|
{
|
|
let background_error_report_kind =
|
|
resolve_local_sync_error_background_report_kind(payload.report_kind.as_str());
|
|
if let Some(error_report_kind) = background_error_report_kind {
|
|
payload.report_kind = error_report_kind.to_string();
|
|
}
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
payload.report_context.as_ref(),
|
|
&payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
if background_error_report_kind.is_some() {
|
|
spawn_sync_report(state.clone(), payload);
|
|
} else {
|
|
warn!(
|
|
event_name = "local_video_finalize_missing_error_report_mapping",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %request_id_for_log,
|
|
candidate_id = ?candidate_id,
|
|
report_kind = %payload.report_kind,
|
|
"gateway local video finalize produced response without background error report mapping"
|
|
);
|
|
}
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
payload.report_context.as_ref(),
|
|
&payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
let response =
|
|
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
|
|
return Ok(Some(attach_control_metadata_headers(
|
|
response,
|
|
request_id,
|
|
candidate_id,
|
|
)?));
|
|
}
|
|
|
|
let usage_payload = build_sync_report_payload(
|
|
trace_id,
|
|
report_kind.unwrap_or_default(),
|
|
report_context,
|
|
status_code,
|
|
client_headers,
|
|
body_json,
|
|
body_base64,
|
|
telemetry,
|
|
);
|
|
if status_code < 400 {
|
|
apply_sync_success_effects(
|
|
state,
|
|
&plan,
|
|
usage_payload.report_context.as_ref(),
|
|
&usage_payload,
|
|
)
|
|
.await;
|
|
}
|
|
record_sync_terminal_usage_and_disarm_guard(
|
|
state,
|
|
&plan,
|
|
usage_payload.report_context.as_ref(),
|
|
&usage_payload,
|
|
candidate_started_at,
|
|
candidate_first_byte_elapsed_ms,
|
|
&mut terminal_guard,
|
|
)
|
|
.await;
|
|
let response = attach_control_metadata_headers(
|
|
build_client_response_from_parts(
|
|
status_code,
|
|
&usage_payload.headers,
|
|
Body::from(body_bytes),
|
|
trace_id,
|
|
Some(decision),
|
|
)?,
|
|
request_id,
|
|
candidate_id,
|
|
)?;
|
|
if !usage_payload.report_kind.trim().is_empty() {
|
|
if status_code >= 400 {
|
|
let report_kind = usage_payload.report_kind.clone();
|
|
if let Err(err) = submit_sync_report(state, usage_payload).await {
|
|
warn!(
|
|
event_name = "local_sync_error_report_submit_failed",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
report_kind = %report_kind,
|
|
"gateway failed to submit local sync error report before returning response: {err:?}"
|
|
);
|
|
}
|
|
} else {
|
|
spawn_sync_report(state.clone(), usage_payload);
|
|
}
|
|
}
|
|
|
|
Ok(Some(response))
|
|
})
|
|
.await;
|
|
if let Err(error) = result.as_ref() {
|
|
terminal_guard.fail_and_disarm(error).await;
|
|
} else {
|
|
terminal_guard.disarm();
|
|
}
|
|
result
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)] // mirrors sync execution context
|
|
fn maybe_build_implicit_sync_finalize_outcome(
|
|
trace_id: &str,
|
|
decision: &GatewayControlDecision,
|
|
plan_kind: &str,
|
|
report_context: &Option<serde_json::Value>,
|
|
status_code: u16,
|
|
headers: &BTreeMap<String, String>,
|
|
body_json: &Option<serde_json::Value>,
|
|
body_base64: &Option<String>,
|
|
telemetry: &Option<ExecutionTelemetry>,
|
|
) -> Result<Option<ImplicitSyncFinalizeOutcome>, GatewayError> {
|
|
let needs_conversion = report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("needs_conversion"))
|
|
.and_then(serde_json::Value::as_bool)
|
|
.unwrap_or(false);
|
|
let has_captured_stream_body = body_json.is_none() && body_base64.is_some();
|
|
let has_cross_format_sync_body = needs_conversion && body_json.is_some();
|
|
if status_code >= 400 || (!has_captured_stream_body && !has_cross_format_sync_body) {
|
|
return Ok(None);
|
|
}
|
|
|
|
let Some(report_kind) = implicit_sync_finalize_report_kind(plan_kind) else {
|
|
return Ok(None);
|
|
};
|
|
|
|
let payload = GatewaySyncReportRequest {
|
|
trace_id: trace_id.to_string(),
|
|
report_kind: report_kind.to_string(),
|
|
report_context: report_context.clone(),
|
|
status_code,
|
|
headers: headers.clone(),
|
|
body_json: body_json.clone(),
|
|
client_body_json: None,
|
|
body_base64: body_base64.clone(),
|
|
telemetry: telemetry.clone(),
|
|
};
|
|
let Some(outcome) = maybe_build_sync_finalize_outcome(trace_id, decision, &payload)? else {
|
|
return Ok(None);
|
|
};
|
|
|
|
Ok(Some(ImplicitSyncFinalizeOutcome { payload, outcome }))
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)] // internal helper mirroring execute path context
|
|
#[cfg(test)]
|
|
async fn execute_sync_via_remote_execution_runtime(
|
|
state: &AppState,
|
|
remote_execution_runtime_base_url: &str,
|
|
trace_id: &str,
|
|
decision: &GatewayControlDecision,
|
|
plan: &ExecutionPlan,
|
|
plan_request_id: &str,
|
|
plan_candidate_id: Option<&str>,
|
|
report_context: Option<&serde_json::Value>,
|
|
candidate_started_unix_secs: u64,
|
|
candidate_started_at: Instant,
|
|
) -> Result<RemoteSyncFallbackOutcome, GatewayError> {
|
|
let remote_request_started_at_unix_ms = current_request_candidate_unix_ms();
|
|
let remote_request_order_id = uuid::Uuid::now_v7().to_string();
|
|
let response = match post_sync_plan_to_remote_execution_runtime(
|
|
state,
|
|
remote_execution_runtime_base_url,
|
|
Some(trace_id),
|
|
plan,
|
|
)
|
|
.await
|
|
{
|
|
Ok(response) => response,
|
|
Err(err) => {
|
|
warn!(
|
|
event_name = "sync_execution_runtime_remote_unavailable",
|
|
log_type = "ops",
|
|
trace_id = %trace_id,
|
|
request_id = %short_request_id(plan_request_id),
|
|
candidate_id = ?plan_candidate_id,
|
|
error = ?err,
|
|
"gateway remote execution runtime sync unavailable"
|
|
);
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: None,
|
|
error_type: Some("execution_runtime_unavailable".to_string()),
|
|
error_message: Some(persisted_sync_abort_message(&err).to_string()),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
return Ok(RemoteSyncFallbackOutcome::Unavailable);
|
|
}
|
|
};
|
|
|
|
if response.status() != http::StatusCode::OK {
|
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
|
record_local_request_candidate_status(
|
|
state,
|
|
plan,
|
|
report_context,
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Failed,
|
|
status_code: Some(response.status().as_u16()),
|
|
error_type: Some("execution_runtime_http_error".to_string()),
|
|
error_message: Some(format!(
|
|
"execution runtime returned HTTP {}",
|
|
response.status()
|
|
)),
|
|
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
|
|
started_at_unix_ms: Some(candidate_started_unix_secs),
|
|
finished_at_unix_ms: Some(terminal_unix_secs),
|
|
},
|
|
)
|
|
.await;
|
|
return Ok(RemoteSyncFallbackOutcome::ClientResponse(
|
|
attach_control_metadata_headers(
|
|
build_client_response(response, trace_id, Some(decision))?,
|
|
Some(plan_request_id),
|
|
plan_candidate_id,
|
|
)?,
|
|
));
|
|
}
|
|
|
|
let remote_response_observed_at_unix_ms = current_request_candidate_unix_ms();
|
|
let response_body = aether_http::read_response_bytes_with_limit(
|
|
response,
|
|
crate::execution_runtime::transport::execution_result_envelope_limit_bytes(
|
|
crate::headers::max_internal_buffered_body_bytes(),
|
|
),
|
|
)
|
|
.await
|
|
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
|
let mut result = serde_json::from_slice::<ExecutionResult>(&response_body)
|
|
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
|
result
|
|
.response_observation
|
|
.get_or_insert(ExecutionResponseObservation {
|
|
request_started_at_unix_ms: remote_request_started_at_unix_ms,
|
|
response_headers_observed_at_unix_ms: remote_response_observed_at_unix_ms,
|
|
request_order_id: remote_request_order_id,
|
|
});
|
|
Ok(RemoteSyncFallbackOutcome::Executed(result))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
|
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
|
use aether_data_contracts::repository::candidates::RequestCandidateReadRepository;
|
|
use aether_data_contracts::repository::usage::UsageReadRepository;
|
|
use aether_usage_runtime::UsageRuntimeConfig;
|
|
use futures_util::{pin_mut, StreamExt as _};
|
|
use serde_json::json;
|
|
use std::collections::BTreeMap;
|
|
use std::sync::Arc;
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
|
|
fn test_openai_image_plan(stream: bool) -> ExecutionPlan {
|
|
ExecutionPlan {
|
|
request_id: "req-1".to_string(),
|
|
candidate_id: Some("candidate-1".to_string()),
|
|
provider_name: Some("OpenAI".to_string()),
|
|
provider_id: "provider-1".to_string(),
|
|
endpoint_id: "endpoint-1".to_string(),
|
|
key_id: "key-1".to_string(),
|
|
method: "POST".to_string(),
|
|
url: "https://chatgpt.com/backend-api/codex/responses".to_string(),
|
|
headers: BTreeMap::new(),
|
|
content_type: Some("application/json".to_string()),
|
|
content_encoding: None,
|
|
body: aether_contracts::RequestBody::from_json(json!({"stream": true})),
|
|
stream,
|
|
client_api_format: "openai:image".to_string(),
|
|
provider_api_format: "openai:image".to_string(),
|
|
model_name: Some("gpt-image-2".to_string()),
|
|
proxy: None,
|
|
transport_profile: None,
|
|
timeouts: None,
|
|
}
|
|
}
|
|
|
|
fn test_gemini_chat_plan() -> ExecutionPlan {
|
|
let mut plan = test_openai_image_plan(false);
|
|
plan.client_api_format = "openai:chat".to_string();
|
|
plan.provider_api_format = "gemini:generate_content".to_string();
|
|
plan.model_name = Some("gemini-3-flash-preview".to_string());
|
|
plan
|
|
}
|
|
|
|
fn test_decision() -> GatewayControlDecision {
|
|
GatewayControlDecision::synthetic(
|
|
"/v1/chat/completions",
|
|
Some("ai_public".to_string()),
|
|
Some("openai".to_string()),
|
|
Some("chat".to_string()),
|
|
Some("openai:chat".to_string()),
|
|
)
|
|
.with_execution_runtime_candidate(true)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn implicit_sync_finalize_converts_chat_json_to_namespaced_responses() {
|
|
let decision = GatewayControlDecision::synthetic(
|
|
"/v1/responses",
|
|
Some("ai_public".to_string()),
|
|
Some("openai".to_string()),
|
|
Some("responses".to_string()),
|
|
Some("openai:responses".to_string()),
|
|
)
|
|
.with_execution_runtime_candidate(true);
|
|
let report_context = Some(json!({
|
|
"provider_api_format": "openai:chat",
|
|
"client_api_format": "openai:responses",
|
|
"needs_conversion": true,
|
|
"mapped_model": "qwen-upstream",
|
|
"original_request_body": {
|
|
"model": "qwen",
|
|
"tools": [{
|
|
"type": "namespace",
|
|
"name": "mcp__vulnerability_report",
|
|
"description": "reporting tools",
|
|
"tools": [{
|
|
"type": "function",
|
|
"name": "vulnerability_report",
|
|
"description": "write the confirmed report",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"report_path": {"type": "string"}
|
|
},
|
|
"required": ["report_path"]
|
|
},
|
|
"strict": true
|
|
}]
|
|
}]
|
|
}
|
|
}));
|
|
let provider_body = Some(json!({
|
|
"id": "chatcmpl_namespace_sync",
|
|
"object": "chat.completion",
|
|
"created": 1_777_777_777,
|
|
"model": "qwen-upstream",
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": null,
|
|
"tool_calls": [{
|
|
"id": "call_report_1",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "vulnerability_report",
|
|
"arguments": "{\"report_path\":\"reports/sql-001-c1.md\"}"
|
|
}
|
|
}]
|
|
},
|
|
"finish_reason": "tool_calls"
|
|
}],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 4,
|
|
"total_tokens": 14
|
|
}
|
|
}));
|
|
|
|
let implicit = maybe_build_implicit_sync_finalize_outcome(
|
|
"trace-namespace-sync",
|
|
&decision,
|
|
"openai_responses_sync",
|
|
&report_context,
|
|
StatusCode::OK.as_u16(),
|
|
&BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
|
|
&provider_body,
|
|
&None,
|
|
&None,
|
|
)
|
|
.expect("cross-format sync JSON finalize should not error")
|
|
.expect("cross-format sync JSON should be finalized");
|
|
let response_body = axum::body::to_bytes(implicit.outcome.response.into_body(), usize::MAX)
|
|
.await
|
|
.expect("response body should read");
|
|
let response_json: Value =
|
|
serde_json::from_slice(&response_body).expect("response body should be JSON");
|
|
|
|
assert_eq!(response_json["object"], "response");
|
|
assert!(response_json.get("choices").is_none());
|
|
assert_eq!(response_json["output"][0]["type"], "function_call");
|
|
assert_eq!(response_json["output"][0]["name"], "vulnerability_report");
|
|
assert_eq!(
|
|
response_json["output"][0]["namespace"],
|
|
"mcp__vulnerability_report"
|
|
);
|
|
assert_eq!(response_json["output"][0]["call_id"], "call_report_1");
|
|
}
|
|
|
|
#[test]
|
|
fn implicit_sync_finalize_leaves_same_format_json_on_passthrough_path() {
|
|
let report_context = Some(json!({
|
|
"provider_api_format": "openai:responses",
|
|
"client_api_format": "openai:responses",
|
|
"needs_conversion": false
|
|
}));
|
|
let body_json = Some(json!({
|
|
"id": "resp_same_format",
|
|
"object": "response",
|
|
"status": "completed",
|
|
"output": []
|
|
}));
|
|
|
|
let outcome = maybe_build_implicit_sync_finalize_outcome(
|
|
"trace-same-format-sync",
|
|
&GatewayControlDecision::synthetic(
|
|
"/v1/responses",
|
|
Some("ai_public".to_string()),
|
|
Some("openai".to_string()),
|
|
Some("responses".to_string()),
|
|
Some("openai:responses".to_string()),
|
|
),
|
|
"openai_responses_sync",
|
|
&report_context,
|
|
StatusCode::OK.as_u16(),
|
|
&BTreeMap::new(),
|
|
&body_json,
|
|
&None,
|
|
&None,
|
|
)
|
|
.expect("same-format sync JSON guard should not error");
|
|
|
|
assert!(outcome.is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn oversized_upstream_response_builds_claude_502_retry_fallback() {
|
|
let mut plan = test_openai_image_plan(false);
|
|
plan.client_api_format = "claude:messages".to_string();
|
|
plan.provider_api_format = "claude:messages".to_string();
|
|
let decision = GatewayControlDecision::synthetic(
|
|
"/v1/messages",
|
|
Some("ai_public".to_string()),
|
|
Some("claude".to_string()),
|
|
Some("messages".to_string()),
|
|
Some("claude:messages".to_string()),
|
|
)
|
|
.with_execution_runtime_candidate(true);
|
|
let failure = SyncExecutionFailure::from_transport(
|
|
ExecutionRuntimeTransportError::UpstreamResponseTooLarge {
|
|
phase: crate::execution_runtime::transport::UpstreamResponseBodyPhase::Wire,
|
|
limit_bytes: 8,
|
|
},
|
|
);
|
|
|
|
assert_eq!(failure.status_code, Some(StatusCode::BAD_GATEWAY.as_u16()));
|
|
assert_eq!(
|
|
failure.fallback_kind,
|
|
Some(SyncExecutionFailureFallbackKind::UpstreamResponseTooLarge)
|
|
);
|
|
let mut retry_scope = AiAttemptRetryScope::Provider;
|
|
let mut retry_fallback = None;
|
|
{
|
|
let mut retry_scope_out = Some(&mut retry_scope);
|
|
let mut retry_fallback_out = Some(&mut retry_fallback);
|
|
maybe_store_sync_execution_failure_fallback(
|
|
&failure,
|
|
&plan,
|
|
"trace-too-large",
|
|
&decision,
|
|
&mut retry_scope_out,
|
|
&mut retry_fallback_out,
|
|
)
|
|
.expect("fallback response should build");
|
|
}
|
|
|
|
assert_eq!(retry_scope, AiAttemptRetryScope::Candidate);
|
|
let response = retry_fallback.expect("oversized response should provide a fallback");
|
|
|
|
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
|
|
let body = to_bytes(response.into_body(), 1024)
|
|
.await
|
|
.expect("fallback body should read");
|
|
let body: Value = serde_json::from_slice(&body).expect("fallback body should be json");
|
|
assert_eq!(body["type"], "error");
|
|
assert_eq!(body["error"]["type"], "upstream_error");
|
|
assert_eq!(body["error"]["message"], "Upstream response too large");
|
|
}
|
|
|
|
fn test_kiro_sync_plan() -> ExecutionPlan {
|
|
ExecutionPlan {
|
|
request_id: "req-kiro-sync-cache-1".to_string(),
|
|
candidate_id: Some("candidate-kiro-sync-cache-1".to_string()),
|
|
provider_name: Some("Kiro".to_string()),
|
|
provider_id: "provider-kiro-sync-1".to_string(),
|
|
endpoint_id: "endpoint-kiro-sync-1".to_string(),
|
|
key_id: "key-kiro-sync-1".to_string(),
|
|
method: "POST".to_string(),
|
|
url: "https://kiro.example/generateAssistantResponse".to_string(),
|
|
headers: BTreeMap::new(),
|
|
content_type: Some("application/json".to_string()),
|
|
content_encoding: None,
|
|
body: aether_contracts::RequestBody::from_json(json!({
|
|
"model": "claude-sonnet-4",
|
|
"messages": [{"role": "user", "content": "hello kiro"}],
|
|
})),
|
|
stream: false,
|
|
client_api_format: "claude:messages".to_string(),
|
|
provider_api_format: "claude:messages".to_string(),
|
|
model_name: Some("claude-sonnet-4".to_string()),
|
|
proxy: None,
|
|
transport_profile: None,
|
|
timeouts: None,
|
|
}
|
|
}
|
|
|
|
fn test_kiro_sync_cacheable_request_body() -> serde_json::Value {
|
|
json!({
|
|
"model": "claude-sonnet-4",
|
|
"system": [{
|
|
"type": "text",
|
|
"text": format!("sync cacheable prompt {}", "cacheable prompt chunk ".repeat(300)),
|
|
"cache_control": {"type": "ephemeral"}
|
|
}],
|
|
"messages": [{"role": "user", "content": "reuse this Kiro prompt"}]
|
|
})
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_gemini_provider_success_uses_plan_format_when_context_is_missing() {
|
|
let plan = test_gemini_chat_plan();
|
|
let body = json!({
|
|
"candidates": [{
|
|
"content": {"role": "model"},
|
|
"finishReason": "MAX_TOKENS"
|
|
}],
|
|
"usageMetadata": {
|
|
"promptTokenCount": 8,
|
|
"candidatesTokenCount": 1,
|
|
"thoughtsTokenCount": 25,
|
|
"totalTokenCount": 34
|
|
}
|
|
});
|
|
|
|
let message = invalid_gemini_provider_success_message(
|
|
&plan,
|
|
None,
|
|
StatusCode::OK.as_u16(),
|
|
Some(&body),
|
|
)
|
|
.expect("empty Gemini 200 response should be rejected from plan format");
|
|
|
|
assert!(message.contains("visible model output"));
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_gemini_provider_success_accepts_thought_only_max_tokens() {
|
|
let plan = test_gemini_chat_plan();
|
|
let body = json!({
|
|
"candidates": [{
|
|
"content": {
|
|
"role": "model",
|
|
"parts": [{"text": "hidden plan", "thought": true}]
|
|
},
|
|
"finishReason": "MAX_TOKENS"
|
|
}],
|
|
"usageMetadata": {
|
|
"promptTokenCount": 8,
|
|
"candidatesTokenCount": 0,
|
|
"thoughtsTokenCount": 24,
|
|
"totalTokenCount": 32
|
|
}
|
|
});
|
|
|
|
let message = invalid_gemini_provider_success_message(
|
|
&plan,
|
|
None,
|
|
StatusCode::OK.as_u16(),
|
|
Some(&body),
|
|
);
|
|
|
|
assert!(message.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_gemini_provider_stream_success_accepts_signature_only_reasoning_exhaustion() {
|
|
let plan = test_gemini_chat_plan();
|
|
let report_context = json!({
|
|
"has_envelope": true,
|
|
"envelope_name": "antigravity:v1internal",
|
|
"provider_api_format": "gemini:generate_content",
|
|
});
|
|
let body = concat!(
|
|
"data: {\"response\":{\"responseId\":\"resp_signature_only_123\",\"modelVersion\":\"gemini-3.7-flash-tiered\",",
|
|
"\"candidates\":[{\"index\":0,\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"\",\"thoughtSignature\":\"opaque-thought-signature\"}]},\"finishReason\":\"MAX_TOKENS\"}],",
|
|
"\"usageMetadata\":{\"promptTokenCount\":22,\"thoughtsTokenCount\":29,\"totalTokenCount\":51}},",
|
|
"\"traceId\":\"trace-signature-only\"}\n\n",
|
|
);
|
|
|
|
let message = invalid_gemini_provider_stream_success_message(
|
|
&plan,
|
|
Some(&report_context),
|
|
StatusCode::OK.as_u16(),
|
|
None,
|
|
body.as_bytes(),
|
|
true,
|
|
);
|
|
|
|
assert!(message.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_gemini_provider_success_error_is_retryable_candidate_failure() {
|
|
let error = invalid_gemini_provider_success_execution_error(
|
|
INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE,
|
|
);
|
|
|
|
assert_eq!(error.kind, ExecutionErrorKind::Upstream5xx);
|
|
assert_eq!(error.phase, ExecutionPhase::Finalize);
|
|
assert_eq!(error.upstream_status, Some(StatusCode::OK.as_u16()));
|
|
assert!(error.retryable);
|
|
assert!(error.failover_recommended);
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_gemini_provider_success_accepts_antigravity_chunks_with_visible_output() {
|
|
let plan = test_gemini_chat_plan();
|
|
let report_context = json!({
|
|
"has_envelope": true,
|
|
"envelope_name": "antigravity:v1internal",
|
|
"provider_api_format": "gemini:generate_content",
|
|
});
|
|
let body = json!({
|
|
"chunks": [{
|
|
"response": {
|
|
"responseId": "resp_antigravity_chunks_123",
|
|
"candidates": [{
|
|
"content": {
|
|
"parts": [{"text": "Hello Gemini"}],
|
|
"role": "model"
|
|
},
|
|
"finishReason": "STOP",
|
|
"index": 0
|
|
}],
|
|
"modelVersion": "gemini-3-flash-agent",
|
|
"usageMetadata": {
|
|
"promptTokenCount": 2,
|
|
"candidatesTokenCount": 2,
|
|
"totalTokenCount": 4
|
|
}
|
|
},
|
|
"traceId": "trace-antigravity-chunks"
|
|
}],
|
|
"metadata": {
|
|
"stream": true,
|
|
"stored_chunks": 1,
|
|
"total_chunks": 1
|
|
}
|
|
});
|
|
|
|
let message = invalid_gemini_provider_success_message(
|
|
&plan,
|
|
Some(&report_context),
|
|
StatusCode::OK.as_u16(),
|
|
Some(&body),
|
|
);
|
|
|
|
assert!(message.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_gemini_provider_success_unwraps_gemini_cli_v1internal_envelope() {
|
|
let plan = test_gemini_chat_plan();
|
|
let report_context = json!({
|
|
"has_envelope": true,
|
|
"envelope_name": "gemini_cli:v1internal",
|
|
"provider_api_format": "gemini:generate_content",
|
|
});
|
|
let body = json!({
|
|
"response": {
|
|
"candidates": [{
|
|
"content": {
|
|
"role": "model",
|
|
"parts": [{"text": "Hello from Gemini CLI"}]
|
|
},
|
|
"finishReason": "STOP"
|
|
}]
|
|
},
|
|
"remainingCredits": 41,
|
|
"consumedCredits": 1,
|
|
"traceId": "trace-upstream-sync-1"
|
|
});
|
|
|
|
let message = invalid_gemini_provider_success_message(
|
|
&plan,
|
|
Some(&report_context),
|
|
StatusCode::OK.as_u16(),
|
|
Some(&body),
|
|
);
|
|
|
|
assert!(message.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn forced_sync_abort_message_does_not_copy_internal_error_details() {
|
|
let secret = "https://user:[email protected]/v1?api_key=should-not-persist";
|
|
let error = GatewayError::Internal(secret.to_string());
|
|
|
|
let message = persisted_sync_abort_message(&error);
|
|
assert_eq!(
|
|
message,
|
|
"local sync attempt aborted before terminal finalization: internal error"
|
|
);
|
|
assert!(!message.contains("password"));
|
|
assert!(!message.contains("api_key"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sync_attempt_terminal_guard_marks_dropped_pending_attempt_cancelled() {
|
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
|
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(
|
|
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
|
Arc::clone(&request_candidate_repository),
|
|
Arc::clone(&usage_repository),
|
|
),
|
|
)
|
|
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
|
enabled: true,
|
|
..UsageRuntimeConfig::default()
|
|
});
|
|
let mut plan = test_openai_image_plan(false);
|
|
plan.request_id = "sync-cancel-guard-request".to_string();
|
|
plan.candidate_id = None;
|
|
let mut report_context = Some(json!({
|
|
"candidate_index": 0,
|
|
"retry_index": 0,
|
|
"user_id": "user-cancel",
|
|
"api_key_id": "api-key-cancel",
|
|
"client_api_format": "openai:image",
|
|
"provider_api_format": "openai:image",
|
|
"request_path": "/v1/images/generations",
|
|
"request_path_and_query": "/v1/images/generations",
|
|
"upstream_url": "https://example.test/v1/images/generations",
|
|
"mapped_model": "gpt-image-2",
|
|
}));
|
|
|
|
ensure_execution_request_candidate_slot(&state, &mut plan, &mut report_context).await;
|
|
let candidate_started_at = Instant::now();
|
|
let started_at = current_request_candidate_unix_ms();
|
|
state.usage_runtime.record_pending(
|
|
state.usage_lifecycle_data_state().as_ref(),
|
|
build_lifecycle_usage_seed(&plan, report_context.as_ref()),
|
|
);
|
|
record_local_request_candidate_status(
|
|
&state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Pending,
|
|
status_code: None,
|
|
error_type: None,
|
|
error_message: None,
|
|
latency_ms: None,
|
|
started_at_unix_ms: Some(started_at),
|
|
finished_at_unix_ms: None,
|
|
},
|
|
)
|
|
.await;
|
|
|
|
crate::request_diagnostics::scope_request_diagnostics(async {
|
|
crate::request_diagnostics::record_request_accepted_at(
|
|
Instant::now() - Duration::from_millis(25),
|
|
);
|
|
let _guard = SyncAttemptTerminalGuard::new(
|
|
&state,
|
|
&plan,
|
|
report_context.clone(),
|
|
started_at,
|
|
candidate_started_at,
|
|
);
|
|
})
|
|
.await;
|
|
|
|
let mut stored_usage = None;
|
|
for _ in 0..50 {
|
|
if let Some(usage) = usage_repository
|
|
.find_by_request_id("sync-cancel-guard-request")
|
|
.await
|
|
.expect("usage should read")
|
|
{
|
|
if usage.status == "cancelled" {
|
|
stored_usage = Some(usage);
|
|
break;
|
|
}
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
}
|
|
let stored_usage = stored_usage.expect("cancelled usage should be recorded");
|
|
assert_eq!(stored_usage.status, "cancelled");
|
|
assert_eq!(stored_usage.billing_status, "void");
|
|
assert_eq!(stored_usage.status_code, Some(499));
|
|
assert_eq!(stored_usage.error_category.as_deref(), Some("cancelled"));
|
|
let request_metadata = stored_usage
|
|
.request_metadata
|
|
.as_ref()
|
|
.expect("cancelled usage should retain request diagnostics");
|
|
assert!(request_metadata
|
|
.get("end_to_end_time_ms")
|
|
.and_then(Value::as_u64)
|
|
.is_some());
|
|
assert!(request_metadata
|
|
.get("end_to_end_first_byte_time_ms")
|
|
.is_none());
|
|
|
|
let stored_candidates = request_candidate_repository
|
|
.list_by_request_id("sync-cancel-guard-request")
|
|
.await
|
|
.expect("request candidates should read");
|
|
assert_eq!(stored_candidates.len(), 1);
|
|
assert_eq!(
|
|
stored_candidates[0].status,
|
|
RequestCandidateStatus::Cancelled
|
|
);
|
|
assert_eq!(stored_candidates[0].status_code, Some(499));
|
|
assert_eq!(
|
|
stored_candidates[0].error_type.as_deref(),
|
|
Some("local_sync_attempt_cancelled")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sync_direct_response_start_marks_usage_and_candidate_active_before_body_finishes() {
|
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
|
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(
|
|
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
|
Arc::clone(&request_candidate_repository),
|
|
Arc::clone(&usage_repository),
|
|
),
|
|
)
|
|
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
|
enabled: true,
|
|
..UsageRuntimeConfig::default()
|
|
});
|
|
|
|
let listener = crate::test_support::bind_loopback_listener()
|
|
.await
|
|
.expect("listener should bind");
|
|
let addr = listener.local_addr().expect("local addr should resolve");
|
|
let (headers_tx, headers_rx) = tokio::sync::oneshot::channel();
|
|
let (body_tx, body_rx) = tokio::sync::oneshot::channel();
|
|
let server = tokio::spawn(async move {
|
|
let (mut socket, _) = listener.accept().await.expect("client should connect");
|
|
let mut request = [0_u8; 4096];
|
|
let _ = socket
|
|
.read(&mut request)
|
|
.await
|
|
.expect("request should read");
|
|
socket
|
|
.write_all(
|
|
b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 11\r\n\r\n",
|
|
)
|
|
.await
|
|
.expect("headers should write");
|
|
socket.flush().await.expect("headers should flush");
|
|
let _ = headers_tx.send(());
|
|
let _ = body_rx.await;
|
|
socket
|
|
.write_all(br#"{"ok":true}"#)
|
|
.await
|
|
.expect("body should write");
|
|
});
|
|
|
|
let mut plan = test_gemini_chat_plan();
|
|
plan.request_id = "sync-response-start-active-request".to_string();
|
|
plan.candidate_id = Some("sync-response-start-active-candidate".to_string());
|
|
plan.url = format!("http://{addr}/chat");
|
|
plan.provider_api_format = "openai:chat".to_string();
|
|
plan.model_name = Some("gpt-5".to_string());
|
|
plan.body = aether_contracts::RequestBody::from_json(json!({
|
|
"model": "gpt-5",
|
|
"messages": [{"role": "user", "content": "slow body"}],
|
|
}));
|
|
let report_context = Some(json!({
|
|
"candidate_index": 0,
|
|
"retry_index": 0,
|
|
"user_id": "user-active",
|
|
"api_key_id": "api-key-active",
|
|
"client_api_format": "openai:chat",
|
|
"provider_api_format": "openai:chat",
|
|
"request_path": "/v1/chat/completions",
|
|
"request_path_and_query": "/v1/chat/completions",
|
|
"upstream_url": plan.url.clone(),
|
|
"mapped_model": "gpt-5",
|
|
}));
|
|
let started_at = current_request_candidate_unix_ms();
|
|
state
|
|
.usage_runtime
|
|
.record_pending_direct(
|
|
state.usage_lifecycle_data_state().as_ref(),
|
|
build_lifecycle_usage_seed(&plan, report_context.as_ref()),
|
|
)
|
|
.await;
|
|
record_local_request_candidate_status(
|
|
&state,
|
|
&plan,
|
|
report_context.as_ref(),
|
|
SchedulerRequestCandidateStatusUpdate {
|
|
status: RequestCandidateStatus::Pending,
|
|
status_code: None,
|
|
error_type: None,
|
|
error_message: None,
|
|
latency_ms: None,
|
|
started_at_unix_ms: Some(started_at),
|
|
finished_at_unix_ms: None,
|
|
},
|
|
)
|
|
.await;
|
|
|
|
let state_for_exec = state.clone();
|
|
let plan_for_exec = plan.clone();
|
|
let report_context_for_exec = report_context.clone();
|
|
let exec = tokio::spawn(async move {
|
|
execute_direct_sync_runtime_candidate(
|
|
&state_for_exec,
|
|
&plan_for_exec,
|
|
report_context_for_exec.as_ref(),
|
|
"trace-response-start-active",
|
|
"openai_chat_sync",
|
|
started_at,
|
|
"sync-response-start-active-request",
|
|
plan_for_exec.candidate_id.as_deref(),
|
|
"openai",
|
|
"endpoint-1",
|
|
"key-1",
|
|
"gpt-5",
|
|
"0",
|
|
None,
|
|
)
|
|
.await
|
|
});
|
|
|
|
headers_rx.await.expect("headers should be written");
|
|
let mut active_usage = None;
|
|
for _ in 0..50 {
|
|
if let Some(usage) = usage_repository
|
|
.find_by_request_id("sync-response-start-active-request")
|
|
.await
|
|
.expect("usage should read")
|
|
{
|
|
if usage.status == "streaming" {
|
|
active_usage = Some(usage);
|
|
break;
|
|
}
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
}
|
|
let active_usage = active_usage.expect("usage should become active before body finishes");
|
|
assert_eq!(active_usage.status_code, Some(200));
|
|
assert!(active_usage.first_byte_time_ms.is_some());
|
|
assert!(active_usage.response_time_ms.is_some());
|
|
|
|
let stored_candidates = request_candidate_repository
|
|
.list_by_request_id("sync-response-start-active-request")
|
|
.await
|
|
.expect("candidate should read");
|
|
let active_candidate = stored_candidates
|
|
.iter()
|
|
.find(|candidate| candidate.id == "sync-response-start-active-candidate")
|
|
.expect("candidate should exist");
|
|
assert_eq!(active_candidate.status, RequestCandidateStatus::Streaming);
|
|
assert_eq!(active_candidate.status_code, Some(200));
|
|
|
|
let _ = body_tx.send(());
|
|
let result = tokio::time::timeout(Duration::from_secs(2), exec)
|
|
.await
|
|
.expect("sync execution should finish")
|
|
.expect("sync execution task should not panic")
|
|
.expect("sync execution should succeed");
|
|
assert_eq!(result.status_code, 200);
|
|
server.abort();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn sync_execution_active_marks_usage_before_response_headers() {
|
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
|
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(
|
|
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
|
Arc::clone(&request_candidate_repository),
|
|
Arc::clone(&usage_repository),
|
|
),
|
|
)
|
|
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
|
enabled: true,
|
|
..UsageRuntimeConfig::default()
|
|
});
|
|
|
|
let listener = crate::test_support::bind_loopback_listener()
|
|
.await
|
|
.expect("listener should bind");
|
|
let addr = listener.local_addr().expect("local addr should resolve");
|
|
let (request_seen_tx, request_seen_rx) = tokio::sync::oneshot::channel();
|
|
let (finish_tx, finish_rx) = tokio::sync::oneshot::channel();
|
|
let server = tokio::spawn(async move {
|
|
let (mut socket, _) = listener.accept().await.expect("client should connect");
|
|
let mut request = [0_u8; 4096];
|
|
let _ = socket
|
|
.read(&mut request)
|
|
.await
|
|
.expect("request should read");
|
|
let _ = request_seen_tx.send(());
|
|
let _ = finish_rx.await;
|
|
socket
|
|
.write_all(
|
|
b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 11\r\n\r\n{\"ok\":true}",
|
|
)
|
|
.await
|
|
.expect("response should write");
|
|
});
|
|
|
|
let mut plan = test_gemini_chat_plan();
|
|
plan.request_id = "sync-active-before-headers-request".to_string();
|
|
plan.candidate_id = Some("sync-active-before-headers-candidate".to_string());
|
|
plan.url = format!("http://{addr}/chat");
|
|
plan.provider_name = Some("OpenAI".to_string());
|
|
plan.provider_api_format = "openai:chat".to_string();
|
|
plan.model_name = Some("gpt-5".to_string());
|
|
plan.body = aether_contracts::RequestBody::from_json(json!({
|
|
"model": "gpt-5",
|
|
"messages": [{"role": "user", "content": "slow headers"}],
|
|
}));
|
|
let report_context = Some(json!({
|
|
"candidate_index": 0,
|
|
"retry_index": 0,
|
|
"user_id": "user-active-before-headers",
|
|
"api_key_id": "api-key-active-before-headers",
|
|
"candidate_id": "sync-active-before-headers-candidate",
|
|
"provider_id": "provider-1",
|
|
"endpoint_id": "endpoint-1",
|
|
"key_id": "key-1",
|
|
"provider_name": "OpenAI",
|
|
"client_api_format": "openai:chat",
|
|
"provider_api_format": "openai:chat",
|
|
"request_path": "/v1/chat/completions",
|
|
"request_path_and_query": "/v1/chat/completions",
|
|
"upstream_url": plan.url.clone(),
|
|
"mapped_model": "gpt-5",
|
|
}));
|
|
let state_for_exec = state.clone();
|
|
let plan_for_exec = plan.clone();
|
|
let report_context_for_exec = report_context.clone();
|
|
let exec = tokio::spawn(async move {
|
|
execute_execution_runtime_sync(
|
|
&state_for_exec,
|
|
"/v1/chat/completions",
|
|
plan_for_exec,
|
|
"trace-active-before-headers",
|
|
&test_decision(),
|
|
"openai_chat_sync",
|
|
Some("openai_chat_sync".to_string()),
|
|
report_context_for_exec,
|
|
)
|
|
.await
|
|
});
|
|
|
|
request_seen_rx
|
|
.await
|
|
.expect("upstream request should be observed");
|
|
let mut active_usage = None;
|
|
for _ in 0..50 {
|
|
if let Some(usage) = usage_repository
|
|
.find_by_request_id("sync-active-before-headers-request")
|
|
.await
|
|
.expect("usage should read")
|
|
{
|
|
if usage.status == "streaming" {
|
|
active_usage = Some(usage);
|
|
break;
|
|
}
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
}
|
|
let active_usage =
|
|
active_usage.expect("usage should become active before upstream headers");
|
|
assert_eq!(active_usage.status_code, None);
|
|
assert_eq!(active_usage.first_byte_time_ms, None);
|
|
assert_eq!(active_usage.response_time_ms, None);
|
|
|
|
let stored_candidates = request_candidate_repository
|
|
.list_by_request_id("sync-active-before-headers-request")
|
|
.await
|
|
.expect("candidate should read");
|
|
let active_candidate = stored_candidates
|
|
.iter()
|
|
.find(|candidate| candidate.id == "sync-active-before-headers-candidate")
|
|
.expect("candidate should exist");
|
|
assert_eq!(active_candidate.status, RequestCandidateStatus::Streaming);
|
|
assert_eq!(active_candidate.status_code, None);
|
|
assert!(active_candidate.started_at_unix_ms.is_some());
|
|
assert!(active_candidate.finished_at_unix_ms.is_none());
|
|
|
|
let _ = finish_tx.send(());
|
|
let response = tokio::time::timeout(Duration::from_secs(2), exec)
|
|
.await
|
|
.expect("sync execution should finish")
|
|
.expect("sync execution task should not panic")
|
|
.expect("sync execution should succeed")
|
|
.expect("sync execution should produce a response");
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
server.abort();
|
|
}
|
|
|
|
#[test]
|
|
fn kiro_sync_report_context_seeds_input_tokens_from_original_request_body() {
|
|
let plan = test_kiro_sync_plan();
|
|
let mut report_context = Some(json!({
|
|
"original_request_body": test_kiro_sync_cacheable_request_body(),
|
|
}));
|
|
|
|
seed_kiro_sync_report_context_input_tokens(&plan, &mut report_context);
|
|
|
|
assert!(report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("input_tokens"))
|
|
.and_then(Value::as_u64)
|
|
.is_some_and(|tokens| tokens > 0));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn kiro_sync_report_context_applies_prompt_cache_usage_from_tracker() {
|
|
let state = AppState::new().expect("gateway state should build");
|
|
let plan = test_kiro_sync_plan();
|
|
|
|
let mut first_report_context = Some(json!({
|
|
"original_request_body": test_kiro_sync_cacheable_request_body(),
|
|
"kiro_simulated_cache_enabled": true,
|
|
}));
|
|
seed_kiro_sync_report_context_input_tokens(&plan, &mut first_report_context);
|
|
seed_kiro_sync_report_context_prompt_cache_usage(&state, &plan, &mut first_report_context)
|
|
.await;
|
|
let first_creation = first_report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("cache_creation_input_tokens"))
|
|
.and_then(Value::as_u64)
|
|
.unwrap_or_default();
|
|
let first_read = first_report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("cache_read_input_tokens"))
|
|
.and_then(Value::as_u64)
|
|
.unwrap_or_default();
|
|
assert!(first_creation > 0);
|
|
assert_eq!(first_read, 0);
|
|
|
|
let mut second_report_context = Some(json!({
|
|
"original_request_body": test_kiro_sync_cacheable_request_body(),
|
|
"kiro_simulated_cache_enabled": true,
|
|
}));
|
|
seed_kiro_sync_report_context_input_tokens(&plan, &mut second_report_context);
|
|
seed_kiro_sync_report_context_prompt_cache_usage(&state, &plan, &mut second_report_context)
|
|
.await;
|
|
let second_creation = second_report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("cache_creation_input_tokens"))
|
|
.and_then(Value::as_u64)
|
|
.unwrap_or_default();
|
|
let second_read = second_report_context
|
|
.as_ref()
|
|
.and_then(|value| value.get("cache_read_input_tokens"))
|
|
.and_then(Value::as_u64)
|
|
.unwrap_or_default();
|
|
assert_eq!(second_creation, 0);
|
|
assert!(second_read > 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn json_whitespace_heartbeat_stream_prefixes_final_json() {
|
|
let (tx, rx) = mpsc::channel::<Result<Bytes, IoError>>(1);
|
|
tx.send(Ok(Bytes::from_static(br#"{"data":[]}"#)))
|
|
.await
|
|
.expect("final body should send");
|
|
drop(tx);
|
|
|
|
let body = to_bytes(
|
|
Body::from_stream(build_json_whitespace_heartbeat_stream(
|
|
rx,
|
|
Duration::from_secs(60),
|
|
None,
|
|
)),
|
|
usize::MAX,
|
|
)
|
|
.await
|
|
.expect("body should collect");
|
|
|
|
assert!(body.starts_with(b"\n"));
|
|
let parsed: Value =
|
|
serde_json::from_slice(&body).expect("leading whitespace is valid JSON");
|
|
assert_eq!(parsed, json!({"data": []}));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn json_whitespace_heartbeat_stream_emits_interval_whitespace() {
|
|
let (tx, rx) = mpsc::channel::<Result<Bytes, IoError>>(1);
|
|
let stream = build_json_whitespace_heartbeat_stream(rx, Duration::from_millis(5), None);
|
|
pin_mut!(stream);
|
|
|
|
let first = stream
|
|
.next()
|
|
.await
|
|
.expect("initial whitespace")
|
|
.expect("initial whitespace ok");
|
|
assert_eq!(first, Bytes::from_static(b"\n"));
|
|
|
|
let second = tokio::time::timeout(Duration::from_millis(100), stream.next())
|
|
.await
|
|
.expect("interval heartbeat")
|
|
.expect("interval heartbeat item")
|
|
.expect("interval heartbeat ok");
|
|
assert_eq!(second, Bytes::from_static(b"\n"));
|
|
|
|
tx.send(Ok(Bytes::from_static(br#"{"data":[{"b64_json":"x"}]}"#)))
|
|
.await
|
|
.expect("final body should send");
|
|
let final_body = tokio::time::timeout(Duration::from_millis(100), stream.next())
|
|
.await
|
|
.expect("final body")
|
|
.expect("final body item")
|
|
.expect("final body ok");
|
|
assert_eq!(
|
|
serde_json::from_slice::<Value>(&final_body).expect("final body json"),
|
|
json!({"data": [{"b64_json": "x"}]})
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn openai_image_sync_sse_parser_tracks_partial_and_completed_frames() {
|
|
let partial = parse_openai_image_sync_sse_frame(
|
|
concat!(
|
|
"event: response.image_generation_call.partial_image\n",
|
|
"data: {\"type\":\"response.image_generation_call.partial_image\",\"partial_image_index\":0}\n\n"
|
|
)
|
|
.as_bytes(),
|
|
)
|
|
.expect("partial frame");
|
|
assert_eq!(
|
|
partial.event_name,
|
|
"response.image_generation_call.partial_image"
|
|
);
|
|
assert!(partial.is_partial_image);
|
|
assert_eq!(
|
|
partial.client_visible_event,
|
|
Some("image_generation.partial_image")
|
|
);
|
|
|
|
let completed = parse_openai_image_sync_sse_frame(
|
|
b"data: {\"type\":\"response.completed\",\"response\":{}}\n\n",
|
|
)
|
|
.expect("completed frame");
|
|
assert_eq!(completed.event_name, "response.completed");
|
|
assert!(completed.is_completed);
|
|
assert_eq!(
|
|
completed.client_visible_event,
|
|
Some("image_generation.completed")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn openai_image_sync_progress_tracks_upstream_stream_without_json_heartbeat_wrapper() {
|
|
let plan = test_openai_image_plan(false);
|
|
let report_context = json!({"upstream_is_stream": true});
|
|
|
|
assert!(should_track_openai_image_sync_upstream_sse(
|
|
OPENAI_IMAGE_SYNC_PLAN_KIND,
|
|
&plan,
|
|
Some(&report_context),
|
|
));
|
|
assert!(!should_enable_openai_image_sync_json_heartbeat(
|
|
OPENAI_IMAGE_SYNC_PLAN_KIND,
|
|
&plan,
|
|
Some(&report_context),
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn openai_image_sync_progress_ignores_non_stream_upstream() {
|
|
let plan = test_openai_image_plan(false);
|
|
let report_context = json!({"upstream_is_stream": false});
|
|
|
|
assert!(!should_track_openai_image_sync_upstream_sse(
|
|
OPENAI_IMAGE_SYNC_PLAN_KIND,
|
|
&plan,
|
|
Some(&report_context),
|
|
));
|
|
assert!(!should_enable_openai_image_sync_json_heartbeat(
|
|
OPENAI_IMAGE_SYNC_PLAN_KIND,
|
|
&plan,
|
|
Some(&report_context),
|
|
));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn openai_image_sync_json_heartbeat_hides_internal_error_details() {
|
|
let secret = "https://user:[email protected]/private?api_key=top-secret";
|
|
let bytes = openai_image_sync_json_heartbeat_final_bytes(Err(GatewayError::Internal(
|
|
secret.to_string(),
|
|
)))
|
|
.await;
|
|
|
|
let body: Value = serde_json::from_slice(&bytes).expect("heartbeat error body is JSON");
|
|
assert_eq!(
|
|
body.pointer("/error/message").and_then(Value::as_str),
|
|
Some("sync image execution failed")
|
|
);
|
|
let body_text = String::from_utf8(bytes).expect("heartbeat body is UTF-8");
|
|
assert!(!body_text.contains(secret));
|
|
assert!(!body_text.contains("top-secret"));
|
|
}
|
|
}
|